Update app.py

app.py

@@ -1,6 +1,19 @@
 import gradio as gr
+import os
+from typing import List, Dict, Any, Optional
+import hashlib
+import json
+from datetime import datetime
 
-# Custom CSS for gradient background and styling
+# PDF 처리 라이브러리
+import pymupdf  # PyMuPDF
+import chromadb
+from chromadb.utils import embedding_functions
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+import numpy as np
+
+# Custom CSS (기존 CSS + 추가 스타일)
 custom_css = """
 .gradio-container {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #4facfe 75%, #00f2fe 100%);
@@ -8,20 +21,16 @@ custom_css = """
     animation: gradient-animation 15s ease infinite;
     min-height: 100vh;
 }
-
 @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;
 }
-
-/* Style for content areas */
 .main-container {
     background-color: rgba(255, 255, 255, 0.95);
     backdrop-filter: blur(10px);
@@ -31,51 +40,380 @@ custom_css = """
     border: 1px solid rgba(255, 255, 255, 0.18);
     margin: 10px;
 }
-
 .dark .main-container {
     background-color: rgba(30, 30, 30, 0.95);
     border: 1px solid rgba(255, 255, 255, 0.1);
 }
+.pdf-status {
+    padding: 10px;
+    border-radius: 10px;
+    margin: 10px 0;
+    font-size: 0.9em;
+}
+.pdf-success {
+    background-color: rgba(52, 211, 153, 0.2);
+    border: 1px solid rgba(52, 211, 153, 0.5);
+    color: #10b981;
+}
+.pdf-error {
+    background-color: rgba(248, 113, 113, 0.2);
+    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;
+}
+.document-card {
+    padding: 12px;
+    margin: 8px 0;
+    border-radius: 8px;
+    background: rgba(255, 255, 255, 0.1);
+    border: 1px solid rgba(255, 255, 255, 0.2);
+    cursor: pointer;
+    transition: all 0.3s ease;
+}
+.document-card:hover {
+    background: rgba(255, 255, 255, 0.2);
+    transform: translateX(5px);
+}
 """
 
-# State variable to track current model
+class PDFRAGSystem:
+    """PDF 기반 RAG 시스템 클래스"""
+    
+    def __init__(self):
+        self.documents = {}
+        self.embedder = None
+        self.vector_store = None
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
+        )
+        self.initialize_vector_store()
+    
+    def initialize_vector_store(self):
+        """벡터 저장소 초기화"""
+        try:
+            # Sentence Transformer 모델 로드
+            self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+            
+            # ChromaDB 클라이언트 초기화
+            self.chroma_client = chromadb.Client()
+            self.collection = self.chroma_client.create_collection(
+                name="pdf_documents",
+                metadata={"hnsw:space": "cosine"}
+            )
+        except Exception as e:
+            print(f"Vector store initialization error: {e}")
+    
+    def extract_text_from_pdf(self, pdf_path: str) -> Dict[str, Any]:
+        """PDF에서 텍스트 추출"""
+        try:
+            doc = pymupdf.open(pdf_path)
+            text_content = []
+            metadata = {
+                "title": doc.metadata.get("title", "Untitled"),
+                "author": doc.metadata.get("author", "Unknown"),
+                "pages": len(doc),
+                "creation_date": doc.metadata.get("creationDate", ""),
+                "file_name": os.path.basename(pdf_path)
+            }
+            
+            for page_num, page in enumerate(doc):
+                text = page.get_text()
+                if text.strip():
+                    text_content.append({
+                        "page": page_num + 1,
+                        "content": text
+                    })
+            
+            doc.close()
+            
+            return {
+                "metadata": metadata,
+                "pages": text_content,
+                "full_text": "\n\n".join([p["content"] for p in text_content])
+            }
+        except Exception as e:
+            raise Exception(f"PDF 처리 오류: {str(e)}")
+    
+    def process_and_index_pdf(self, pdf_path: str, doc_id: str) -> Dict[str, Any]:
+        """PDF 처리 및 벡터 인덱싱"""
+        try:
+            # PDF 텍스트 추출
+            pdf_data = self.extract_text_from_pdf(pdf_path)
+            
+            # 텍스트를 청크로 분할
+            chunks = self.text_splitter.split_text(pdf_data["full_text"])
+            
+            # 각 청크에 대한 임베딩 생성
+            embeddings = self.embedder.encode(chunks)
+            
+            # ChromaDB에 저장
+            ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
+            metadatas = [
+                {
+                    "doc_id": doc_id,
+                    "chunk_index": i,
+                    "source": pdf_data["metadata"]["file_name"],
+                    "page_count": pdf_data["metadata"]["pages"]
+                }
+                for i in range(len(chunks))
+            ]
+            
+            self.collection.add(
+                ids=ids,
+                embeddings=embeddings.tolist(),
+                documents=chunks,
+                metadatas=metadatas
+            )
+            
+            # 문서 정보 저장
+            self.documents[doc_id] = {
+                "metadata": pdf_data["metadata"],
+                "chunk_count": len(chunks),
+                "upload_time": datetime.now().isoformat()
+            }
+            
+            return {
+                "success": True,
+                "doc_id": doc_id,
+                "chunks": len(chunks),
+                "pages": pdf_data["metadata"]["pages"],
+                "title": pdf_data["metadata"]["title"]
+            }
+            
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e)
+            }
+    
+    def search_relevant_chunks(self, query: str, top_k: int = 5) -> List[Dict]:
+        """쿼리와 관련된 청크 검색"""
+        try:
+            # 쿼리 임베딩 생성
+            query_embedding = self.embedder.encode([query])
+            
+            # 유사한 문서 검색
+            results = self.collection.query(
+                query_embeddings=query_embedding.tolist(),
+                n_results=top_k
+            )
+            
+            if results and results['documents']:
+                chunks = []
+                for i in range(len(results['documents'][0])):
+                    chunks.append({
+                        "content": results['documents'][0][i],
+                        "metadata": results['metadatas'][0][i],
+                        "distance": results['distances'][0][i] if 'distances' in results else None
+                    })
+                return chunks
+            return []
+            
+        except Exception as e:
+            print(f"Search error: {e}")
+            return []
+    
+    def generate_rag_prompt(self, query: str, context_chunks: List[Dict]) -> str:
+        """RAG 프롬프트 생성"""
+        context = "\n\n---\n\n".join([
+            f"[출처: {chunk['metadata']['source']}, 청크 {chunk['metadata']['chunk_index']+1}]\n{chunk['content']}"
+            for chunk in context_chunks
+        ])
+        
+        prompt = f"""다음 문서 내용을 참고하여 질문에 답변해주세요. 
+답변은 제공된 문서 내용을 바탕으로 작성하되, 필요시 추가 설명을 포함할 수 있습니다.
+문서에서 관련 정보를 찾을 수 없는 경우, 그 사실을 명시해주세요.
+
+📚 참고 문서:
+{context}
+
+❓ 질문: {query}
+
+💡 답변:"""
+        
+        return prompt
+
+# RAG 시스템 인스턴스 생성
+rag_system = PDFRAGSystem()
+
+# State variables
 current_model = gr.State("openai/gpt-oss-120b")
