Update app.py

app.py

@@ -3,9 +3,8 @@ import sys
 sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 
 import os
-import json
 import logging
-import hashlib
+import traceback
 import gradio as gr
 import pandas as pd
 import docx2txt
@@ -14,332 +13,251 @@ from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers.ensemble import EnsembleRetriever
+from langchain.chains import create_retrieval_chain, create_history_aware_retriever
+from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.documents import Document
-from langchain_core.output_parsers import StrOutputParser
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.chains.combine_documents import create_stuff_documents_chain
 
-# --- CẤU HÌNH ---
-# Thay API Key của bạn vào đây hoặc set biến môi trường
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") 
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 DATA_PATH = "medical_data"
 DB_PATH = "chroma_db"
-TRACKING_FILE = "processed_files.json" # File lưu trạng thái để không học lại dữ liệu cũ
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-
-# --- HÀM XỬ LÝ DỮ LIỆU (INCREMENTAL) ---
+MAX_HISTORY_TURNS = 6
 
-def calculate_md5(file_path):
-    """Tính mã hash MD5 của file"""
-    hash_md5 = hashlib.md5()
-    try:
-        with open(file_path, "rb") as f:
-            for chunk in iter(lambda: f.read(4096), b""):
-                hash_md5.update(chunk)
-        return hash_md5.hexdigest()
-    except Exception:
-        return None
+USE_BM25 = True
+USE_MMR = True
 
-def load_excel_enhanced(file_path, filename):
-    """Đọc Excel giữ ngữ cảnh từng dòng"""
-    try:
-        df = pd.read_excel(file_path)
-        text_parts = []
-        columns = df.columns.tolist()
-        for index, row in df.iterrows():
-            row_str = f"Nguồn: {filename}, Dòng {index+2}: " 
-            items = []
-            for col in columns:
-                val = row[col]
-                if pd.notna(val):
-                    items.append(f"{col}: {val}")
-            if items:
-                text_parts.append(row_str + " | ".join(items))
-        return "\n".join(text_parts)
-    except Exception as e:
-        logging.error(f"Lỗi đọc Excel {filename}: {e}")
-        return ""
-
-def process_single_file(file_path, filename):
-    """Chọn loader phù hợp cho từng loại file"""
-    docs = []
-    filename_lower = filename.lower()
-    try:
-        if filename_lower.endswith(".pdf"):
-            loader = PyPDFLoader(file_path)
-            loaded_docs = loader.load()
-            for d in loaded_docs: 
-                d.metadata["source"] = filename
-            docs.extend(loaded_docs)
-            
-        elif filename_lower.endswith((".xlsx", ".xls")):
-            text = load_excel_enhanced(file_path, filename)
-            if text: 
-                docs.append(Document(page_content=text, metadata={"source": filename}))
-                
-        elif filename_lower.endswith(".docx"):
-            text = docx2txt.process(file_path)
-            if text.strip(): 
-                docs.append(Document(page_content=text, metadata={"source": filename}))
-                
-        elif filename_lower.endswith(".txt"):
-            with open(file_path, "r", encoding="utf-8") as f: 
-                text = f.read()
-            if text.strip(): 
-                docs.append(Document(page_content=text, metadata={"source": filename}))
-                
-    except Exception as e:
-        logging.error(f"Không thể đọc file {filename}: {e}")
-    
-    return docs
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 
-def load_and_process_new_files(folder_path, embedding_model):
-    """
-    Phiên bản sửa lỗi Cache: Tự động học lại nếu phát hiện DB rỗng.
-    """
+def load_documents_from_folder(folder_path: str) -> list[Document]:
+    logging.info(f"--- Bắt đầu quét thư mục: {folder_path} ---")
+    documents: list[Document] = []
     if not os.path.exists(folder_path):
         os.makedirs(folder_path, exist_ok=True)
