Update app.py

app.py

@@ -4,7 +4,6 @@ from typing import List, Dict, Any, Optional, Tuple
 import hashlib
 from datetime import datetime
 import numpy as np
-from gradio_client import Client
 
 # PDF 처리 라이브러리
 try:
@@ -38,11 +37,6 @@ custom_css = """
     margin: 12px;
 }
 
-.main-container:hover {
-    box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
-    transition: all 0.3s ease;
-}
-
 /* Status messages styling */
 .pdf-status {
     padding: 12px 16px;
@@ -50,7 +44,6 @@ custom_css = """
     margin: 12px 0;
     font-size: 0.95rem;
     font-weight: 500;
-    backdrop-filter: blur(10px);
 }
 
 .pdf-success {
@@ -71,40 +64,14 @@ custom_css = """
     color: #0c5460;
 }
 
-.pdf-warning {
-    background: linear-gradient(135deg, #fff3cd 0%, #ffeeba 100%);
-    border: 1px solid #ffeaa7;
-    color: #856404;
-}
-
-/* RAG context display */
 .rag-context {
     background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
     border-left: 4px solid #f59e0b;
-    padding: 16px;
-    margin: 16px 0;
+    padding: 12px;
+    margin: 12px 0;
     border-radius: 8px;
     font-size: 0.9rem;
 }
-
-/* Chat message styling */
-.message {
-    padding: 12px 16px;
-    margin: 8px 4px;
-    border-radius: 12px;
-    max-width: 80%;
-}
-
-.user-message {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    color: white;
-    margin-left: auto;
-}
-
-.bot-message {
-    background: #f3f4f6;
-    color: #1f2937;
-}
 """
 
 class SimpleTextSplitter:
@@ -256,7 +223,7 @@ class PDFRAGSystem:
             for doc_id in doc_ids:
                 if doc_id in self.document_chunks:
                     chunks = self.document_chunks[doc_id]
-                    for chunk in chunks[:top_k]:
+                    for i, chunk in enumerate(chunks[:5]):  # 처음 5개만
                         chunk_lower = chunk.lower()
                         score = sum(1 for keyword in query_keywords if keyword in chunk_lower)
                         if score > 0:
@@ -279,33 +246,33 @@ class PDFRAGSystem:
         
         # 프롬프트 구성
         prompt_parts = []
-        prompt_parts.append("아래 참고 문서를 바탕으로 질문에 답변해주세요.\n")
-        prompt_parts.append("=" * 50)
+        prompt_parts.append("다음 문서 내용을 참고하여 답변해주세요:\n")
+        prompt_parts.append("=" * 40)
         
         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(f"\n[참고 {i} - {chunk['doc_name']}]")
+            content = chunk['content'][:300] if len(chunk['content']) > 300 else chunk['content']
             prompt_parts.append(content)
-            prompt_parts.append("")
         
-        prompt_parts.append("=" * 50)
+        prompt_parts.append("\n" + "=" * 40)
         prompt_parts.append(f"\n질문: {query}")
-        prompt_parts.append("\n위 참고문서의 내용을 바탕으로 정확하고 상세하게 답변해주세요:")
         
         return "\n".join(prompt_parts)
 
 # RAG 시스템 인스턴스 생성
 rag_system = PDFRAGSystem()
 
-# State variables
+# State variable to track current model and RAG settings
 current_model = gr.State("openai/gpt-oss-120b")
-conversation_history = gr.State([])
+rag_enabled_state = gr.State(False)
+selected_docs_state = gr.State([])
+top_k_state = gr.State(3)
 
 def upload_pdf(file):
     """PDF 파일 업로드 처리"""
     if file is None:
         return (
-            gr.update(value="<div class='pdf-status pdf-warning'>📁 파일을 선택해주세요</div>"),
+            gr.update(value="<div class='pdf-status pdf-info'>📁 파일을 선택해주세요</div>"),
             gr.update(choices=[]),
             gr.update(value=False)
         )
@@ -323,10 +290,9 @@ def upload_pdf(file):
         if result["success"]:
             status_html = f"""
             <div class="pdf-status pdf-success">
-                ✅ PDF 업로드 완료<br>
-                📄 파일: {result['title']}<br>
-                📑 페이지: {result['pages']}페이지<br>
-                🔍 청크: {result['chunks']}개 생성
+                ✅ PDF 업로드 완료!<br>
+                📄 {result['title']}<br>
+                📑 {result['pages']} 페이지 | 🔍 {result['chunks']} 청크
             </div>
             """
             
@@ -340,12 +306,11 @@ def upload_pdf(file):
                 gr.update(value=True)
             )
         else:
-            status_html = f"""
-            <div class="pdf-status pdf-error">
-                ❌ 업로드 실패: {result['error']}
-            </div>
-            """
-            return status_html, gr.update(), gr.update(value=False)
+            return (
+                f"<div class='pdf-status pdf-error'>❌ 오류: {result['error']}</div>",
+                gr.update(),
+                gr.update(value=False)
+            )
             
     except Exception as e:
         return (
@@ -368,55 +333,32 @@ def clear_documents():
 
 def switch_model(model_choice):
     """Function to switch between models"""
-    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
+    return gr.update(visible=False), gr.update(visible=True), model_choice
 
-def chat_with_rag(message, history, model_name, enable_rag, selected_docs, top_k, temperature, max_tokens):
-    """RAG를 적용한 채팅 함수"""
-    
-    if not message:
-        return history
-    
-    # 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)
-        
-        # 디버그: RAG 적용 확인
-        print(f"RAG 적용됨 - 원본: {len(message)}자, 강화: {len(enhanced_message)}자")
-    else:
-        enhanced_message = message
-    
-    try:
-        # 여기서 실제 모델 API를 호출해야 합니다
-        # 임시로 모의 응답 생성
-        if enable_rag and selected_docs:
-            response = f"""📚 [RAG 기반 답변]
+def create_rag_wrapper(original_fn, model_name):
+    """원본 모델 함수를 RAG로 감싸는 래퍼 생성"""
+    def wrapped_fn(message, history=None):
+        # RAG 설정 가져오기
+        if rag_enabled_state.value and selected_docs_state.value:
+            doc_ids = [doc.split(":")[0] for doc in selected_docs_state.value]
+            enhanced_message = rag_system.create_rag_prompt(message, doc_ids, top_k_state.value)
             
-문서를 참고하여 답변드립니다:
-
-{enhanced_message[:500]}...
-
-[참고: 실제 모델 API 연결 필요]
-"""
-        else:
-            response = f"""💬 [일반 답변]
+            # RAG 적용 알림
+            print(f"🔍 RAG 적용: {len(message)}자 → {len(enhanced_message)}자")
             
-질문: {message}
-
-[참고: 실제 모델 API 연결 필요]
-"""
-        
-        # 대화 기록에 추가
-        history.append([message, response])
-        
-    except Exception as e:
-        response = f"❌ 오류 발생: {str(e)}"
-        history.append([message, response])
+            # 원본 모델에 강화된 메시지 전달
+            if history is not None:
+                return original_fn(enhanced_message, history)
+            else:
+                return original_fn(enhanced_message)
+        else:
+            # RAG 미적용시 원본 메시지 그대로 전달
+            if history is not None:
+                return original_fn(message, history)
+            else:
+                return original_fn(message)
     
-    return history
+    return wrapped_fn
 
 # Main interface with soft theme
 with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo:
@@ -425,50 +367,50 @@ with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo
         # Sidebar
         with gr.Column(scale=1):
             with gr.Group(elem_classes="main-container"):
-                gr.Markdown("# 🤖 AI Chat + RAG")
+                gr.Markdown("# 🚀 Inference Provider + RAG")
                 gr.Markdown(
-                    "OpenAI GPT-OSS 모델과 PDF 문서 기반 답변 시스템입니다."
+                    "OpenAI GPT-OSS models served by Cerebras API. "
+                    "Upload PDF documents for context-aware responses."
                 )
                 
-                # Login button
-                login_button = gr.LoginButton("🔐 Hugging Face 로그인", size="lg")
-                
                 # Model selection
                 model_dropdown = gr.Dropdown(
                     choices=["openai/gpt-oss-120b", "openai/gpt-oss-20b"],
                     value="openai/gpt-oss-120b",
-                    label="📊 모델 선택",
-                    info="원하는 모델 크기를 선택하세요"
+                    label="📊 Select Model",
+                    info="Choose between different model sizes"
                 )
                 
+                # Login button
+                login_button = gr.LoginButton("Sign in with Hugging Face", size="lg")
+                
                 # Reload button to apply model change
-                reload_btn = gr.Button("🔄 모델 변경 적용", variant="primary", size="lg")
+                reload_btn = gr.Button("🔄 Apply Model Change", variant="primary", size="lg")
                 
                 # RAG Settings
-                with gr.Accordion("📚 PDF RAG 설정", open=True):
+                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-info'>📁 PDF를 업로드하여 문서 기반 답변을 받으세요</div>"
+                        value="<div class='pdf-status pdf-info'>📤 Upload a PDF to enable document-based answers</div>"
                     )
                     
                     document_list = gr.CheckboxGroup(
                         choices=[],
-                        label="📄 업로드된 문서",
-                        info="참고할 문서를 선택하세요"
+                        label="📄 Uploaded Documents",
+                        info="Select documents to use as context"
                     )
                     
