Update app.py

app.py

@@ -4,14 +4,33 @@ from typing import List, Dict, Any, Optional
 import hashlib
 import json
 from datetime import datetime
+import tempfile
+
+# 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")
+
+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")
 
-# 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
+from typing import Tuple
 
 # Custom CSS (기존 CSS + 추가 스타일)
 custom_css = """
@@ -65,84 +84,96 @@ custom_css = """
     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);
-}
 """
 
-class PDFRAGSystem:
-    """PDF 기반 RAG 시스템 클래스"""
+class SimpleTextSplitter:
+    """간단한 텍스트 분할기"""
+    def __init__(self, chunk_size=1000, chunk_overlap=200):
+        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)
+        
+        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
+        
+        return chunks
+
+class SimplePDFRAGSystem:
+    """간단한 PDF 기반 RAG 시스템"""
     
     def __init__(self):
         self.documents = {}
+        self.document_chunks = {}
+        self.embeddings_store = {}
+        self.text_splitter = SimpleTextSplitter(chunk_size=1000, chunk_overlap=200)
+        
+        # 임베딩 모델 초기화 (가능한 경우)
         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}")
+        if ST_AVAILABLE:
+            try:
+                self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+                print("Embedding model loaded successfully")
+            except Exception as e:
+                print(f"Failed to load embedding model: {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'를 실행해주세요."
+            }
+        
         try:
-            doc = pymupdf.open(pdf_path)
+            doc = fitz.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
-                    })
+                    text_content.append(text)
             
             doc.close()
             
             return {
                 "metadata": metadata,
-                "pages": text_content,
-                "full_text": "\n\n".join([p["content"] for p in text_content])
+                "full_text": "\n\n".join(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 처리 및 벡터 인덱싱"""
+    def process_and_store_pdf(self, pdf_path: str, doc_id: str) -> Dict[str, Any]:
+        """PDF 처리 및 저장"""
         try:
             # PDF 텍스트 추출
             pdf_data = self.extract_text_from_pdf(pdf_path)
@@ -150,33 +181,20 @@ class PDFRAGSystem:
             # 텍스트를 청크로 분할
             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.document_chunks[doc_id] = chunks
             
-            self.collection.add(
-                ids=ids,
-                embeddings=embeddings.tolist(),
-                documents=chunks,
-                metadatas=metadatas
-            )
+            # 임베딩 생성 (가능한 경우)
+            if self.embedder:
+                embeddings = self.embedder.encode(chunks)
+                self.embeddings_store[doc_id] = embeddings
             
             # 문서 정보 저장
             self.documents[doc_id] = {
                 "metadata": pdf_data["metadata"],
                 "chunk_count": len(chunks),
-                "upload_time": datetime.now().isoformat()
+                "upload_time": datetime.now().isoformat(),
+                "full_text": pdf_data["full_text"][:500]  # 처음 500자 저장
             }
             
             return {
@@ -193,59 +211,92 @@ class PDFRAGSystem:
                 "error": str(e)
             }
     
-    def search_relevant_chunks(self, query: str, top_k: int = 5) -> List[Dict]:
+    def search_relevant_chunks(self, query: str, doc_ids: List[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
-            )
+        all_relevant_chunks = []
+        
+        if self.embedder and self.embeddings_store:
+            # 임베딩 기반 검색
+            query_embedding = self.embedder.encode([query])[0]
             
-            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 []
+            for doc_id in doc_ids:
+                if doc_id in self.embeddings_store and doc_id in self.document_chunks:
+                    doc_embeddings = self.embeddings_store[doc_id]
+                    chunks = self.document_chunks[doc_id]
+                    
+                    # 코사인 유사도 계산
+                    similarities = []
+                    for emb in doc_embeddings:
+                        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]
+                        })
+        else:
+            # 임베딩이 없는 경우 키워드 기반 간단한 검색
+            query_lower = query.lower()
+            query_words = set(query_lower.split())
             
-        except Exception as e:
-            print(f"Search error: {e}")
-            return []
+            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):
+                        chunk_lower = chunk.lower()
+                        # 쿼리 단어가 청크에 포함되어 있는지 확인
+                        matching_words = sum(1 for word in query_words if word in chunk_lower)
+                        if matching_words > 0:
+                            all_relevant_chunks.append({
+                                "content": chunk,
+                                "doc_id": doc_id,
+                                "doc_name": self.documents[doc_id]["metadata"]["file_name"],
+                                "chunk_index": idx,
+                                "similarity": matching_words / len(query_words)
+                            })
+        
+        # 유사도 순으로 정렬하고 상위 k개 반환
+        all_relevant_chunks.sort(key=lambda x: x.get('similarity', 0), reverse=True)
+        return all_relevant_chunks[:top_k]
     