-        return Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
-
-    # 1. Load file theo dõi cũ (nếu có)
-    processed_files = {}
-    if os.path.exists(TRACKING_FILE):
-        try:
-            with open(TRACKING_FILE, "r") as f:
-                processed_files = json.load(f)
-        except Exception:
-            processed_files = {} # Nếu file lỗi json thì coi như chưa có
-
-    # 2. --- KẾT NỐI & KIỂM TRA THỰC TẾ DATABASE (Đoạn Code Mới) ---
-    vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
-    
-    db_is_empty = False
-    try:
-        # Đếm số lượng dữ liệu thực tế trong Chroma
-        count = vectorstore._collection.count()
-        if count == 0:
-            db_is_empty = True
-    except Exception:
-        db_is_empty = True
-
-    # QUAN TRỌNG: Nếu DB rỗng nhưng file theo dõi lại bảo "đã học" -> LỖI CACHE
-    if db_is_empty and processed_files:
-        logging.warning("⚠️ PHÁT HIỆN LỖI CACHE: Database rỗng nhưng file theo dõi báo đã học.")
-        logging.warning("🔄 Hệ thống sẽ tự động RESET và học lại toàn bộ dữ liệu ngay bây giờ...")
-        processed_files = {} # Xóa sạch ký ức giả để ép hệ thống học lại
-
-    # 3. Quét file và học như bình thường
-    new_docs = []
-    current_files_status = {}
-    files_changed = False
-
-    logging.info("--- Bắt đầu quét dữ liệu ---")
+        return []
     for root, _, files in os.walk(folder_path):
         for filename in files:
-            if filename.startswith("~$"): continue 
-            
             file_path = os.path.join(root, filename)
-            file_hash = calculate_md5(file_path)
-            
-            if not file_hash: continue
-            current_files_status[filename] = file_hash
-
-            # Nếu file chưa có trong danh sách đã xử lý -> Load
-            if filename not in processed_files or processed_files[filename] != file_hash:
-                logging.info(f"⏳ Đang đọc tài liệu: {filename}")
-                file_docs = process_single_file(file_path, filename)
-                new_docs.extend(file_docs)
-                files_changed = True
+            filename_lower = filename.lower()
+            try:
+                if filename_lower.endswith(".pdf"):
+                    loader = PyPDFLoader(file_path)
+                    docs = loader.load()
+                    for d in docs: d.metadata["source"] = filename
+                    documents.extend(docs)
+                elif filename_lower.endswith(".docx"):
+                    text = docx2txt.process(file_path)
+                    if text.strip(): documents.append(Document(page_content=text, metadata={"source": filename}))
+                elif filename_lower.endswith((".xlsx", ".xls", ".csv")):
+                    # Xử lý đơn giản cho excel/csv
+                    if filename_lower.endswith(".csv"): df = pd.read_csv(file_path)
+                    else: df = pd.read_excel(file_path)
+                    text_data = "\n".join([" | ".join(f"{c}:{v}" for c,v in r.items() if pd.notna(v)) for _, r in df.iterrows()])
+                    if text_data.strip(): documents.append(Document(page_content=text_data, metadata={"source": filename}))
+                elif filename_lower.endswith((".txt", ".md")):
+                    with open(file_path, "r", encoding="utf-8") as f: text = f.read()
+                    if text.strip(): documents.append(Document(page_content=text, metadata={"source": filename}))
+            except Exception as e:
+                logging.error(f"Lỗi file {filename}: {e}")
+    return documents
+
+def build_vectorstore_and_corpus(embedding_model):
+    from shutil import rmtree
+    splits: list[Document] = []
+    vectorstore = None
+
+    if os.path.exists(DB_PATH) and os.listdir(DB_PATH):
+        try:
+            vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
+            existing = vectorstore.get()
+            if existing.get("documents"):
+                for text, meta in zip(existing["documents"], existing["metadatas"]):
+                    splits.append(Document(page_content=text, metadata=meta))
+                logging.info(f"Đã load {len(splits)} chunks từ DB cũ.")
+            else:
+                splits = []
+        except Exception:
+            rmtree(DB_PATH, ignore_errors=True)
+            splits = []
 
