Update app.py

app.py

@@ -1,308 +1,135 @@
 import os
 import sys
 import logging
-import traceback
-import pandas as pd
-import docx2txt
 import chromadb
-from pathlib import Path
+import gradio as gr
 
-# --- HACK FIX CHO CHROMA DB (Nếu chạy trên môi trường cũ) ---
+# Fix lỗi SQLite cũ trên Hugging Face
 try:
     __import__("pysqlite3")
     sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 except ImportError:
     pass
 
-import gradio as gr
-from chromadb.config import Settings
 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_huggingface import HuggingFaceEmbeddings
 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 import create_retrieval_chain
 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_huggingface import HuggingFaceEmbeddings
+from langchain_core.prompts import ChatPromptTemplate
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain_core.documents import Document
 
-# --- CẤU HÌNH ---
-# Hãy chắc chắn bạn đã set biến môi trường GOOGLE_API_KEY hoặc điền trực tiếp vào đây (không khuyến khích share key)
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") 
-DATA_PATH = "medical_data"
-DB_PATH = "chroma_db"
-MAX_HISTORY_TURNS = 4
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+DB_PATH = "chroma_db" 
 
-# --- 1. PHÂN LOẠI DỮ LIỆU (THEO ẢNH CỦA BẠN) ---
-def classify_by_folder_path(file_path_str: str) -> str:
-    path_lower = file_path_str.lower().replace("\\", "/")
-    
-    if "thông tin thuốc nội bộ" in path_lower:
-        return "drug_info"
-    elif "phác đồ tại ttytkv thanh ba" in path_lower:
-        return "local_regimen"
-    elif "phác đồ bộ y tế" in path_lower:
-        return "moh_regimen"
-    elif "các hiệp hội" in path_lower:
-        return "association"
-    elif "general_knowledge" in path_lower:
-        return "general_knowledge"
-    else:
-        return "general_knowledge"
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
 
 def get_category_vn_name(cat_code):
-    mapping = {
+    return {
         "drug_info": "💊 Thuốc Nội Bộ",
         "local_regimen": "🏥 Phác Đồ Thanh Ba",
         "moh_regimen": "🏛️ Bộ Y Tế",
-        "association": "🌐 Hiệp Hội",
-        "general_knowledge": "📚 Kiến thức chung"
-    }
-    return mapping.get(cat_code, "Khác")
-
-# --- XỬ LÝ FILE ---
-def process_excel_file(file_path: str, filename: str, category: str) -> list[Document]:
-    docs = []
-    try:
-        df = pd.read_csv(file_path) if file_path.endswith(".csv") else pd.read_excel(file_path)
-        df.dropna(how='all', inplace=True)
-        df.fillna("Không có thông tin", inplace=True)
-        cat_vn = get_category_vn_name(category)
-        
-        for idx, row in df.iterrows():
-            content_parts = []
-            for col_name, val in row.items():
-                clean_val = str(val).strip()
-                if clean_val and clean_val.lower() != "nan":
-                    content_parts.append(f"{col_name}: {clean_val}")
-            
-            if content_parts:
-                page_content = f"[{cat_vn}] Nguồn: {filename}\n" + "\n".join(content_parts)
-                metadata = {"source": filename, "row": idx+1, "type": "excel", "category": category}
-                docs.append(Document(page_content=page_content, metadata=metadata))
-    except Exception as e:
-        logging.error(f"Lỗi Excel {filename}: {e}")
-    return docs
-
-def load_documents_from_folder(root_folder: str) -> list[Document]:
-    logging.info(f"--- Quét dữ liệu từ: {root_folder} ---")
-    documents = []
-    if not os.path.exists(root_folder):
-        os.makedirs(root_folder, exist_ok=True)
-        return []
-        
-    for root, _, files in os.walk(root_folder):
-        for filename in files:
-            file_path = os.path.join(root, filename)
-            
-            # --- Phân loại ---
-            category = classify_by_folder_path(file_path)
-            cat_vn = get_category_vn_name(category)
-            logging.info(f"Load: {filename} -> Nhóm: {cat_vn}")
-            
-            try:
-                new_docs = []
-                if filename.lower().endswith(".pdf"):
-                    loader = PyPDFLoader(file_path)
-                    new_docs = loader.load()
-                elif filename.lower().endswith(".docx"):
-                    text = docx2txt.process(file_path)
-                    if text.strip():
-                        new_docs = [Document(page_content=text)]
-                elif filename.lower().endswith((".xlsx", ".xls", ".csv")):
-                    new_docs = process_excel_file(file_path, filename, category)
-                elif filename.lower().endswith((".txt", ".md")):
-                    with open(file_path, "r", encoding="utf-8") as f:
-                        new_docs = [Document(page_content=f.read())]
-                
-                # Gắn metadata và Label vào text
-                for doc in new_docs:
-                    doc.metadata["source"] = filename
-                    doc.metadata["category"] = category 
-                    if not doc.page_content.startswith("["): 
-                        doc.page_content = f"[{cat_vn}] Nguồn {filename}:\n{doc.page_content}"
-                    
-                documents.extend(new_docs)
-            except Exception as e:
-                logging.error(f"Lỗi đọc {filename}: {e}")
-    return documents
+        "association": "🌐 Hiệp Hội"
+    }.get(cat_code, "Khác")
 