+uploaded_documents = gr.State({})
+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)
+    
+    try:
+        # 파일 해시를 ID로 사용
+        with open(file.name, 'rb') as f:
+            file_hash = hashlib.md5(f.read()).hexdigest()[:8]
+        
+        doc_id = f"doc_{file_hash}"
+        
+        # PDF 처리 및 인덱싱
+        result = rag_system.process_and_index_pdf(file.name, doc_id)
+        
+        if result["success"]:
+            status_html = f"""
+            <div class="pdf-status pdf-success">
+                ✅ PDF 업로드 성공!<br>
+                📄 제목: {result.get('title', 'Unknown')}<br>
+                📑 페이지: {result['pages']}페이지<br>
+                🔍 생성된 청크: {result['chunks']}개<br>
+                🆔 문서 ID: {doc_id}
+            </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]
+            
+            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']}
+            </div>
+            """
+            return status_html, gr.update(choices=[]), 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)
+
+def clear_documents():
+    """업로드된 문서 초기화"""
+    try:
+        # ChromaDB 컬렉션 재생성
+        rag_system.chroma_client.delete_collection("pdf_documents")
+        rag_system.collection = rag_system.chroma_client.create_collection(
+            name="pdf_documents",
+            metadata={"hnsw:space": "cosine"}
+        )
+        rag_system.documents = {}
+        
+        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()
+
+def process_with_rag(message: str, enable_rag: bool, selected_docs: List[str], top_k: int = 5):
+    """RAG를 활용한 메시지 처리"""
+    if not enable_rag or not selected_docs:
+        return message  # RAG 비활성화시 원본 메시지 반환
+    
+    try:
+        # 관련 청크 검색
+        relevant_chunks = rag_system.search_relevant_chunks(message, top_k=top_k)
+        
+        if relevant_chunks:
+            # 선택된 문서의 청크만 필터링
+            selected_doc_ids = [doc.split(":")[0] for doc in selected_docs]
+            filtered_chunks = [
+                chunk for chunk in relevant_chunks 
+                if chunk['metadata']['doc_id'] in selected_doc_ids
+            ]
+            
+            if filtered_chunks:
+                # RAG 프롬프트 생성
+                rag_prompt = rag_system.generate_rag_prompt(message, filtered_chunks[:top_k])
+                return rag_prompt
+        
+        return message
+        
+    except Exception as e:
+        print(f"RAG processing error: {e}")
+        return message
 
 def switch_model(model_choice):