-    # 4. Nạp dữ liệu mới vào DB
-    if new_docs:
-        logging.info(f"📥 Đang nạp {len(new_docs)} đoạn văn bản vào trí nhớ...")
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(new_docs)
+    if not splits:
+        logging.info("--- Tạo index mới ---")
+        documents = load_documents_from_folder(DATA_PATH)
+        if not documents: return None, []
         
-        # Batch insert để tránh lỗi memory
-        batch_size = 100
-        for i in range(0, len(splits), batch_size):
-            batch = splits[i:i + batch_size]
-            vectorstore.add_documents(batch)
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+        splits = text_splitter.split_documents(documents)
         
-        logging.info("✅ Đã học xong dữ liệu mới.")
-    else:
-        if db_is_empty:
-             logging.warning("⚠️ Thư mục dữ liệu trống! Vui lòng copy file vào folder 'medical_data'.")
-        else:
-             logging.info("✅ Dữ liệu đã cập nhật đầy đủ.")
+        vectorstore = Chroma.from_documents(documents=splits, embedding=embedding_model, persist_directory=DB_PATH)
+        
+    return vectorstore, splits
 
-    # 5. Chỉ lưu file theo dõi khi DB chắc chắn có dữ liệu
-    # (Tránh trường hợp code chạy lỗi giữa chừng mà lại lưu là "đã xong")
-    if files_changed or (not db_is_empty and len(current_files_status) != len(processed_files)):
-        with open(TRACKING_FILE, "w") as f:
-            json.dump(current_files_status, f)
+def get_retriever():
+    logging.info("--- Tải Embedding Model ---")
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    
+    vectorstore, splits = build_vectorstore_and_corpus(embedding_model)
+    if not vectorstore: return None
 
-    return vectorstore
+    if USE_MMR:
+        base_retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 10})
+    else:
+        base_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
+
+    if USE_BM25:
+        bm25_retriever = BM25Retriever.from_documents(splits)
+        bm25_retriever.k = 10
+        ensemble_retriever = EnsembleRetriever(
+            retrievers=[bm25_retriever, base_retriever],
+            weights=[0.4, 0.6]
+        )
+        first_stage_retriever = ensemble_retriever
+    else:
+        first_stage_retriever = base_retriever
 
-# --- BOT CLASS ---
+    logging.info("--- Tải Reranker Model (BGE-M3) ---")
+    reranker_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
+
+    compressor = CrossEncoderReranker(model=reranker_model, top_n=5)
+    
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor, 
+        base_retriever=first_stage_retriever
+    )
+    
+    logging.info("Đã tích hợp Reranker thành công.")
+    return compression_retriever
 
 class DeepMedBot:
     def __init__(self):
         self.retriever = None
-        self.qa_chain = None
-        self.history_chain = None
+        self.rag_chain = None
         self.ready = False
         
         if not GOOGLE_API_KEY:
-            logging.error("LỖI: Chưa cấu hình GOOGLE_API_KEY")
+            logging.error("Thiếu GOOGLE_API_KEY")
             return
-        
+
         try:
-            logging.info("--- Khởi tạo Embedding Model ---")
-            # Model đa ngôn ngữ tốt cho tiếng Việt
-            embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-            
-            # Load DB (Thông minh)
-            self.vectorstore = load_and_process_new_files(DATA_PATH, embedding_model)
-            
-            # Thiết lập Retriever
-            self.retriever = self.vectorstore.as_retriever(search_kwargs={"k": 5})
-            
-            logging.info("--- Khởi tạo LLM ---")
+            self.retriever = get_retriever()
+            if not self.retriever: return
+
             self.llm = ChatGoogleGenerativeAI(
                 model="gemini-2.5-flash",
                 temperature=0.2,
-                google_api_key=GOOGLE_API_KEY,
-                streaming=True # Quan trọng cho hiệu ứng gõ chữ
+                google_api_key=GOOGLE_API_KEY
             )
