Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import os
-from typing import List, Dict, Any, Optional
+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:
@@ -20,64 +21,89 @@ except ImportError:
     ST_AVAILABLE = False
     print("⚠️ Sentence Transformers not installed. Install with: pip install sentence-transformers")
 
-# Custom CSS for gradient background and styling
+# Soft and bright custom CSS
 custom_css = """
 .gradio-container {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #4facfe 75%, #00f2fe 100%);
-    background-size: 400% 400%;
-    animation: gradient-animation 15s ease infinite;
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
     min-height: 100vh;
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
 }
-@keyframes gradient-animation {
-    0% { background-position: 0% 50%; }
-    50% { background-position: 100% 50%; }
-    100% { background-position: 0% 50%; }
-}
-.dark .gradio-container {
-    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 25%, #0f3460 50%, #533483 75%, #e94560 100%);
-    background-size: 400% 400%;
-    animation: gradient-animation 15s ease infinite;
-}
+
 .main-container {
-    background-color: rgba(255, 255, 255, 0.95);
-    backdrop-filter: blur(10px);
-    border-radius: 20px;
-    padding: 20px;
-    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
-    border: 1px solid rgba(255, 255, 255, 0.18);
-    margin: 10px;
+    background: rgba(255, 255, 255, 0.98);
+    border-radius: 16px;
+    padding: 24px;
+    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
+    border: 1px solid rgba(0, 0, 0, 0.05);
+    margin: 12px;
 }
-.dark .main-container {
-    background-color: rgba(30, 30, 30, 0.95);
-    border: 1px solid rgba(255, 255, 255, 0.1);
+
+.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: 10px;
-    border-radius: 10px;
-    margin: 10px 0;
-    font-size: 0.9em;
+    padding: 12px 16px;
+    border-radius: 12px;
+    margin: 12px 0;
+    font-size: 0.95rem;
+    font-weight: 500;
+    backdrop-filter: blur(10px);
 }
+
 .pdf-success {
-    background-color: rgba(52, 211, 153, 0.2);
-    border: 1px solid rgba(52, 211, 153, 0.5);
-    color: #10b981;
+    background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%);
+    border: 1px solid #b1dfbb;
+    color: #155724;
 }
+
 .pdf-error {
-    background-color: rgba(248, 113, 113, 0.2);
-    border: 1px solid rgba(248, 113, 113, 0.5);
-    color: #ef4444;
+    background: linear-gradient(135deg, #f8d7da 0%, #f5c6cb 100%);
+    border: 1px solid #f1aeb5;
+    color: #721c24;
 }
+
 .pdf-info {
-    background-color: rgba(59, 130, 246, 0.2);
-    border: 1px solid rgba(59, 130, 246, 0.5);
-    color: #3b82f6;
+    background: linear-gradient(135deg, #d1ecf1 0%, #bee5eb 100%);
+    border: 1px solid #9ec5d8;
+    color: #0c5460;
+}
+
+.pdf-warning {
+    background: linear-gradient(135deg, #fff3cd 0%, #ffeeba 100%);
+    border: 1px solid #ffeaa7;
+    color: #856404;
 }
+
+/* RAG context display */
 .rag-context {
-    background-color: rgba(251, 191, 36, 0.1);
+    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
     border-left: 4px solid #f59e0b;
-    padding: 10px;
-    margin: 10px 0;
-    border-radius: 5px;
+    padding: 16px;
+    margin: 16px 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;
 }
 """
 
@@ -230,12 +256,12 @@ 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 chunk in chunks[:top_k]:
                         chunk_lower = chunk.lower()
                         score = sum(1 for keyword in query_keywords if keyword in chunk_lower)
                         if score > 0:
                             all_relevant_chunks.append({
-                                "content": chunk[:500],  # 길이 제한
+                                "content": chunk[:500],
                                 "doc_name": self.documents[doc_id]["metadata"]["file_name"],
                                 "similarity": score / len(query_keywords) if query_keywords else 0
                             })
@@ -253,7 +279,7 @@ class PDFRAGSystem:
         
         # 프롬프트 구성
         prompt_parts = []
-        prompt_parts.append("다음 문서 내용을 참고하여 질문에 답변해주세요:\n")
+        prompt_parts.append("아래 참고 문서를 바탕으로 질문에 답변해주세요.\n")
         prompt_parts.append("=" * 50)
         
         for i, chunk in enumerate(relevant_chunks, 1):
@@ -264,21 +290,22 @@ class PDFRAGSystem:
         
         prompt_parts.append("=" * 50)
         prompt_parts.append(f"\n질문: {query}")