-    """Function to switch between models"""
+    """모델 전환 함수"""
     return gr.update(visible=False), gr.update(visible=True), model_choice
 
+# Gradio 인터페이스
 with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", 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")
+                gr.Markdown("# 🚀 AI Chat with RAG")
                 gr.Markdown(
-                    "This Space showcases OpenAI GPT-OSS models, served by the Cerebras API. "
-                    "Sign in with your Hugging Face account to use this API."
+                    "PDF 문서를 업로드하여 AI가 문서 내용을 참고해 답변하도록 할 수 있습니다."
                 )
                 
-                # Model selection
+                # 모델 선택
                 model_dropdown = gr.Dropdown(
-                    choices=[
-                        "openai/gpt-oss-120b",
-                        "openai/gpt-oss-20b"
-                    ],
+                    choices=["openai/gpt-oss-120b", "openai/gpt-oss-20b"],
                     value="openai/gpt-oss-120b",
-                    label="📊 Select Model",
-                    info="Choose between different model sizes"
+                    label="📊 모델 선택"
                 )
                 
-                # Login button
                 login_button = gr.LoginButton("Sign in with Hugging Face", size="lg")
+                reload_btn = gr.Button("🔄 모델 변경 적용", variant="primary", size="lg")
                 
-                # Reload button to apply model change
-                reload_btn = gr.Button("🔄 Apply Model Change", variant="primary", size="lg")
+                # RAG 설정
+                with gr.Accordion("📚 PDF RAG 설정", open=True):
+                    pdf_upload = gr.File(
+                        label="PDF 업로드",
+                        file_types=[".pdf"],
+                        type="filepath"
+                    )
+                    
+                    upload_status = gr.HTML(
+                        value="<div class='pdf-status'>PDF를 업로드하여 RAG를 활성화하세요</div>"
+                    )
+                    
+                    document_list = gr.CheckboxGroup(
+                        choices=[],
+                        label="📄 업로드된 문서",
+                        info="질문에 참고할 문서를 선택하세요"
+                    )
+                    
+                    with gr.Row():
+                        clear_btn = gr.Button("🗑️ 모든 문서 삭제", size="sm")
+                        refresh_btn = gr.Button("🔄 목록 새로고침", size="sm")
+                    
+                    enable_rag = gr.Checkbox(
+                        label="RAG 활성화",
+                        value=False,
+                        info="문서 기반 답변 생성 활성화"
+                    )
+                    
+                    with gr.Accordion("⚙️ RAG 고급 설정", open=False):
+                        top_k_chunks = gr.Slider(
+                            minimum=1,
+                            maximum=10,
+                            value=5,
+                            step=1,
+                            label="참조할 청크 수",
+                            info="답변 생성시 참고할 문서 청크의 개수"
+                        )
+                        
+                        chunk_size = gr.Slider(
+                            minimum=500,
+                            maximum=2000,
+                            value=1000,
+                            step=100,
+                            label="청크 크기",
+                            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,
@@ -91,42 +429,129 @@ 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")
                 
-                # Container for model interfaces
+                # 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>"
+                    )
+                
+                # 모델 인터페이스 컨테이너
                 with gr.Column(visible=True) as model_120b_container:
                     gr.Markdown("### Model: openai/gpt-oss-120b")
