Update app.py

app.py

@@ -1,13 +1,16 @@
 import os
 import sys
+import logging
+import gradio as gr
+
+# --- 1. SỬA LỖI SQLITE TRÊN HUGGING FACE (BẮT BUỘC ĐỂ ĐẦU FILE) ---
 try:
     __import__("pysqlite3")
     sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 except ImportError:
-    pass
-import logging
+    pass # Nếu chạy local không có pysqlite3 thì bỏ qua
+
 import chromadb
-import gradio as gr
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -21,8 +24,9 @@ from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain_core.documents import Document
 
+# --- CẤU HÌNH ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
-DB_PATH = "chroma_db" 
+DB_PATH = "chroma_db"
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
 
@@ -34,61 +38,71 @@ def get_category_vn_name(cat_code):
         "association": "🌐 Hiệp Hội"
     }.get(cat_code, "Khác")
 
+# --- 2. LOAD DB ĐÃ CÓ (KHÔNG BUILD LẠI) ---
 def get_retrievers():
     if not os.path.exists(DB_PATH):
-        raise FileNotFoundError("❌ Chưa upload folder 'chroma_db'!")
+        raise FileNotFoundError("❌ LỖI: Chưa upload folder 'chroma_db' lên Hugging Face!")
 
+    logging.info("--- Đang tải dữ liệu... ---")
     embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
     vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding)
     
+    # Tái tạo BM25 từ VectorStore
     all_data = vectorstore.get()
     splits = [Document(page_content=txt, metadata=m) for txt, m in zip(all_data['documents'], all_data['metadatas'])]
 
+    # Mode 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
+    if bm25_fast: bm25_fast.k = 5
     
     fast_retriever = EnsembleRetriever(retrievers=[bm25_fast, vec_fast], weights=[0.4, 0.6]) if bm25_fast else vec_fast
 
+    # Mode 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
+    if bm25_deep: bm25_deep.k = 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)
     deep_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=ensemble)
     
     return fast_retriever, deep_retriever
 
+# --- 3. BOT LOGIC ---
 class DeepMedBot:
     def __init__(self):
         self.ready = False
+        if not GOOGLE_API_KEY:
+            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._build_chains()
             self.ready = True
         except Exception as e:
-            logging.error(f"Lỗi: {e}")
+            logging.error(f"Lỗi khởi tạo: {e}")
 
     def _build_chains(self):
-        # Prompt Bảng cho Thuốc
+        # Prompt Nhanh (Bảng 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"
+            "Bạn là Dược sĩ Lâm sàng.\n"
+            "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 | Đơn vị | Ghi chú |\n"
             "| --- | --- | --- | --- | --- |\n"
-            "Nếu không thấy, báo: '❌ Không có trong kho'."
+            "Nếu không thấy, báo: '❌ Không tìm thấy trong kho'."
             "Context:\n{context}"
         )
         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)
         
-        # Prompt Phác đồ ưu tiên Thanh Ba
+        # Prompt Chuyên sâu (Phác đồ Thanh Ba + Bảng)
         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"
@@ -105,7 +119,7 @@ class DeepMedBot:
         self.deep_chain = create_retrieval_chain(self.deep_retriever, deep_chain)
 
     def chat(self, msg, history, mode):
-        if not self.ready: return "⚠️ Đang khởi động... Vui lòng đợi 1 phút."
+        if not self.ready: return "⚠️ Đang khởi động hoặc Lỗi (Xem Logs trên Hugging Face)..."
         chain = self.deep_chain if mode == "Chuyên sâu" else self.fast_chain
         res = chain.invoke({"input": msg})
         
@@ -120,9 +134,13 @@ bot = DeepMedBot()
 def respond(message, history, mode):
     return bot.chat(message, history, mode)
 
-gr.ChatInterface(
+demo = 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ế độ")],
     title="TTYT Thanh Ba - Hỗ trợ Lâm sàng",
+    description="Hệ thống tra cứu Phác đồ & Thuốc nội bộ.",
     css=".gradio-container {min_height: 600px}"
-).launch()
+)
+
+if __name__ == "__main__":
+    demo.launch()