-                    with gr.Row():
-                        clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm", variant="secondary")
+                    clear_btn = gr.Button("🗑️ Clear All Documents", size="sm", variant="secondary")
                     
                     enable_rag = gr.Checkbox(
-                        label="✨ RAG 활성화",
+                        label="✨ Enable RAG",
                         value=False,
-                        info="선택한 문서를 참고하여 답변 생성"
+                        info="Use documents for context-aware responses"
                     )
                     
                     top_k_chunks = gr.Slider(
@@ -476,27 +418,26 @@ with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo
                         maximum=5,
                         value=3,
                         step=1,
-                        label="참조 청크 수",
-                        info="답변시 참고할 문서 조각 개수"
+                        label="Context Chunks",
+                        info="Number of document chunks to use"
                     )
                 
                 # Additional options
-                with gr.Accordion("🎛️ 모델 옵션", open=False):
+                with gr.Accordion("⚙️ Advanced Options", open=False):
+                    gr.Markdown("*These options will be available after model implementation*")
                     temperature = gr.Slider(
                         minimum=0,
                         maximum=2,
                         value=0.7,
                         step=0.1,
-                        label="Temperature",
-                        info="낮을수록 일관성 있고, 높을수록 창의적입니다"
+                        label="Temperature"
                     )
                     max_tokens = gr.Slider(
                         minimum=1,
                         maximum=4096,
                         value=512,
                         step=1,
-                        label="Max Tokens",
-                        info="생성할 최대 토큰 수"
+                        label="Max Tokens"
                     )
         
         # Main chat area
@@ -504,187 +445,114 @@ with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo
             with gr.Group(elem_classes="main-container"):
                 gr.Markdown("## 💬 Chat Interface")
                 
-                # RAG 상태 표시
+                # RAG status
                 rag_status = gr.HTML(
-                    value="<div class='pdf-status pdf-info'>🔍 RAG: <strong>비활성화</strong></div>"
+                    value="<div class='pdf-status pdf-info'>🔍 RAG: <strong>Disabled</strong></div>"
                 )
                 
-                # 통합된 채팅 인터페이스 (모델별로 하나씩)
+                # RAG context preview
+                context_preview = 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")
+                    gr.Markdown("### Model: openai/gpt-oss-120b")
                     