-                    gr.load("models/openai/gpt-oss-120b", accept_token=login_button, provider="fireworks-ai")
+                    # 실제 모델 로드는 gr.load()로 처리
+                    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("🗑️ 대화 초기화")
                 
                 with gr.Column(visible=False) as model_20b_container:
                     gr.Markdown("### Model: openai/gpt-oss-20b")
-                    gr.load("models/openai/gpt-oss-20b", accept_token=login_button, provider="fireworks-ai")
+                    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("🗑️ 대화 초기화")
+    
+    # 이벤트 핸들러
+    
+    # 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]
+    )
+    
+    # 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>"
+        ),
+        inputs=[enable_rag],
+        outputs=[rag_status]
+    )
     
-    # 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("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)
+    # 채팅 기능 (RAG 통합)
+    def chat_with_rag(message, history, enable_rag, selected_docs, top_k):
+        """RAG를 활용한 채팅"""
+        # RAG 처리
+        processed_message = process_with_rag(message, enable_rag, selected_docs, top_k)
+        
+        # 여기에 실제 모델 API 호출 코드가 들어가야 함
+        # 현재는 예시 응답
+        if enable_rag and selected_docs:
+            response = f"[RAG 활성화] 선택된 {len(selected_docs)}개 문서를 참고하여 답변합니다:\n\n{processed_message[:200]}..."
         else:
-            return gr.update(visible=False), gr.update(visible=True)
+            response = f"[일반 모드] {message}에 대한 답변입니다."
+        
+        history.append((message, response))
+        return "", history
     
-    model_dropdown.change(
-        fn=update_visibility,
-        inputs=[model_dropdown],
-        outputs=[model_120b_container, model_20b_container]
+    # 120b 모델 채팅
+    msg_box_120b.submit(
+        fn=chat_with_rag,
+        inputs=[msg_box_120b, chatbot_120b, enable_rag, document_list, top_k_chunks],
+        outputs=[msg_box_120b, chatbot_120b]
+    )
+    
+    send_btn_120b.click(
+        fn=chat_with_rag,
+        inputs=[msg_box_120b, chatbot_120b, enable_rag, document_list, top_k_chunks],
+        outputs=[msg_box_120b, chatbot_120b]
+    )
+    
+    clear_btn_120b.click(
+        lambda: ([], ""),
+        outputs=[chatbot_120b, msg_box_120b]
+    )
+    
+    # 20b 모델 채팅
+    msg_box_20b.submit(
+        fn=chat_with_rag,
+        inputs=[msg_box_20b, chatbot_20b, enable_rag, document_list, top_k_chunks],
+        outputs=[msg_box_20b, chatbot_20b]
+    )
+    
+    send_btn_20b.click(
+        fn=chat_with_rag,
+        inputs=[msg_box_20b, chatbot_20b, enable_rag, document_list, top_k_chunks],
+        outputs=[msg_box_20b, chatbot_20b]
+    )
+    
+    clear_btn_20b.click(
+        lambda: ([], ""),
+        outputs=[chatbot_20b, msg_box_20b]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()