Update app.py

app.py

@@ -3,8 +3,9 @@ import sys
 sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 
 import os
+import json
 import logging
-import traceback
+import hashlib
 import gradio as gr
 import pandas as pd
 import docx2txt
@@ -13,251 +14,307 @@ 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
 
-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")
+# --- 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") 
 DATA_PATH = "medical_data"
 DB_PATH = "chroma_db"
-MAX_HISTORY_TURNS = 6
-
-USE_BM25 = True
-USE_MMR = True
+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")
 
-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] = []
+# --- HÀM XỬ LÝ DỮ LIỆU (INCREMENTAL) ---
+
+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
+
+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
+
+def load_and_process_new_files(folder_path, embedding_model):
+    """Chỉ xử lý file mới hoặc file bị thay đổi"""
     if not os.path.exists(folder_path):
         os.makedirs(folder_path, exist_ok=True)
-        return []
+        return Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
+
+    # Load file tracking cũ
+    if os.path.exists(TRACKING_FILE):
+        with open(TRACKING_FILE, "r") as f:
+            processed_files = json.load(f)
+    else:
+        processed_files = {}
+
+    new_docs = []
+    current_files_status = {}
+    files_changed = False
+
+    logging.info("--- Đang quét thay đổi trong thư mục dữ liệu ---")
+
     for root, _, files in os.walk(folder_path):
         for filename in files:
+            if filename.startswith("~$"): continue # Bỏ qua file tạm của Word/Excel
+            
             file_path = os.path.join(root, filename)
-            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 = []
-
-    if not splits:
-        logging.info("--- Tạo index mới ---")
-        documents = load_documents_from_folder(DATA_PATH)
-        if not documents: return None, []
-        
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
-        splits = text_splitter.split_documents(documents)
-        
-        vectorstore = Chroma.from_documents(documents=splits, embedding=embedding_model, persist_directory=DB_PATH)
-        
-    return vectorstore, splits
+            file_hash = calculate_md5(file_path)
+            
+            if not file_hash: continue
+            
+            current_files_status[filename] = file_hash
 
-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
+            # Kiểm tra: Nếu file chưa có trong DB hoặc Hash đã đổi -> Load lại
+            if filename not in processed_files or processed_files[filename] != file_hash:
+                logging.info(f"Phát hiện mới/cập nhật: {filename}")
+                file_docs = process_single_file(file_path, filename)
+                new_docs.extend(file_docs)
+                files_changed = True
 
-    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
+    vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
+
+    if new_docs:
+        logging.info(f"Đang Embed {len(new_docs)} chunks mới...")
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.split_documents(new_docs)
+        vectorstore.add_documents(splits)
+        logging.info("Hoàn tất cập nhật DB.")
     else:
-        first_stage_retriever = base_retriever
+        logging.info("Không có dữ liệu mới cần cập nhật.")
 
-    logging.info("--- Tải Reranker Model (BGE-M3) ---")
-    reranker_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
+    # Lưu lại trạng thái mới (chỉ lưu những file hiện đang tồn tại)
+    if files_changed or len(current_files_status) != len(processed_files):
+        with open(TRACKING_FILE, "w") as f:
+            json.dump(current_files_status, f)
 
-    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
+    return vectorstore
+
+# --- BOT CLASS ---
 
 class DeepMedBot:
     def __init__(self):
         self.retriever = None
-        self.rag_chain = None
+        self.qa_chain = None
+        self.history_chain = None
         self.ready = False
         
         if not GOOGLE_API_KEY:
-            logging.error("Thiếu GOOGLE_API_KEY")
+            logging.error("LỖI: Chưa cấu hình GOOGLE_API_KEY")
             return
-
+        
         try:
-            self.retriever = get_retriever()
-            if not self.retriever: return
-
+            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.llm = ChatGoogleGenerativeAI(
                 model="gemini-2.5-flash",
                 temperature=0.2,
-                google_api_key=GOOGLE_API_KEY
+                google_api_key=GOOGLE_API_KEY,
+                streaming=True # Quan trọng cho hiệu ứng gõ chữ
             )