-            
             self._build_chains()
             self.ready = True
             logging.info("Bot đã sẵn sàng!")
-            
         except Exception as e:
-            logging.error(f"Lỗi khởi tạo Bot: {e}")
+            logging.error(f"Lỗi init bot: {e}")
+            logging.debug(traceback.format_exc())
 
     def _build_chains(self):
-        # 1. Chain viết lại câu hỏi dựa trên lịch sử
-        context_system_prompt = (
-            "Dựa trên lịch sử trò chuyện và câu hỏi mới nhất của người dùng, "
-            "hãy viết lại câu hỏi thành một câu hoàn chỉnh để tìm kiếm thông tin.\n"
-            "KHÔNG trả lời câu hỏi, chỉ viết lại câu hỏi thôi."
+        contextualize_q_system_prompt = (
+            "Dựa trên lịch sử trò chuyện và câu hỏi mới nhất, "
+            "hãy viết lại câu hỏi thành một câu hoàn chỉnh nếu cần. "
+            "Nếu không liên quan, giữ nguyên. KHÔNG trả lời, chỉ viết lại."
         )
-        context_prompt = ChatPromptTemplate.from_messages([
-            ("system", context_system_prompt),
+        contextualize_q_prompt = ChatPromptTemplate.from_messages([
+            ("system", contextualize_q_system_prompt),
             MessagesPlaceholder("chat_history"),
             ("human", "{input}"),
         ])
-        self.history_chain = context_prompt | self.llm | StrOutputParser()
+        history_aware_retriever = create_history_aware_retriever(
+            self.llm, self.retriever, contextualize_q_prompt
+        )
 
-        # 2. Chain trả lời câu hỏi (RAG)
         qa_system_prompt = (
-            "Bạn là trợ lý y tế DeepMed tại TTYT Thanh Ba. Dựa vào Context sau để trả lời.\n"
+            "Bạn là trợ lý y tế DeepMed. Dựa vào Context sau để trả lời.\n"
             "Context:\n{context}\n\n"
-            "Quy tắc:\n"
-            "- Trả lời ngắn gọn, đúng trọng tâm.\n"
-            "- Nếu là thông tin thuốc (liều, chống chỉ định), phải giữ nguyên số liệu.\n"
-            "- Nếu không có thông tin trong Context, hãy nói 'Tôi chưa có dữ liệu về vấn đề này'."
+            "Yêu cầu:\n"
+            "- Chỉ trả lời dựa trên Context. Nếu không có thông tin, nói 'Tôi không tìm thấy thông tin trong tài liệu'.\n"
+            "- Trả lời ngắn gọn, có cấu trúc.\n"
         )
         qa_prompt = ChatPromptTemplate.from_messages([
             ("system", qa_system_prompt),
-            ("human", "{question}"),
+            MessagesPlaceholder("chat_history"),
+            ("human", "{input}"),
         ])
         
-        # create_stuff_documents_chain nhận vào 'context' (list docs) và 'question'
-        self.qa_chain = create_stuff_documents_chain(self.llm, qa_prompt)
+        question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt)
+        self.rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
 
     def chat_stream(self, message: str, history: list):
+        """
+        Hàm Generator dùng cho Streaming
+        """
         if not self.ready:
-            yield "Hệ thống đang khởi động hoặc lỗi cấu hình..."
+            yield "Hệ thống đang khởi động hoặc lỗi dữ liệu."
             return
 
-        # Bước 1: Contextualize - Hiểu ngữ cảnh nếu có lịch sử
-        actual_question = message
-        if history:
-            yield "🔄 Đang phân tích ngữ cảnh..."
-            chat_history_objs = []
-            # Lấy 4 turn gần nhất để làm context
-            for u, a in history[-4:]:
-                chat_history_objs.append(HumanMessage(content=u))
-                chat_history_objs.append(AIMessage(content=a))
-            
-            try:
-                actual_question = self.history_chain.invoke({
-                    "chat_history": chat_history_objs,
-                    "input": message
-                })
-            except Exception:
-                actual_question = message # Fallback nếu lỗi
+        chat_history = []
+        for u, b in history[-MAX_HISTORY_TURNS:]:
+            chat_history.append(HumanMessage(content=u))
+            chat_history.append(AIMessage(content=b))
 
-        # Bước 2: Retrieval - Tìm kiếm
-        yield "📚 Đang tra cứu tài liệu y khoa..."
-        docs = self.retriever.invoke(actual_question)
-        
-        if not docs:
-            yield "Không tìm thấy thông tin phù hợp trong cơ sở dữ liệu nội bộ."
-            return
-
-        # Bước 3: Generation - Trả lời Streaming
         full_response = ""
-        input_data = {"context": docs, "question": actual_question}
+        retrieved_docs = []
         
         try:
-            # Loop qua từng token được sinh ra
-            for chunk in self.qa_chain.stream(input_data):
-                # Xử lý format output (đôi khi là string, đôi khi là object)
-                if isinstance(chunk, str):
-                    content = chunk
-                elif hasattr(chunk, "content"):
-                    content = chunk.content
-                else:
-                    content = str(chunk)
-                
-                full_response += content
-                yield full_response 
+            for chunk in self.rag_chain.stream({"input": message, "chat_history": chat_history}):
+                if "answer" in chunk:
+                    content = chunk["answer"]
+                    full_response += content
+                    yield full_response 
+
+                if "context" in chunk:
+                    retrieved_docs = chunk["context"]
+
+            if retrieved_docs:
+                refs = self._build_references_text(retrieved_docs)
+                if refs:
+                    full_response += f"\n\n---\n📚 **Nguồn:**\n{refs}"
+                    yield full_response
 
-            # Bước 4: Thêm nguồn tham khảo
-            refs = self._build_references_text(docs)
-            if refs:
-                full_response += f"\n\n---\n📚 **Nguồn tham khảo:**\n{refs}"
-                yield full_response
-                
         except Exception as e:
-            logging.error(f"Lỗi Stream: {e}")
-            yield "Đã xảy ra lỗi trong quá trình tạo câu trả lời."
+            logging.error(f"Lỗi stream: {e}")
+            yield "Gặp lỗi trong quá trình xử lý."
 
     @staticmethod
-    def _build_references_text(docs):
-        sources = set()
+    def _build_references_text(docs) -> str:
+        from collections import defaultdict
+        source_pages = defaultdict(set)
         for doc in docs:
             src = os.path.basename(doc.metadata.get("source", "Tài liệu"))
-            # Nếu có số trang (PDF)
             if "page" in doc.metadata:
-                src += f" (Trang {doc.metadata['page'] + 1})"
-            sources.add(src)
+                source_pages[src].add(doc.metadata["page"] + 1)
+            else:
+                source_pages[src]
         
-        return "\n".join([f"- {s}" for s in sources])
-
-# --- GIAO DIỆN GRADIO ---
+        lines = []
+        for src, pages in source_pages.items():
+            if pages:
+                p_str = ", ".join(str(p) for p in sorted(pages))
+                lines.append(f"- {src} (Trang {p_str})")
+            else:
+                lines.append(f"- {src}")
+        return "\n".join(lines)
 
 bot = DeepMedBot()
 
-def gradio_chat_generator(message, history):
-    # Hàm này đóng vai trò cầu nối, gọi generator của bot
+def gradio_chat_stream(message, history):
     for partial_response in bot.chat_stream(message, history):
         yield partial_response
 
-custom_css = """
-#component-0 {height: 90vh !important;} 
-footer {visibility: hidden}
-"""
-
 demo = gr.ChatInterface(
-    fn=gradio_chat_generator,
-    title="🏥 DeepMed AI - Trợ lý Y tế Trung Tâm Y Tế Khu Vực Thanh Ba",
+    fn=gradio_chat_stream,
+    title="🏥 DeepMed AI - Trợ lý Y tế tại Trung Tâm Y Tế Khu Vực Thanh Ba",
 )
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()