-                    chatbot_120b = gr.Chatbot(
-                        height=400,
-                        show_label=False,
-                        elem_classes="chatbot"
+                    # Load the original model and wrap it with RAG
+                    original_interface_120b = gr.load(
+                        "models/openai/gpt-oss-120b",
+                        accept_token=login_button,
+                        provider="fireworks-ai"
                     )
                     
-                    with gr.Row():
-                        msg_120b = gr.Textbox(
-                            placeholder="메시지를 입력하세요... (Enter로 전송)",
-                            show_label=False,
-                            scale=4,
-                            container=False
-                        )
-                        send_btn_120b = gr.Button("📤 전송", variant="primary", scale=1)
-                    
-                    with gr.Row():
-                        clear_btn_120b = gr.Button("🗑️ 대화 ��기화", variant="secondary", size="sm")
-                        
-                        # 예제 질문들
-                        gr.Examples(
-                            examples=[
-                                "문서의 주요 내용을 요약해주세요",
-                                "이 문서에서 가장 중요한 포인트는 무엇인가요?",
-                                "문서에 언급된 날짜와 일정을 알려주세요"
-                            ],
-                            inputs=msg_120b
-                        )
+                    # Note: The loaded interface will have its own chat components
+                    # We'll intercept the messages through our wrapper function
                 
                 with gr.Column(visible=False) as model_20b_container:
-                    gr.Markdown("### 🚀 Model: openai/gpt-oss-20b")
+                    gr.Markdown("### Model: openai/gpt-oss-20b")
                     
-                    chatbot_20b = gr.Chatbot(
-                        height=400,
-                        show_label=False,
-                        elem_classes="chatbot"
+                    # Load the original model
+                    original_interface_20b = gr.load(
+                        "models/openai/gpt-oss-20b",
+                        accept_token=login_button,
+                        provider="fireworks-ai"
                     )
-                    
-                    with gr.Row():
-                        msg_20b = gr.Textbox(
-                            placeholder="메시지를 입력하세요... (Enter로 전송)",
-                            show_label=False,
-                            scale=4,
-                            container=False
-                        )
-                        send_btn_20b = gr.Button("📤 전송", variant="primary", scale=1)
-                    
-                    with gr.Row():
-                        clear_btn_20b = gr.Button("🗑️ 대화 초기화", variant="secondary", size="sm")
-                        
-                        # 예제 질문들
-                        gr.Examples(
-                            examples=[
-                                "문서의 주요 내용을 요약해주세요",
-                                "이 문서에서 가장 중요한 포인트는 무엇인가요?",
-                                "문서에 언급된 날짜와 일정을 알려주세요"
-                            ],
-                            inputs=msg_20b
-                        )
     
     # Event Handlers
     
-    # PDF 업로드
+    # PDF upload
     pdf_upload.upload(
         fn=upload_pdf,
         inputs=[pdf_upload],
         outputs=[upload_status, document_list, enable_rag]
     )
     
-    # 문서 삭제
+    # Clear documents
     clear_btn.click(
         fn=clear_documents,
         outputs=[upload_status, document_list, enable_rag]
     )
     
-    # RAG 상태 업데이트
+    # Update RAG state when settings change
+    def update_rag_state(enabled, docs, k):
+        rag_enabled_state.value = enabled
+        selected_docs_state.value = docs if docs else []
+        top_k_state.value = k
+        
+        status = "✅ Enabled" if enabled and docs else "⭕ Disabled"
+        status_html = f"<div class='pdf-status pdf-info'>🔍 RAG: <strong>{status}</strong></div>"
+        
+        # Show context preview if RAG is enabled
+        if enabled and docs:
+            preview = f"<div class='rag-context'>📚 Using {len(docs)} document(s) with {k} chunks per query</div>"
+            return gr.update(value=status_html), gr.update(value=preview, visible=True)
+        else:
+            return gr.update(value=status_html), gr.update(value="", visible=False)
+    
+    # Connect RAG state updates
     enable_rag.change(
-        fn=lambda x: gr.update(
-            value=f"<div class='pdf-status pdf-info'>🔍 RAG: <strong>{'✅ 활성화' if x else '⭕ 비활성화'}</strong></div>"
-        ),
-        inputs=[enable_rag],
-        outputs=[rag_status]
+        fn=update_rag_state,
+        inputs=[enable_rag, document_list, top_k_chunks],
+        outputs=[rag_status, context_preview]
+    )
+    
+    document_list.change(
+        fn=update_rag_state,
+        inputs=[enable_rag, document_list, top_k_chunks],
+        outputs=[rag_status, context_preview]
+    )
+    
+    top_k_chunks.change(
+        fn=update_rag_state,
+        inputs=[enable_rag, document_list, top_k_chunks],
+        outputs=[rag_status, context_preview]
     )
     
-    # 모델 전환
+    # Handle model switching
     reload_btn.click(
         fn=switch_model,
         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_120b.submit(
-        fn=lambda msg, hist: chat_with_rag(
-            msg, hist, "openai/gpt-oss-120b", 
-            enable_rag.value, document_list.value, top_k_chunks.value,
-            temperature.value, max_tokens.value
-        ),
-        inputs=[msg_120b, chatbot_120b],
-        outputs=[chatbot_120b]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg_120b]
-    )
-    
-    send_btn_120b.click(
-        fn=lambda msg, hist: chat_with_rag(
-            msg, hist, "openai/gpt-oss-120b",
-            enable_rag.value, document_list.value, top_k_chunks.value,
-            temperature.value, max_tokens.value
-        ),
-        inputs=[msg_120b, chatbot_120b],
-        outputs=[chatbot_120b]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg_120b]
-    )
-    
-    clear_btn_120b.click(
-        fn=lambda: [],
-        outputs=[chatbot_120b]
-    )
-    
-    # 20b 모델 채팅 이벤트
-    msg_20b.submit(
-        fn=lambda msg, hist: chat_with_rag(
-            msg, hist, "openai/gpt-oss-20b",
-            enable_rag.value, document_list.value, top_k_chunks.value,
-            temperature.value, max_tokens.value
-        ),
-        inputs=[msg_20b, chatbot_20b],
-        outputs=[chatbot_20b]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg_20b]
-    )
-    
-    send_btn_20b.click(
-        fn=lambda msg, hist: chat_with_rag(
-            msg, hist, "openai/gpt-oss-20b",
-            enable_rag.value, document_list.value, top_k_chunks.value,
-            temperature.value, max_tokens.value
-        ),
-        inputs=[msg_20b, chatbot_20b],
-        outputs=[chatbot_20b]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg_20b]
-    )
+    # 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)
     
-    clear_btn_20b.click(
-        fn=lambda: [],
-        outputs=[chatbot_20b]
+    model_dropdown.change(
+        fn=update_visibility,
+        inputs=[model_dropdown],
+        outputs=[model_120b_container, model_20b_container]
     )
-
-# 실제 모델 API 연결을 위한 함수 (구현 필요)
-def connect_to_model_api(model_name, message, temperature, max_tokens):
-    """
-    실제 모델 API에 연결하는 함수
-    TODO: 여기에 실제 API 호출 코드를 구현해야 합니다
     
-    예시:
-    - OpenAI API
-    - Hugging Face Inference API
-    - Custom model endpoint
-    """
-    # client = Client(f"models/{model_name}")
-    # response = client.predict(message, temperature=temperature, max_tokens=max_tokens)
-    # return response
-    pass
+    # Monkey-patch the loaded interfaces to add RAG support
+    # This is done after the interface is loaded
+    demo.load = lambda: print("📚 RAG System Ready!")
 
 demo.launch()