+            
             self._build_chains()
             self.ready = True
             logging.info("Bot đã sẵn sàng!")
+            
         except Exception as e:
-            logging.error(f"Lỗi init bot: {e}")
-            logging.debug(traceback.format_exc())
+            logging.error(f"Lỗi khởi tạo Bot: {e}")
 
     def _build_chains(self):
-        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."
+        # 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_prompt = ChatPromptTemplate.from_messages([
-            ("system", contextualize_q_system_prompt),
+        context_prompt = ChatPromptTemplate.from_messages([
+            ("system", context_system_prompt),
             MessagesPlaceholder("chat_history"),
             ("human", "{input}"),
         ])
-        history_aware_retriever = create_history_aware_retriever(
-            self.llm, self.retriever, contextualize_q_prompt
-        )
+        self.history_chain = context_prompt | self.llm | StrOutputParser()
 
+        # 2. Chain trả lời câu hỏi (RAG)
         qa_system_prompt = (
-            "Bạn là trợ lý y tế DeepMed. Dựa vào Context sau để trả lời.\n"
+            "Bạn là trợ lý y tế DeepMed tại TTYT Thanh Ba. Dựa vào Context sau để trả lời.\n"
             "Context:\n{context}\n\n"
-            "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"
+            "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'."
         )
         qa_prompt = ChatPromptTemplate.from_messages([
             ("system", qa_system_prompt),
-            MessagesPlaceholder("chat_history"),
-            ("human", "{input}"),
+            ("human", "{question}"),
         ])
         
-        question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt)
-        self.rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
+        # create_stuff_documents_chain nhận vào 'context' (list docs) và 'question'
+        self.qa_chain = create_stuff_documents_chain(self.llm, qa_prompt)
 
     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 dữ liệu."
+            yield "Hệ thống đang khởi động hoặc lỗi cấu hình..."
             return
 
-        chat_history = []
-        for u, b in history[-MAX_HISTORY_TURNS:]:
-            chat_history.append(HumanMessage(content=u))
-            chat_history.append(AIMessage(content=b))
+        # 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
+
+        # 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 = ""
-        retrieved_docs = []
+        input_data = {"context": docs, "question": actual_question}
         
         try:
-            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
+            # 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 
 
+            # 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 "Gặp lỗi trong quá trình xử lý."
+            logging.error(f"Lỗi Stream: {e}")
+            yield "Đã xảy ra lỗi trong quá trình tạo câu trả lời."
 
     @staticmethod
-    def _build_references_text(docs) -> str:
-        from collections import defaultdict
-        source_pages = defaultdict(set)
+    def _build_references_text(docs):
+        sources = 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:
-                source_pages[src].add(doc.metadata["page"] + 1)
-            else:
-                source_pages[src]
+                src += f" (Trang {doc.metadata['page'] + 1})"
+            sources.add(src)
         
-        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)
+        return "\n".join([f"- {s}" for s in sources])
+
+# --- GIAO DIỆN GRADIO ---
 
 bot = DeepMedBot()
 
-def gradio_chat_stream(message, history):
+def gradio_chat_generator(message, history):
+    # Hàm này đóng vai trò cầu nối, gọi generator của bot
     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_stream,
-    title="🏥 DeepMed AI - Trợ lý Y tế tại Trung Tâm Y Tế Khu Vực Thanh Ba",
+    fn=gradio_chat_generator,
+    title="🏥 DeepMed AI - Trợ lý Y tế TTYT Thanh Ba",
+    description="Hỗ trợ tra cứu Phác đồ điều trị, Dược thư và Quy trình kỹ thuật.",
+    theme="soft",
+    css=custom_css,
+    examples=["Liều dùng Paracetamol cho trẻ em?", "Phác đồ điều trị sốt xuất huyết?", "Quy trình rửa tay ngoại khoa?"],
+    cache_examples=False
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=False)