-# --- RETRIEVER CONFIG (Bộ lọc thông minh) ---
 def get_retrievers():
-    logging.info("--- Tải Embedding Model ---")
-    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-    chroma_settings = Settings(anonymized_telemetry=False)
-    
-    # Logic nạp lại dữ liệu (Re-index nếu chưa có DB)
-    if not os.path.exists(DB_PATH) or not os.listdir(DB_PATH):
-        raw_docs = load_documents_from_folder(DATA_PATH)
-        if not raw_docs: return None, None
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(raw_docs)
-        vectorstore = Chroma.from_documents(documents=splits, embedding=embedding_model, persist_directory=DB_PATH, client_settings=chroma_settings)
-    else:
-        logging.info("--- Sử dụng DB đã lưu ---")
-        vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding_model, client_settings=chroma_settings)
-        # Load lại raw để build BM25 (In-memory)
-        raw_docs = load_documents_from_folder(DATA_PATH)
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(raw_docs)
-
-    # 1. TRA CỨU NHANH (Chỉ lấy 'drug_info')
-    vector_retriever_fast = vectorstore.as_retriever(
-        search_kwargs={"k": 5, "filter": {"category": "drug_info"}}
-    )
-    # BM25 cho thuốc
-    drug_splits = [d for d in splits if d.metadata.get("category") == "drug_info"]
-    if drug_splits:
-        bm25_fast = BM25Retriever.from_documents(drug_splits)
-        bm25_fast.k = 5
-        fast_retriever = EnsembleRetriever(retrievers=[bm25_fast, vector_retriever_fast], weights=[0.4, 0.6])
-    else:
-        fast_retriever = vector_retriever_fast
+    if not os.path.exists(DB_PATH):
+        raise FileNotFoundError("❌ Chưa upload folder 'chroma_db'!")
 
-    # 2. TRA CỨU CHUYÊN SÂU (Bỏ general_knowledge, ưu tiên các nhóm chuyên môn)
-    target_categories = ["local_regimen", "moh_regimen", "association", "drug_info"]
+    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding)
     
