Update app.py

app.py

@@ -3,25 +3,25 @@ import sys
 import logging
 import gradio as gr
 
-# --- 1. SỬA LỖI SQLITE TRÊN HUGGING FACE (BẮT BUỘC ĐỂ ĐẦU FILE) ---
+# --- 1. SỬA LỖI SQLITE ---
 try:
     __import__("pysqlite3")
     sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 except ImportError:
-    pass # Nếu chạy local không có pysqlite3 thì bỏ qua
+    pass 
 
 import chromadb
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
+# CẬP NHẬT IMPORT: Import trực tiếp từ langchain.retrievers để tránh lỗi ModuleNotFoundError
+from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
 from langchain_community.retrievers import BM25Retriever
-from langchain.retrievers.ensemble import EnsembleRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.chains import create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 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 ---
@@ -38,58 +38,75 @@ 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) ---
+# --- 2. LOAD DB VỚI CƠ CHẾ AN TOÀN (SAFE LOAD) ---
 def get_retrievers():
     if not os.path.exists(DB_PATH):
-        # Lỗi phổ biến nhất: Quên upload hoặc upload sai chỗ
         raise FileNotFoundError(f"❌ LỖI: Không tìm thấy thư mục '{DB_PATH}'. Bạn đã upload folder này vào phần Files chưa?")
 
-    logging.info("--- Đang tải dữ liệu... ---")
+    logging.info("--- Đang tải dữ liệu từ ChromaDB... ---")
     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()
-    
-    # Kiểm tra nếu DB rỗng
-    if not all_data['documents']:
-        raise ValueError("❌ LỖI: Database rỗng! Có thể quá trình build_db ở máy local bị lỗi hoặc chưa có file nào được xử lý.")
-
-    splits = [Document(page_content=txt, metadata=m) for txt, m in zip(all_data['documents'], all_data['metadatas'])]
+    # Kiểm tra dữ liệu
+    try:
+        all_data = vectorstore.get()
+        if not all_data['documents']:
+            raise ValueError("Database rỗng")
+        # Tái tạo Documents để dùng cho BM25
+        splits = [Document(page_content=txt, metadata=m) for txt, m in zip(all_data['documents'], all_data['metadatas'])]
+    except Exception as e:
+        logging.error(f"Lỗi đọc dữ liệu Chroma: {e}")
+        raise ValueError(f"Không thể đọc dữ liệu từ ChromaDB: {e}")
+
+    # --- HÀM TẠO RETRIEVER AN TOÀN ---
+    def create_safe_retriever(k_val, filter_dict, doc_subset):
+        # 1. Luôn tạo Vector Retriever (Cái này ít lỗi nhất)
+        vec_retriever = vectorstore.as_retriever(search_kwargs={"k": k_val, "filter": filter_dict})
+        
+        # 2. Cố gắng tạo Hybrid (BM25 + Vector)
+        try:
+            if not doc_subset:
+                return vec_retriever # Không có docs thì dùng vector thôi
+            
+            bm25 = BM25Retriever.from_documents(doc_subset)
+            bm25.k = k_val
+            # Thử tạo Ensemble (Đây là chỗ hay gây lỗi '_type')
+            ensemble = EnsembleRetriever(retrievers=[bm25, vec_retriever], weights=[0.4, 0.6])
+            return ensemble
+        except Exception as e:
+            logging.warning(f"⚠️ Không thể tạo Hybrid Search (Lỗi: {e}). Đang chuyển sang dùng Vector Search thuần túy.")
+            return vec_retriever # Fallback về Vector nếu Hybrid lỗi
 
-    # Mode 1: FAST (Chỉ thuốc)
-    vec_fast = vectorstore.as_retriever(search_kwargs={"k": 5, "filter": {"category": "drug_info"}})
+    # Mode 1: FAST
+    logging.info("--- Khởi tạo Fast Retriever ---")
     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
-    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
+    fast_retriever = create_safe_retriever(5, {"category": "drug_info"}, drug_docs)
 
-    # Mode 2: DEEP (Ưu tiên Thanh Ba)
+    # Mode 2: DEEP
+    logging.info("--- Khởi tạo Deep Retriever ---")
     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
-    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)
+    base_deep_retriever = create_safe_retriever(25, {"category": {"$in": cats}}, deep_docs)
     
+    # 3. Thử tạo Reranker (Cũng có thể gây lỗi '_type')
+    try:
+        reranker = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
+        compressor = CrossEncoderReranker(model=reranker, top_n=10)
+        deep_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=base_deep_retriever)
+    except Exception as e:
+        logging.warning(f"⚠️ Không thể tải Reranker (Lỗi: {e}). Dùng bộ tìm kiếm cơ bản.")
+        deep_retriever = base_deep_retriever
+
     return fast_retriever, deep_retriever
 
 # --- 3. BOT LOGIC ---
 class DeepMedBot:
     def __init__(self):
         self.ready = False
-        self.init_error = "Đang khởi động..." # Lưu lỗi để hiển thị ra màn hình chat
+        self.init_error = "Đang khởi động..." 
         
         if not GOOGLE_API_KEY:
-            self.init_error = "❌ LỖI: Chưa có API Key.\nVui lòng vào Settings -> Variables and secrets -> New secret.\nName: GOOGLE_API_KEY\nValue: (Key của bạn)"
-            logging.error(self.init_error)
+            self.init_error = "❌ LỖI: Chưa cấu hình GOOGLE_API_KEY trong Settings."
             return
 
         try:
@@ -97,13 +114,14 @@ class DeepMedBot:
             self.llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.2, google_api_key=GOOGLE_API_KEY)
             self._build_chains()
             self.ready = True
-            self.init_error = "" # Xóa lỗi nếu thành công
+            self.init_error = ""
+            logging.info("✅ BOT KHỞI ĐỘNG THÀNH CÔNG!")
         except Exception as e:
-            self.init_error = f"❌ LỖI KHỞI TẠO: {str(e)}"
+            self.init_error = f"❌ LỖI KHỞI TẠO NGHIÊM TRỌNG: {str(e)}"
             logging.error(self.init_error)
 
     def _build_chains(self):
-        # Prompt Nhanh (Bảng Thuốc)
+        # Prompt Nhanh
         fast_sys = (
             "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"
@@ -115,7 +133,7 @@ class DeepMedBot:
         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 Chuyên sâu (Phác đồ Thanh Ba + Bảng)
+        # Prompt Chuyên sâu
         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"
@@ -132,18 +150,19 @@ class DeepMedBot:
         self.deep_chain = create_retrieval_chain(self.deep_retriever, deep_chain)
 
     def chat(self, msg, history, mode):
-        # Nếu chưa sẵn sàng, trả về chính xác lỗi gì cho người dùng biết
         if not self.ready: 
-            return f"⚠️ HỆ THỐNG CHƯA SẴN SÀNG.\n\nNguyên nhân lỗi:\n{self.init_error}\n\nVui lòng kiểm tra lại cấu hình trên Hugging Face."
+            return f"⚠️ HỆ THỐNG GẶP LỖI.\n\nChi tiết lỗi:\n{self.init_error}\n\nHãy thử Restart Space trong phần Settings."
             
         chain = self.deep_chain if mode == "Chuyên sâu" else self.fast_chain
-        res = chain.invoke({"input": msg})
-        
-        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
+        try:
+            res = chain.invoke({"input": msg})
+            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
+        except Exception as e:
+            return f"❌ Lỗi khi trả lời: {str(e)}"
 
 bot = DeepMedBot()