-        prompt_parts.append("\n위 참고문서를 바탕으로 자세하고 정확하게 답변해주세요:")
+        prompt_parts.append("\n위 참고문서의 내용을 바탕으로 정확하고 상세하게 답변해주세요:")
         
         return "\n".join(prompt_parts)
 
 # RAG 시스템 인스턴스 생성
 rag_system = PDFRAGSystem()
 
-# State variable to track current model
+# State variables
 current_model = gr.State("openai/gpt-oss-120b")
+conversation_history = gr.State([])
 
 def upload_pdf(file):
     """PDF 파일 업로드 처리"""
     if file is None:
         return (
-            gr.update(value="<div class='pdf-status pdf-error'>파일을 선택해주세요</div>"),
+            gr.update(value="<div class='pdf-status pdf-warning'>📁 파일을 선택해주세요</div>"),
             gr.update(choices=[]),
             gr.update(value=False)
         )
@@ -296,7 +323,7 @@ def upload_pdf(file):
         if result["success"]:
             status_html = f"""
             <div class="pdf-status pdf-success">
-                ✅ PDF 업로드 성공!<br>
+                ✅ PDF 업로드 완료<br>
                 📄 파일: {result['title']}<br>
                 📑 페이지: {result['pages']}페이지<br>
                 🔍 청크: {result['chunks']}개 생성
@@ -334,80 +361,99 @@ def clear_documents():
     rag_system.embeddings_store = {}
     
     return (
-        gr.update(value="<div class='pdf-status pdf-success'>✅ 모든 문서가 삭제되었습니다</div>"),
+        gr.update(value="<div class='pdf-status pdf-info'>🗑️ 모든 문서가 삭제되었습니다</div>"),
         gr.update(choices=[], value=[]),
         gr.update(value=False)
     )
 
 def switch_model(model_choice):
     """Function to switch between models"""
-    return gr.update(visible=False), gr.update(visible=True), 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
 
-def create_rag_context_display(query, selected_docs, top_k):
-    """RAG 컨텍스트 표시용 HTML 생성"""
-    if not selected_docs:
-        return ""
+def chat_with_rag(message, history, model_name, enable_rag, selected_docs, top_k, temperature, max_tokens):
+    """RAG를 적용한 채팅 함수"""
     
-    doc_ids = [doc.split(":")[0] for doc in selected_docs]
-    chunks = rag_system.search_relevant_chunks(query, doc_ids, top_k)
+    if not message:
+        return history
     
-    if not chunks:
-        return ""
+    # 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
     
-    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>"
+    try:
+        # 여기서 실제 모델 API를 호출해야 합니다
+        # 임시로 모의 응답 생성
+        if enable_rag and selected_docs:
+            response = f"""📚 [RAG 기반 답변]
+            
+문서를 참고하여 답변드립니다:
+
+{enhanced_message[:500]}...
+
+[참고: 실제 모델 API 연결 필요]
+"""
+        else:
+            response = f"""💬 [일반 답변]
+            
+질문: {message}
+
+[참고: 실제 모델 API 연결 필요]
+"""
+        
+        # 대화 기록에 추가
+        history.append([message, response])
+        
+    except Exception as e:
+        response = f"❌ 오류 발생: {str(e)}"
+        history.append([message, response])
     
-    return html
+    return history
 
-# 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>
-    """)
+# Main interface with soft theme
+with gr.Blocks(fill_height=True, theme=gr.themes.Soft(), css=custom_css) as demo:
     
     with gr.Row():
         # Sidebar
         with gr.Column(scale=1):
             with gr.Group(elem_classes="main-container"):
-                gr.Markdown("# 🚀 Inference Provider + RAG")
+                gr.Markdown("# 🤖 AI Chat + RAG")
                 gr.Markdown(
-                    "OpenAI GPT-OSS models with PDF RAG support. "
-                    "Sign in with your Hugging Face account to use this API."
+                    "OpenAI GPT-OSS 모델과 PDF 문서 기반 답변 시스템입니다."
                 )
                 
+                # 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="📊 Select Model",
-                    info="Choose between different model sizes"
+                    label="📊 모델 선택",
+                    info="원하는 모델 크기를 선택하세요"
                 )
                 
-                # Login button
-                login_button = gr.LoginButton("Sign in with Hugging Face", size="lg")
-                
                 # Reload button to apply model change
-                reload_btn = gr.Button("🔄 Apply Model Change", variant="primary", size="lg")
+                reload_btn = gr.Button("🔄 모델 변경 적용", variant="primary", size="lg")
                 
                 # RAG Settings
-                with gr.Accordion("📚 PDF RAG Settings", open=True):
+                with gr.Accordion("📚 PDF RAG 설정", open=True):
                     pdf_upload = gr.File(
-                        label="Upload PDF",
+                        label="📤 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'>📁 PDF를 업로드하여 문서 기반 답변을 받으세요</div>"
                     )
                     
                     document_list = gr.CheckboxGroup(
@@ -416,10 +462,11 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                         info="참고할 문서를 선택하세요"
                     )
                     