-    vector_retriever_deep = vectorstore.as_retriever(
-        search_kwargs={"k": 25, "filter": {"category": {"$in": target_categories}}}
-    )
-    deep_splits = [d for d in splits if d.metadata.get("category") in target_categories]
-    if deep_splits:
-        bm25_deep = BM25Retriever.from_documents(deep_splits)
-        bm25_deep.k = 25
-        ensemble_deep = EnsembleRetriever(retrievers=[bm25_deep, vector_retriever_deep], weights=[0.5, 0.5])
-    else:
-        ensemble_deep = vector_retriever_deep
-
-    # Reranker BGE-M3 (Bước lọc cuối cùng cực quan trọng)
+    # Tái tạo BM25 từ VectorStore (Trick để không phải upload raw data)
+    all_data = vectorstore.get()
+    splits = [Document(page_content=txt, metadata=m) for txt, m in zip(all_data['documents'], all_data['metadatas'])]
+
+    # 1. FAST (Chỉ thuốc)
+    vec_fast = vectorstore.as_retriever(search_kwargs={"k": 5, "filter": {"category": "drug_info"}})
+    drug_docs = [d for d in splits if d.metadata.get("category") == "drug_info"]
+    bm25_fast = BM25Retriever.from_documents(drug_docs) if drug_docs else None
+    bm25_fast.k = 5 if bm25_fast else 5
+    
+    fast_retriever = EnsembleRetriever(retrievers=[bm25_fast, vec_fast], weights=[0.4, 0.6]) if bm25_fast else vec_fast
+
+    # 2. DEEP (Ưu tiên Thanh Ba)
+    cats = ["local_regimen", "moh_regimen", "association", "drug_info"]
+    vec_deep = vectorstore.as_retriever(search_kwargs={"k": 25, "filter": {"category": {"$in": cats}}})
+    deep_docs = [d for d in splits if d.metadata.get("category") in cats]
+    bm25_deep = BM25Retriever.from_documents(deep_docs) if deep_docs else None
+    bm25_deep.k = 25 if bm25_deep else 25
+    
+    ensemble = EnsembleRetriever(retrievers=[bm25_deep, vec_deep], weights=[0.5, 0.5]) if bm25_deep else vec_deep
+    
+    # Rerank
     reranker = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
-    compressor = CrossEncoderReranker(model=reranker, top_n=10) # Lấy top 10 đoạn văn chuẩn nhất
-    deep_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=ensemble_deep)
+    compressor = CrossEncoderReranker(model=reranker, top_n=10)
+    deep_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=ensemble)
     
     return fast_retriever, deep_retriever
 
-# --- BOT SETUP ---
 class DeepMedBot:
     def __init__(self):
         self.ready = False
-        if not GOOGLE_API_KEY:
-            logging.error("❌ Thiếu API Key! Vui lòng kiểm tra biến môi trường.")
-            return
         try:
             self.fast_retriever, self.deep_retriever = get_retrievers()
-            self.llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2, google_api_key=GOOGLE_API_KEY)
+            self.llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.2, google_api_key=GOOGLE_API_KEY)
             self._build_chains()
             self.ready = True
         except Exception as e:
-            logging.error(f"Lỗi khởi động: {e}")
+            logging.error(f"Lỗi: {e}")
 
     def _build_chains(self):