-    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
-        ])
+    def generate_context_prompt(self, query: str, chunks: List[Dict]) -> str:
+        """컨텍스트를 포함한 프롬프트 생성"""
+        if not 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 = f"""다음 문서 내용을 참고하여 질문에 답변해주세요. 
-답변은 제공된 문서 내용을 바탕으로 작성하되, 필요시 추가 설명을 포함할 수 있습니다.
-문서에서 관련 정보를 찾을 수 없는 경우, 그 사실을 명시해주세요.
+        context = "\n---\n".join(context_parts)
+        
+        enhanced_prompt = f"""다음 문서 내용을 참고하여 질문에 답변해주세요.
 
-📚 참고 문서:
+## 참고 문서:
 {context}
 
-❓ 질문: {query}
+## 질문:
+{query}
 
-💡 답변:"""
+## 답변:
+위 문서 내용을 바탕으로 질문에 대해 상세하고 정확하게 답변하겠습니다."""
         
-        return prompt
+        return enhanced_prompt
 
 # RAG 시스템 인스턴스 생성
-rag_system = PDFRAGSystem()
+rag_system = SimplePDFRAGSystem()
 
 # State variables
 current_model = gr.State("openai/gpt-oss-120b")
-uploaded_documents = gr.State({})
 rag_enabled = gr.State(False)
 
 def upload_pdf(file):
@@ -260,8 +311,8 @@ def upload_pdf(file):
         
         doc_id = f"doc_{file_hash}"
         
-        # PDF 처리 및 인덱싱
-        result = rag_system.process_and_index_pdf(file.name, doc_id)
+        # PDF 처리 및 저장
+        result = rag_system.process_and_store_pdf(file.name, doc_id)
         
         if result["success"]:
             status_html = f"""
@@ -300,49 +351,67 @@ def upload_pdf(file):
 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 = {}
+        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()
 
-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 비활성화시 원본 메시지 반환
+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
+
+def chat_with_model(message: str, history: List[Tuple[str, str]], enable_rag: bool, selected_docs: List[str], top_k: int, model: str):
+    """모델과 대화 (RAG 포함)"""
     
-    try:
+    # 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, top_k=top_k)
+        relevant_chunks = rag_system.search_relevant_chunks(message, doc_ids, 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
-            ]
+            # 컨텍스트를 포함한 프롬프트 생성
+            enhanced_message = rag_system.generate_context_prompt(message, relevant_chunks)
             
-            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
+            # 디버그 정보 포함 응답 (실제 구현시 모델 API 호출로 대체)
+            response = f"""📚 RAG 기반 답변 (모델: {model})
 
-def switch_model(model_choice):
-    """모델 전환 함수"""
-    return gr.update(visible=False), gr.update(visible=True), model_choice
+찾은 관련 문서 섹션: {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를 활성화하면 문서 기반 답변을 받을 수 있습니다."""
+    
+    history.append((message, response))
+    return "", history
 
 # Gradio 인터페이스
 with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as demo:
@@ -403,14 +472,18 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                             info="답변 생성시 참고할 문서 청크의 개수"
                         )
                         
-                        chunk_size = gr.Slider(
-                            minimum=500,
-                            maximum=2000,
-                            value=1000,
-                            step=100,
-                            label="청크 크기",
-                            info="문서를 분할하는 청크의 크기 (문자 수)"
-                        )
+                        gr.Markdown("""
+                        ### 📝 RAG 사용 팁:
+                        1. PDF 파일을 업로드하세요
+                        2. 업로드된 문서를 선택하세요
+                        3. RAG를 활성화하세요
+                        4. 문서 내용에 대해 질문하세요
+                        
+                        예시 질문:
+                        - "문서의 주요 내용을 요약해주세요"
+                        - "이 문서에서 언급된 날짜는 언제인가요?"
+                        - "참가 자격 조건은 무엇인가요?"
+                        """)
                 
                 # 고급 옵션
                 with gr.Accordion("⚙️ 모델 설정", open=False):
@@ -443,7 +516,6 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
                 # 모델 인터페이스 컨테이너
                 with gr.Column(visible=True) as model_120b_container:
                     gr.Markdown("### Model: openai/gpt-oss-120b")
-                    # 실제 모델 로드는 gr.load()로 처리
                     chatbot_120b = gr.Chatbot(height=400)
                     msg_box_120b = gr.Textbox(
                         label="메시지 입력",
@@ -501,31 +573,15 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
         outputs=[]
     )
     
-    # 채팅 기능 (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:
-            response = f"[일반 모드] {message}에 대한 답변입니다."
-        
-        history.append((message, response))
-        return "", history
-    
     # 120b 모델 채팅
     msg_box_120b.submit(
-        fn=chat_with_rag,
+        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]
     )
     
     send_btn_120b.click(
-        fn=chat_with_rag,
+        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]
     )
@@ -537,13 +593,13 @@ with gr.Blocks(fill_height=True, theme="Nymbo/Nymbo_Theme", css=custom_css) as d
     
     # 20b 모델 채팅
     msg_box_20b.submit(
-        fn=chat_with_rag,
+        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]
     )
     
     send_btn_20b.click(
-        fn=chat_with_rag,
+        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]
     )