-                    clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm")
+                    with gr.Row():
+                        clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm", variant="secondary")
                     
                     enable_rag = gr.Checkbox(
-                        label="RAG 활성화",
+                        label="✨ RAG 활성화",
                         value=False,
                         info="선택한 문서를 참고하여 답변 생성"
                     )
@@ -430,25 +477,26 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                         value=3,
                         step=1,
                         label="참조 청크 수",
-                        info="답변 생성시 참고할 문서 조각 개수"
+                        info="답변시 참고할 문서 조각 개수"
                     )
                 
                 # Additional options
-                with gr.Accordion("⚙️ Advanced Options", open=False):
-                    gr.Markdown("*These options will be available after model implementation*")
+                with gr.Accordion("🎛️ 모델 옵션", open=False):
                     temperature = gr.Slider(
                         minimum=0,
                         maximum=2,
                         value=0.7,
                         step=0.1,
-                        label="Temperature"
+                        label="Temperature",
+                        info="낮을수록 일관성 있고, 높을수록 창의적입니다"
                     )
                     max_tokens = gr.Slider(
                         minimum=1,
                         maximum=4096,
                         value=512,
                         step=1,
-                        label="Max Tokens"
+                        label="Max Tokens",
+                        info="생성할 최대 토큰 수"
                     )
         
         # Main chat area
@@ -461,53 +509,67 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                     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")
+                    gr.Markdown("### 🚀 Model: openai/gpt-oss-120b")
+                    
+                    chatbot_120b = gr.Chatbot(
+                        height=400,
+                        show_label=False,
+                        elem_classes="chatbot"
+                    )
                     
-                    # RAG 처리를 위한 커스텀 인터페이스
-                    with gr.Group():
-                        # 사용자 입력 텍스트박스
-                        user_input = gr.Textbox(
-                            label="메시지 입력",
-                            placeholder="문서에 대해 질문하거나 일반 대화를 시작하세요...",
-                            lines=2
+                    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")
                         
-                        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"
+                        # 예제 질문들
+                        gr.Examples(
+                            examples=[
+                                "문서의 주요 내용을 요약해주세요",
+                                "이 문서에서 가장 중요한 포인트는 무엇인가요?",
+                                "문서에 언급된 날짜와 일정을 알려주세요"
+                            ],
+                            inputs=msg_120b
                         )
                 
                 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"
+                    )
                     
-                    with gr.Group():
-                        # 사용자 입력 텍스트박스 (20b용)
-                        user_input_20b = gr.Textbox(
-                            label="메시지 입력",
-                            placeholder="문서에 대해 질문하거나 일반 대화를 시작하세요...",
-                            lines=2
+                    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")
                         
-                        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"
+                        # 예제 질문들
+                        gr.Examples(
+                            examples=[
+                                "문서의 주요 내용을 요약해주세요",
+                                "이 문서에서 가장 중요한 포인트는 무엇인가요?",
+                                "문서에 언급된 날짜와 일정을 알려주세요"
+                            ],
+                            inputs=msg_20b
                         )
     
     # Event Handlers
@@ -528,7 +590,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 class='pdf-status pdf-info'>🔍 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]
@@ -540,59 +602,89 @@ 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("Model switched successfully!"),
+        fn=lambda: gr.Info("✅ 모델이 성공적으로 전환되었습니다!"),
         inputs=[],
         outputs=[]
     )
     
-    # 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)
-    
-    model_dropdown.change(
-        fn=update_visibility,
-        inputs=[model_dropdown],
-        outputs=[model_120b_container, model_20b_container]
+    # 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]
     )
     
-    # 메시지 전송 처리 (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)
+    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]
+    )
     
-    # 120b 모델용 이벤트
-    send_btn.click(
-        fn=process_message,
-        inputs=[user_input, enable_rag, document_list, top_k_chunks],
-        outputs=[user_input, rag_context_display]
+    clear_btn_120b.click(
+        fn=lambda: [],
+        outputs=[chatbot_120b]
     )
     
-    user_input.submit(
-        fn=process_message,
-        inputs=[user_input, enable_rag, document_list, top_k_chunks],
-        outputs=[user_input, rag_context_display]
+    # 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]
     )
     
-    # 20b 모델용 이벤트
     send_btn_20b.click(
-        fn=process_message,
-        inputs=[user_input_20b, enable_rag, document_list, top_k_chunks],
-        outputs=[user_input_20b, rag_context_display]
+        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]
     )
     
-    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]
+    clear_btn_20b.click(
+        fn=lambda: [],
+        outputs=[chatbot_20b]
     )
 
+# 실제 모델 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
+
 demo.launch()