Update app.py

app.py

@@ -2,37 +2,25 @@ import gradio as gr
 import os
 from typing import List, Dict, Any, Optional
 import hashlib
-import json
 from datetime import datetime
-import tempfile
+import numpy as np
 
-# PDF 처리 라이브러리 (설치 필요한 경우를 위한 대체 구현 포함)
+# PDF 처리 라이브러리
 try:
     import fitz  # PyMuPDF
     PDF_AVAILABLE = True
 except ImportError:
     PDF_AVAILABLE = False
-    print("PyMuPDF not installed. Install with: pip install pymupdf")
-
-try:
-    import chromadb
-    from chromadb.utils import embedding_functions
-    CHROMA_AVAILABLE = True
-except ImportError:
-    CHROMA_AVAILABLE = False
-    print("ChromaDB not installed. Install with: pip install chromadb")
+    print("⚠️ PyMuPDF not installed. Install with: pip install pymupdf")
 
 try:
     from sentence_transformers import SentenceTransformer
     ST_AVAILABLE = True
 except ImportError:
     ST_AVAILABLE = False
-    print("Sentence Transformers not installed. Install with: pip install sentence-transformers")
+    print("⚠️ Sentence Transformers not installed. Install with: pip install sentence-transformers")
 
-import numpy as np
-from typing import Tuple
-
-# Custom CSS (기존 CSS + 추가 스타일)
+# Custom CSS for gradient background and styling
 custom_css = """
 .gradio-container {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #4facfe 75%, #00f2fe 100%);
@@ -79,81 +67,80 @@ custom_css = """
     border: 1px solid rgba(248, 113, 113, 0.5);
     color: #ef4444;
 }