-        # --- 1. PROMPT TRA CỨU NHANH (Tối ưu hiển thị Bảng thuốc) ---
-        fast_system = (
-            "Bạn là Dược sĩ Lâm sàng phụ trách kho dược nội bộ.\n"
-            "Nhiệm vụ: Tra cứu thông tin thuốc từ dữ liệu [💊 Thuốc Nội Bộ] và trả lời người dùng.\n\n"
-            "YÊU CẦU ĐỊNH DẠNG (BẮT BUỘC):\n"
-            "1. **Luôn trình bày kết quả dưới dạng Bảng (Markdown Table)** với các cột sau (nếu có thông tin):\n"
-            "   | Tên thuốc | Hoạt chất | Hàm lượng | Đơn vị tính | Ghi chú/Tồn kho |\n"
-            "   | --- | --- | --- | --- | --- |\n"
-            "2. Nếu tìm thấy nhiều thuốc tương tự, hãy liệt kê hết vào bảng.\n"
-            "3. Sau bảng, có thể bổ sung ngắn gọn về chỉ định hoặc liều dùng nếu dữ liệu có đề cập.\n"
-            "4. Nếu không tìm thấy thuốc nào khớp, trả lời: '❌ Không tìm thấy thuốc này trong danh mục nội bộ'.\n\n"
+        # Prompt Bảng cho Thuốc
+        fast_sys = (
+            "Bạn là Dược sĩ. Tra cứu [💊 Thuốc Nội Bộ] và trả lời bằng **Bảng Markdown**:\n"
+            "| Tên thuốc | Hoạt chất | Hàm lượng | ĐVT | Ghi chú |\n"
+            "| --- | --- | --- | --- | --- |\n"
+            "Nếu không thấy, báo: '❌ Không có trong kho'."
             "Context:\n{context}"
         )
-        fast_prompt = ChatPromptTemplate.from_messages([("system", fast_system), ("human", "{input}")])
-        self.fast_chain = create_retrieval_chain(self.fast_retriever, create_stuff_documents_chain(self.llm, fast_prompt))
+        fast_chain = create_stuff_documents_chain(self.llm, ChatPromptTemplate.from_messages([("system", fast_sys), ("human", "{input}")]))
+        self.fast_chain = create_retrieval_chain(self.fast_retriever, fast_chain)
         
-        # --- 2. PROMPT CHUYÊN SÂU (Logic Thanh Ba + Bảng đối chiếu) ---
-        deep_system = (
-            "Bạn là Bác sĩ Trưởng khoa, hỗ trợ ra quyết định lâm sàng dựa trên bằng chứng.\n"
-            "Nhiệm vụ: Đề xuất phác đồ điều trị và đối chiếu thuốc có sẵn.\n\n"
-            "QUY TRÌNH TƯ DUY & TRẢ LỜI:\n"
-            "Bước 1: Xác định Phác đồ (Theo thứ tự ưu tiên):\n"
-            "   - Ưu tiên 1: [🏥 Phác Đồ Thanh Ba]. (Nếu có, phải tuân thủ tuyệt đối).\n"
-            "   - Ưu tiên 2: [🏛️ Bộ Y Tế] hoặc [🌐 Hiệp Hội] (Chỉ dùng khi Thanh Ba không quy định).\n\n"
-            "Bước 2: Đối chiếu Kho Dược (Quan trọng):\n"
-            "   - Kiểm tra xem các thuốc trong phác đồ có trong [💊 Thuốc Nội Bộ] hay không.\n\n"
-            "YÊU CẦU ĐỊNH DẠNG ĐẦU RA:\n"
-            "1. **Tóm tắt Chẩn đoán/Nguyên tắc:** (Ngắn gọn, gạch đầu dòng).\n"
-            "2. **Phác đồ Điều trị:** (Nêu rõ nguồn áp dụng là Thanh Ba hay BYT).\n"
-            "3. **Bảng Kê Đơn & Đối Chiếu Thuốc (BẮT BUỘC CÓ BẢNG):**\n"
-            "   | Tên thuốc (Theo phác đồ) | Liều dùng | Trạng thái Kho Dược | Gợi ý thay thế (Nếu thiếu) |\n"
-            "   | --- | --- | --- | --- |\n"
-            "   | Ví dụ: Paracetamol | 500mg | ✅ Có sẵn | - |\n"
-            "   | Ví dụ: Thuốc lạ X | ... | ❌ Không có | Dùng thuốc Y trong kho |\n\n"
-            "Lưu ý: Chỉ đưa ra thông tin có trong Context, không tự bịa đặt thuốc."
+        # Prompt Phác đồ ưu tiên Thanh Ba
+        deep_sys = (
+            "Bạn là Bác sĩ Trưởng khoa.\n"
+            "1. **Tìm phác đồ:** Ưu tiên tuyệt đối [🏥 Phác Đồ Thanh Ba]. Nếu không có mới dùng [Bộ Y Tế].\n"
+            "2. **Đối chiếu thuốc:** Kiểm tra thuốc trong phác đồ có trong [💊 Thuốc Nội Bộ] không.\n"
+            "3. **Định dạng trả lời:**\n"
+            "   - Chẩn đoán/Nguyên tắc.\n"
+            "   - Phác đồ (Ghi rõ nguồn).\n"
+            "   - **Bảng kê đơn:**\n"
+            "     | Tên thuốc | Liều dùng | Có trong kho? | Thay thế |\n"
+            "     | --- | --- | --- | --- |\n"
             "Context:\n{context}"
         )
-        deep_prompt = ChatPromptTemplate.from_messages([("system", deep_system), ("human", "{input}")])
-        self.deep_chain = create_retrieval_chain(self.deep_retriever, create_stuff_documents_chain(self.llm, deep_prompt))
+        deep_chain = create_stuff_documents_chain(self.llm, ChatPromptTemplate.from_messages([("system", deep_sys), ("human", "{input}")]))
+        self.deep_chain = create_retrieval_chain(self.deep_retriever, deep_chain)
 
     def chat(self, msg, history, mode):
-        if not self.ready: return "⚠️ Hệ thống đang khởi động hoặc chưa có dữ liệu..."
-        
+        if not self.ready: return "⚠️ Đang khởi động... Vui lòng đợi 1 phút."
         chain = self.deep_chain if mode == "Chuyên sâu" else self.fast_chain
-        response = chain.invoke({"input": msg})
+        res = chain.invoke({"input": msg})
         
-        answer = response['answer']
-        # Tạo trích dẫn nguồn
-        if 'context' in response and response['context']:
-            refs = []
-            seen = set()
-            for doc in response['context']:
-                cat_vn = get_category_vn_name(doc.metadata.get('category'))
-                src = doc.metadata.get('source')
-                ref_item = f"[{cat_vn}] {src}"
-                if ref_item not in seen:
-                    refs.append(f"- {ref_item}")
-                    seen.add(ref_item)
-            if refs:
-                answer += "\n\n---\n📚 **Nguồn tham khảo:**\n" + "\n".join(refs)
-        return answer
+        ans = res['answer']
+        if 'context' in res and res['context']:
+            refs = list(set([f"- [{get_category_vn_name(d.metadata.get('category'))}] {d.metadata.get('source')}" for d in res['context']]))
+            ans += "\n\n---\n📚 **Nguồn:**\n" + "\n".join(refs)
+        return ans
 
-# --- GRADIO UI ---
 bot = DeepMedBot()
 
 def respond(message, history, mode):
     return bot.chat(message, history, mode)
 
-css = """.gradio-container {min_height: 600px}"""
-
-demo = gr.ChatInterface(
+gr.ChatInterface(
     fn=respond,
-    additional_inputs=[
-        gr.Radio(
-            ["Tra cứu nhanh (Chỉ thuốc)", "Chuyên sâu"], 
-            value="Tra cứu nhanh (Chỉ thuốc)", 
-            label="Chế độ tra cứu"
-        )
-    ],
-    title="Hệ thống Hỗ trợ Lâm sàng (TTYT Thanh Ba)",
-    description="Tra cứu thuốc nội bộ và Phác đồ điều trị (Ưu tiên dữ liệu Thanh Ba).",
-    css=css
-)
-
-if __name__ == "__main__":
-    # Tạo folder mẫu để tránh lỗi nếu thư mục chưa tồn tại (Dựa trên tên trong ảnh)
-    folders = [
-        "thông tin thuốc nội bộ",
-        "phác đồ tại ttytkv thanh ba",
-        "phác đồ bộ y tế",
-        "các hiệp hội",
-        "general_knowledge"
-    ]
-    for f in folders:
-        path = os.path.join(DATA_PATH, f)
-        if not os.path.exists(path):
-            os.makedirs(path)
-            
-    # Lưu ý: Xóa folder 'chroma_db' cũ trước khi chạy lại lần đầu tiên sau khi sửa code!
-    demo.launch()
+    additional_inputs=[gr.Radio(["Tra cứu nhanh (Chỉ thuốc)", "Chuyên sâu"], value="Tra cứu nhanh (Chỉ thuốc)", label="Chế độ")],
+    title="TTYT Thanh Ba - Hỗ trợ Lâm sàng",
+    css=".gradio-container {min_height: 600px}"
+).launch()