-.pdf-processing {
-    background-color: rgba(251, 191, 36, 0.2);
-    border: 1px solid rgba(251, 191, 36, 0.5);
-    color: #f59e0b;
+.pdf-info {
+    background-color: rgba(59, 130, 246, 0.2);
+    border: 1px solid rgba(59, 130, 246, 0.5);
+    color: #3b82f6;
+}
+.rag-context {
+    background-color: rgba(251, 191, 36, 0.1);
+    border-left: 4px solid #f59e0b;
+    padding: 10px;
+    margin: 10px 0;
+    border-radius: 5px;
 }
 """
 
 class SimpleTextSplitter:
-    """간단한 텍스트 분할기"""
-    def __init__(self, chunk_size=1000, chunk_overlap=200):
+    """텍스트 분할기"""
+    def __init__(self, chunk_size=800, chunk_overlap=100):
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
     
     def split_text(self, text: str) -> List[str]:
         """텍스트를 청크로 분할"""
         chunks = []
-        start = 0
-        text_len = len(text)
+        sentences = text.split('. ')
+        current_chunk = ""
         
-        while start < text_len:
-            end = start + self.chunk_size
-            
-            # 문장 끝에서 자르기 위해 마침표 찾기
-            if end < text_len:
-                last_period = text.rfind('.', start, end)
-                if last_period != -1 and last_period > start:
-                    end = last_period + 1
-            
-            chunk = text[start:end].strip()
-            if chunk:
-                chunks.append(chunk)
-            
-            start = end - self.chunk_overlap
-            if start < 0:
-                start = 0
+        for sentence in sentences:
+            if len(current_chunk) + len(sentence) < self.chunk_size:
+                current_chunk += sentence + ". "
+            else:
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                current_chunk = sentence + ". "
+        
+        if current_chunk:
+            chunks.append(current_chunk.strip())
         
         return chunks
 
-class SimplePDFRAGSystem:
-    """간단한 PDF 기반 RAG 시스템"""
+class PDFRAGSystem:
+    """PDF 기반 RAG 시스템"""
     
     def __init__(self):
         self.documents = {}
         self.document_chunks = {}
         self.embeddings_store = {}
-        self.text_splitter = SimpleTextSplitter(chunk_size=1000, chunk_overlap=200)
+        self.text_splitter = SimpleTextSplitter(chunk_size=800, chunk_overlap=100)
         
-        # 임베딩 모델 초기화 (가능한 경우)
+        # 임베딩 모델 초기화
         self.embedder = None
         if ST_AVAILABLE:
             try:
                 self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-                print("Embedding model loaded successfully")
+                print("✅ 임베딩 모델 로드 성공")
             except Exception as e:
-                print(f"Failed to load embedding model: {e}")
+                print(f"⚠️ 임베딩 모델 로드 실패: {e}")
     
     def extract_text_from_pdf(self, pdf_path: str) -> Dict[str, Any]:
         """PDF에서 텍스트 추출"""
         if not PDF_AVAILABLE:
-            # PyMuPDF가 없는 경우 대체 방법
             return {
                 "metadata": {
                     "title": "PDF Reader Not Available",
                     "file_name": os.path.basename(pdf_path),
                     "pages": 0
                 },
-                "full_text": "PDF 처리 라이브러리가 설치되지 않았습니다. 'pip install pymupdf'를 실행해주세요."
+                "full_text": "PDF 처리를 위해 'pip install pymupdf'를 실행해주세요."
             }
         
         try:
             doc = fitz.open(pdf_path)
             text_content = []
             metadata = {
-                "title": doc.metadata.get("title", "Untitled"),
-                "author": doc.metadata.get("author", "Unknown"),
+                "title": doc.metadata.get("title", os.path.basename(pdf_path)),
                 "pages": len(doc),
                 "file_name": os.path.basename(pdf_path)
             }
@@ -184,7 +171,7 @@ class SimplePDFRAGSystem:
             # 청크 저장
             self.document_chunks[doc_id] = chunks
             
-            # 임베딩 생성 (가능한 경우)
+            # 임베딩 생성
             if self.embedder:
                 embeddings = self.embedder.encode(chunks)
                 self.embeddings_store[doc_id] = embeddings
@@ -193,8 +180,7 @@ class SimplePDFRAGSystem:
             self.documents[doc_id] = {
                 "metadata": pdf_data["metadata"],
                 "chunk_count": len(chunks),
-                "upload_time": datetime.now().isoformat(),
-                "full_text": pdf_data["full_text"][:500]  # 처음 500자 저장
+                "upload_time": datetime.now().isoformat()
             }
             
             return {
@@ -206,13 +192,10 @@ class SimplePDFRAGSystem:
             }
             
         except Exception as e:
-            return {
-                "success": False,
-                "error": str(e)
-            }
+            return {"success": False, "error": str(e)}
     
-    def search_relevant_chunks(self, query: str, doc_ids: List[str], top_k: int = 5) -> List[Dict]:
-        """쿼리와 관련된 청크 검색"""
+    def search_relevant_chunks(self, query: str, doc_ids: List[str], top_k: int = 3) -> List[Dict]:
+        """관련 청크 검색"""
         all_relevant_chunks = []
         
         if self.embedder and self.embeddings_store:
@@ -230,79 +213,75 @@ class SimplePDFRAGSystem:
                         sim = np.dot(query_embedding, emb) / (np.linalg.norm(query_embedding) * np.linalg.norm(emb))
                         similarities.append(sim)
                     
-                    # 상위 k개 선택
+                    # 상위 청크 선택
                     top_indices = np.argsort(similarities)[-top_k:][::-1]
                     
                     for idx in top_indices:
-                        all_relevant_chunks.append({
-                            "content": chunks[idx],
-                            "doc_id": doc_id,
-                            "doc_name": self.documents[doc_id]["metadata"]["file_name"],
-                            "chunk_index": idx,
-                            "similarity": similarities[idx]
-                        })
+                        if similarities[idx] > 0.2:
+                            all_relevant_chunks.append({
+                                "content": chunks[idx],
+                                "doc_name": self.documents[doc_id]["metadata"]["file_name"],
+                                "similarity": similarities[idx]
+                            })
         else:
-            # 임베딩이 없는 경우 키워드 기반 간단한 검색
-            query_lower = query.lower()
-            query_words = set(query_lower.split())
+            # 키워드 기반 검색
+            query_keywords = set(query.lower().split())
             
             for doc_id in doc_ids:
                 if doc_id in self.document_chunks:
                     chunks = self.document_chunks[doc_id]
-                    for idx, chunk in enumerate(chunks):
+                    for chunk in chunks[:top_k]:  # 처음 몇 개만 사용
                         chunk_lower = chunk.lower()
-                        # 쿼리 단어가 청크에 포함되어 있는지 확인
-                        matching_words = sum(1 for word in query_words if word in chunk_lower)
-                        if matching_words > 0:
+                        score = sum(1 for keyword in query_keywords if keyword in chunk_lower)
+                        if score > 0:
                             all_relevant_chunks.append({
-                                "content": chunk,
-                                "doc_id": doc_id,
+                                "content": chunk[:500],  # 길이 제한
                                 "doc_name": self.documents[doc_id]["metadata"]["file_name"],
-                                "chunk_index": idx,
-                                "similarity": matching_words / len(query_words)
+                                "similarity": score / len(query_keywords) if query_keywords else 0
                             })
         
-        # 유사도 순으로 정렬하고 상위 k개 반환
+        # 정렬 및 반환
         all_relevant_chunks.sort(key=lambda x: x.get('similarity', 0), reverse=True)
         return all_relevant_chunks[:top_k]
     
-    def generate_context_prompt(self, query: str, chunks: List[Dict]) -> str:
-        """컨텍스트를 포함한 프롬프트 생성"""
-        if not chunks:
+    def create_rag_prompt(self, query: str, doc_ids: List[str], top_k: int = 3) -> str:
+        """RAG 프롬프트 생성"""
+        relevant_chunks = self.search_relevant_chunks(query, doc_ids, top_k)
+        
+        if not relevant_chunks:
             return query
         
-        context_parts = []
-        for i, chunk in enumerate(chunks, 1):
-            context_parts.append(
-                f"[문서: {chunk['doc_name']}, 섹션 {chunk['chunk_index']+1}]\n{chunk['content']}\n"
-            )
+        # 프롬프트 구성
+        prompt_parts = []
+        prompt_parts.append("다음 문서 내용을 참고하여 질문에 답변해주세요:\n")
+        prompt_parts.append("=" * 50)
         
-        context = "\n---\n".join(context_parts)
+        for i, chunk in enumerate(relevant_chunks, 1):
+            prompt_parts.append(f"\n[참고문서 {i} - {chunk['doc_name']}]")
+            content = chunk['content'][:400] if len(chunk['content']) > 400 else chunk['content']
+            prompt_parts.append(content)
+            prompt_parts.append("")
         
-        enhanced_prompt = f"""다음 문서 내용을 참고하여 질문에 답변해주세요.
-
-## 참고 문서:
-{context}
-
-## 질문:
-{query}
-
-## 답변:
-위 문서 내용을 바탕으로 질문에 대해 상세하고 정확하게 답변하겠습니다."""
+        prompt_parts.append("=" * 50)
+        prompt_parts.append(f"\n질문: {query}")
+        prompt_parts.append("\n위 참고문서를 바탕으로 자세하고 정확하게 답변해주세요:")
         
-        return enhanced_prompt
+        return "\n".join(prompt_parts)
 
 # RAG 시스템 인스턴스 생성
-rag_system = SimplePDFRAGSystem()
+rag_system = PDFRAGSystem()
 
-# State variables
+# State variable to track current model
 current_model = gr.State("openai/gpt-oss-120b")
-rag_enabled = gr.State(False)
 
 def upload_pdf(file):
     """PDF 파일 업로드 처리"""
     if file is None:
-        return gr.update(value="파일을 선택해주세요"), gr.update(choices=[]), gr.update(value=False)
+        return (
+            gr.update(value="<div class='pdf-status pdf-error'>파일을 선택해주세요</div>"),
+            gr.update(choices=[]),
+            gr.update(value=False)
+        )
     
     try:
         # 파일 해시를 ID로 사용
@@ -318,175 +297,145 @@ def upload_pdf(file):
             status_html = f"""
             <div class="pdf-status pdf-success">
                 ✅ PDF 업로드 성공!<br>
-                📄 제목: {result.get('title', 'Unknown')}<br>
+                📄 파일: {result['title']}<br>
                 📑 페이지: {result['pages']}페이지<br>
-                🔍 생성된 청크: {result['chunks']}개<br>
-                🆔 문서 ID: {doc_id}
+                🔍 청크: {result['chunks']}개 생성
             </div>
             """
             
             # 문서 목록 업데이트
-            doc_list = list(rag_system.documents.keys())
             doc_choices = [f"{doc_id}: {rag_system.documents[doc_id]['metadata']['file_name']}" 
-                          for doc_id in doc_list]
+                          for doc_id in rag_system.documents.keys()]
             
-            return status_html, gr.update(choices=doc_choices, value=doc_choices), gr.update(value=True)
+            return (
+                status_html,
+                gr.update(choices=doc_choices, value=doc_choices),
+                gr.update(value=True)
+            )
         else:
             status_html = f"""
             <div class="pdf-status pdf-error">
-                ❌ PDF 업로드 실패<br>
-                오류: {result['error']}
+                ❌ 업로드 실패: {result['error']}
             </div>
             """
-            return status_html, gr.update(choices=[]), gr.update(value=False)
+            return status_html, gr.update(), gr.update(value=False)
             
     except Exception as e:
-        status_html = f"""
-        <div class="pdf-status pdf-error">
-            ❌ 오류 발생: {str(e)}
-        </div>
-        """
-        return status_html, gr.update(choices=[]), gr.update(value=False)
+        return (
+            f"<div class='pdf-status pdf-error'>❌ 오류: {str(e)}</div>",
+            gr.update(),
+            gr.update(value=False)
+        )
 
 def clear_documents():
-    """업로드된 문서 초기화"""
-    try:
-        rag_system.documents = {}
-        rag_system.document_chunks = {}
-        rag_system.embeddings_store = {}
-        
-        return gr.update(value="<div class='pdf-status pdf-success'>✅ 모든 문서가 삭제되었습니다</div>"), gr.update(choices=[], value=[]), gr.update(value=False)
-    except Exception as e:
-        return gr.update(value=f"<div class='pdf-status pdf-error'>❌ 삭제 실패: {str(e)}</div>"), gr.update(), gr.update()
+    """문서 초기화"""
+    rag_system.documents = {}
+    rag_system.document_chunks = {}
+    rag_system.embeddings_store = {}
+    
+    return (
+        gr.update(value="<div class='pdf-status pdf-success'>✅ 모든 문서가 삭제되었습니다</div>"),
+        gr.update(choices=[], value=[]),
+        gr.update(value=False)
+    )
 
 def switch_model(model_choice):
-    """모델 전환 함수"""
-    if model_choice == "openai/gpt-oss-120b":
-        return gr.update(visible=True), gr.update(visible=False), model_choice
-    else:
-        return gr.update(visible=False), gr.update(visible=True), model_choice
+    """Function to switch between models"""
+    return gr.update(visible=False), gr.update(visible=True), model_choice
 
-def chat_with_model(message: str, history: List[Tuple[str, str]], enable_rag: bool, selected_docs: List[str], top_k: int, model: str):
-    """모델과 대화 (RAG 포함)"""
+def create_rag_context_display(query, selected_docs, top_k):
+    """RAG 컨텍스트 표시용 HTML 생성"""
+    if not selected_docs:
+        return ""
     
-    # RAG가 활성화되고 문서가 선택된 경우
-    if enable_rag and selected_docs:
-        # 선택된 문서 ID 추출
-        doc_ids = [doc.split(":")[0] for doc in selected_docs]
-        
-        # 관련 청크 검색
-        relevant_chunks = rag_system.search_relevant_chunks(message, doc_ids, top_k)
-        
-        if relevant_chunks:
-            # 컨텍스트를 포함한 프롬프트 생성
-            enhanced_message = rag_system.generate_context_prompt(message, relevant_chunks)
-            
-            # 디버그 정보 포함 응답 (실제 구현시 모델 API 호출로 대체)
-            response = f"""📚 RAG 기반 답변 (모델: {model})
-
-찾은 관련 문서 섹션: {len(relevant_chunks)}개
-
-질문: {message}
-
-답변:
-{enhanced_message[:2000]}...
-
-[참고: 실제 구현시 여기서 모델 API를 호출하여 enhanced_message를 전송하고 응답을 받아야 합니다]
-
-관련 문서 섹션 요약:
-"""
-            for i, chunk in enumerate(relevant_chunks[:3], 1):
-                response += f"\n{i}. {chunk['doc_name']} - 섹션 {chunk['chunk_index']+1} (유사도: {chunk['similarity']:.2f})"
-                response += f"\n   내용: {chunk['content'][:200]}...\n"
-        else:
-            response = f"⚠️ 선택된 문서에서 '{message}'와 관련된 내용을 찾을 수 없습니다. 다른 질문을 시도해보세요."
-    else:
-        # RAG 비활성화 상태
-        response = f"""일반 답변 모드 (모델: {model})
-
-질문: {message}
-
-[참고: 실제 구현시 여기서 모델 API를 호출하여 message를 전송하고 응답을 받아야 합니다]
-
-PDF 문서를 업로드하고 RAG를 활성화하면 문서 기반 답변을 받을 수 있습니다."""
+    doc_ids = [doc.split(":")[0] for doc in selected_docs]
+    chunks = rag_system.search_relevant_chunks(query, doc_ids, top_k)
+    
+    if not chunks:
+        return ""
     
-    history.append((message, response))
-    return "", history
+    html = "<div class='rag-context'><strong>📚 참고 문서:</strong><br>"
+    for i, chunk in enumerate(chunks, 1):
+        html += f"<br>{i}. {chunk['doc_name']} (유사도: {chunk['similarity']:.2f})<br>"
+        html += f"<small>{chunk['content'][:200]}...</small><br>"
+    html += "</div>"
+    
+    return html
 
-# Gradio 인터페이스
+# Main interface
 with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as demo:
+    # JavaScript to handle message passing
+    gr.HTML("""
+    <script>
+    function sendToModel(processedMsg) {
+        // This function would send the processed message to the model
+        console.log("Sending to model:", processedMsg);
+    }
+    </script>
+    """)
+    
     with gr.Row():
-        # 사이드바
+        # Sidebar
         with gr.Column(scale=1):
             with gr.Group(elem_classes="main-container"):
-                gr.Markdown("# 🚀 AI Chat with RAG")
+                gr.Markdown("# 🚀 Inference Provider + RAG")
                 gr.Markdown(
-                    "PDF 문서를 업로드하여 AI가 문서 내용을 참고해 답변하도록 할 수 있습니다."
+                    "OpenAI GPT-OSS models with PDF RAG support. "
+                    "Sign in with your Hugging Face account to use this API."
                 )
                 
-                # 모델 선택
+                # Model selection
                 model_dropdown = gr.Dropdown(
                     choices=["openai/gpt-oss-120b", "openai/gpt-oss-20b"],
                     value="openai/gpt-oss-120b",
-                    label="📊 모델 선택"
+                    label="📊 Select Model",
+                    info="Choose between different model sizes"
                 )
                 
+                # Login button
                 login_button = gr.LoginButton("Sign in with Hugging Face", size="lg")
-                reload_btn = gr.Button("🔄 모델 변경 적용", variant="primary", size="lg")
                 
-                # RAG 설정
-                with gr.Accordion("📚 PDF RAG 설정", open=True):
+                # Reload button to apply model change
+                reload_btn = gr.Button("🔄 Apply Model Change", variant="primary", size="lg")
+                
+                # RAG Settings
+                with gr.Accordion("📚 PDF RAG Settings", open=True):
                     pdf_upload = gr.File(
-                        label="PDF 업로드",
+                        label="Upload PDF",
                         file_types=[".pdf"],
                         type="filepath"
                     )
                     
                     upload_status = gr.HTML(
-                        value="<div class='pdf-status'>PDF를 업로드하여 RAG를 활성화하세요</div>"
+                        value="<div class='pdf-status pdf-info'>📤 PDF를 업로드하여 문서 기반 답변을 받으세요</div>"
                     )
                     
                     document_list = gr.CheckboxGroup(
                         choices=[],
                         label="📄 업로드된 문서",
-                        info="질문에 참고할 문서를 선택하세요"
+                        info="참고할 문서를 선택하세요"
                     )
                     
-                    with gr.Row():
-                        clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm")
-                        refresh_btn = gr.Button("🔄 목록 새로고침", size="sm")
+                    clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm")
                     
                     enable_rag = gr.Checkbox(
                         label="RAG 활성화",
                         value=False,
-                        info="문서 기반 답변 생성 활성화"
+                        info="선택한 문서를 참고하여 답변 생성"
                     )
                     
-                    with gr.Accordion("⚙️ RAG 고급 설정", open=False):
-                        top_k_chunks = gr.Slider(
-                            minimum=1,
-                            maximum=10,
-                            value=5,
-                            step=1,
-                            label="참조할 청크 수",
-                            info="답변 생성시 참고할 문서 청크의 개수"
-                        )
-                        
-                        gr.Markdown("""
-                        ### 📝 RAG 사용 팁:
-                        1. PDF 파일을 업로드하세요
-                        2. 업로드된 문서를 선택하세요
-                        3. RAG를 활성화하세요
-                        4. 문서 내용에 대해 질문하세요
-                        
-                        예시 질문:
-                        - "문서의 주요 내용을 요약해주세요"
-                        - "이 문서에서 언급된 날짜는 언제인가요?"
-                        - "참가 자격 조건은 무엇인가요?"
-                        """)
+                    top_k_chunks = gr.Slider(
+                        minimum=1,
+                        maximum=5,
+                        value=3,
+                        step=1,
+                        label="참조 청크 수",
+                        info="답변 생성시 참고할 문서 조각 개수"
+                    )
                 
-                # 고급 옵션
-                with gr.Accordion("⚙️ 모델 설정", open=False):
+                # Additional options
+                with gr.Accordion("⚙️ Advanced Options", open=False):
+                    gr.Markdown("*These options will be available after model implementation*")
                     temperature = gr.Slider(
                         minimum=0,
                         maximum=2,
@@ -502,52 +451,75 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                         label="Max Tokens"
                     )
         
-        # 메인 채팅 영역
+        # Main chat area
         with gr.Column(scale=3):
             with gr.Group(elem_classes="main-container"):
                 gr.Markdown("## 💬 Chat Interface")
                 
                 # RAG 상태 표시
-                with gr.Row():
-                    rag_status = gr.HTML(
-                        value="<div style='padding: 10px; background: rgba(59, 130, 246, 0.1); border-radius: 8px; margin-bottom: 10px;'>🔍 RAG: <strong>비활성화</strong></div>"
-                    )
+                rag_status = gr.HTML(
+                    value="<div class='pdf-status pdf-info'>🔍 RAG: <strong>비활성화</strong></div>"
+                )
                 
-                # 모델 인터페이스 컨테이너
+                # RAG 컨텍스트 표시 영역
+                rag_context_display = gr.HTML(value="", visible=False)
+                
+                # Container for model interfaces
                 with gr.Column(visible=True) as model_120b_container:
                     gr.Markdown("### Model: openai/gpt-oss-120b")
-                    chatbot_120b = gr.Chatbot(height=400)
-                    msg_box_120b = gr.Textbox(
-                        label="메시지 입력",
-                        placeholder="PDF 내용에 대해 질문해보세요...",
-                        lines=2
-                    )
-                    with gr.Row():
-                        send_btn_120b = gr.Button("📤 전송", variant="primary")
-                        clear_btn_120b = gr.Button("🗑️ 대화 초기화")
+                    
+                    # RAG 처리를 위한 커스텀 인터페이스
+                    with gr.Group():
+                        # 사용자 입력 텍스트박스
+                        user_input = gr.Textbox(
+                            label="메시지 입력",
+                            placeholder="문서에 대해 질문하거나 일반 대화를 시작하세요...",
+                            lines=2
+                        )
+                        
+                        with gr.Row():
+                            send_btn = gr.Button("📤 전송", variant="primary")
+                            clear_chat_btn = gr.Button("🗑️ 대화 초기화")
+                        
+                        # 원본 모델 로드
+                        original_model = gr.load(
+                            "models/openai/gpt-oss-120b",
+                            accept_token=login_button,
+                            provider="fireworks-ai"
+                        )
                 
                 with gr.Column(visible=False) as model_20b_container:
                     gr.Markdown("### Model: openai/gpt-oss-20b")
-                    chatbot_20b = gr.Chatbot(height=400)
-                    msg_box_20b = gr.Textbox(
-                        label="메시지 입력",
-                        placeholder="PDF 내용에 대해 질문해보세요...",
-                        lines=2
-                    )
-                    with gr.Row():
-                        send_btn_20b = gr.Button("📤 전송", variant="primary")
-                        clear_btn_20b = gr.Button("🗑️ 대화 초기화")
+                    
+                    with gr.Group():
+                        # 사용자 입력 텍스트박스 (20b용)
+                        user_input_20b = gr.Textbox(
+                            label="메시지 입력",
+                            placeholder="문서에 대해 질문하거나 일반 대화를 시작하세요...",
+                            lines=2
+                        )
+                        
+                        with gr.Row():
+                            send_btn_20b = gr.Button("📤 전송", variant="primary")
+                            clear_chat_btn_20b = gr.Button("🗑️ 대화 초기화")
+                        
+                        # 원본 모델 로드
+                        original_model_20b = gr.load(
+                            "models/openai/gpt-oss-20b",
+                            accept_token=login_button,
+                            provider="fireworks-ai"
+                        )
     
-    # 이벤트 핸들러
+    # Event Handlers
     
-    # PDF 업로드 처리
+    # PDF 업로드
     pdf_upload.upload(
         fn=upload_pdf,
         inputs=[pdf_upload],
         outputs=[upload_status, document_list, enable_rag]
     )
     
-    # 문서 초기화
+    # 문서 삭제
     clear_btn.click(
         fn=clear_documents,
         outputs=[upload_status, document_list, enable_rag]
@@ -556,7 +528,7 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
     # RAG 상태 업데이트
     enable_rag.change(
         fn=lambda x: gr.update(
-            value=f"<div style='padding: 10px; background: rgba(59, 130, 246, 0.1); border-radius: 8px; margin-bottom: 10px;'>🔍 RAG: <strong>{'활성화' if x else '비활성화'}</strong></div>"
+            value=f"<div class='pdf-status pdf-info'>🔍 RAG: <strong>{'활성화' if x else '비활성화'}</strong></div>"
         ),
         inputs=[enable_rag],
         outputs=[rag_status]
@@ -568,46 +540,59 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
         inputs=[model_dropdown],
         outputs=[model_120b_container, model_20b_container, current_model]
     ).then(
-        fn=lambda: gr.Info("모델이 성공적으로 전환되었습니다!"),
+        fn=lambda: gr.Info("Model switched successfully!"),
         inputs=[],
         outputs=[]
     )
     
-    # 120b 모델 채팅
-    msg_box_120b.submit(
-        fn=lambda msg, hist, rag, docs, k: chat_with_model(msg, hist, rag, docs, k, "openai/gpt-oss-120b"),
-        inputs=[msg_box_120b, chatbot_120b, enable_rag, document_list, top_k_chunks],
-        outputs=[msg_box_120b, chatbot_120b]
-    )
+    # Update visibility based on dropdown selection
+    def update_visibility(model_choice):
+        if model_choice == "openai/gpt-oss-120b":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
     
-    send_btn_120b.click(
-        fn=lambda msg, hist, rag, docs, k: chat_with_model(msg, hist, rag, docs, k, "openai/gpt-oss-120b"),
-        inputs=[msg_box_120b, chatbot_120b, enable_rag, document_list, top_k_chunks],
-        outputs=[msg_box_120b, chatbot_120b]
+    model_dropdown.change(
+        fn=update_visibility,
+        inputs=[model_dropdown],
+        outputs=[model_120b_container, model_20b_container]
     )
     
-    clear_btn_120b.click(
-        lambda: ([], ""),
-        outputs=[chatbot_120b, msg_box_120b]
+    # 메시지 전송 처리 (RAG 포함)
+    def process_message(message, enable_rag, selected_docs, top_k):
+        """메시지를 RAG로 처리하여 모델에 전송"""
+        if enable_rag and selected_docs:
+            doc_ids = [doc.split(":")[0] for doc in selected_docs]
+            enhanced_message = rag_system.create_rag_prompt(message, doc_ids, top_k)
+            context_html = create_rag_context_display(message, selected_docs, top_k)
+            return enhanced_message, gr.update(value=context_html, visible=True)
+        else:
+            return message, gr.update(value="", visible=False)
+    
+    # 120b 모델용 이벤트
+    send_btn.click(
+        fn=process_message,
+        inputs=[user_input, enable_rag, document_list, top_k_chunks],
+        outputs=[user_input, rag_context_display]
     )
     
-    # 20b 모델 채팅
-    msg_box_20b.submit(
-        fn=lambda msg, hist, rag, docs, k: chat_with_model(msg, hist, rag, docs, k, "openai/gpt-oss-20b"),
-        inputs=[msg_box_20b, chatbot_20b, enable_rag, document_list, top_k_chunks],
-        outputs=[msg_box_20b, chatbot_20b]
+    user_input.submit(
+        fn=process_message,
+        inputs=[user_input, enable_rag, document_list, top_k_chunks],
+        outputs=[user_input, rag_context_display]
     )
     
+    # 20b 모델용 이벤트
     send_btn_20b.click(
-        fn=lambda msg, hist, rag, docs, k: chat_with_model(msg, hist, rag, docs, k, "openai/gpt-oss-20b"),
-        inputs=[msg_box_20b, chatbot_20b, enable_rag, document_list, top_k_chunks],
-        outputs=[msg_box_20b, chatbot_20b]
+        fn=process_message,
+        inputs=[user_input_20b, enable_rag, document_list, top_k_chunks],
+        outputs=[user_input_20b, rag_context_display]
     )
     
-    clear_btn_20b.click(
-        lambda: ([], ""),
-        outputs=[chatbot_20b, msg_box_20b]
+    user_input_20b.submit(
+        fn=process_message,
+        inputs=[user_input_20b, enable_rag, document_list, top_k_chunks],
+        outputs=[user_input_20b, rag_context_display]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()