Update app.py

app.py

@@ -12,6 +12,7 @@ import chromadb
 from chromadb.config import Settings
 from shutil import rmtree
 
+# --- CÁC THƯ VIỆN LANGCHAIN ---
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
 from langchain_community.document_loaders import PyPDFLoader
@@ -25,22 +26,31 @@ from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.documents import Document
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.retrievers import ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 
+# --- THƯ VIỆN TỐI ƯU TỐC ĐỘ (CACHE & RERANK) ---
+from langchain.retrievers.document_compressors import FlashrankRerank
+from langchain.globals import set_llm_cache
+from langchain_community.cache import SQLiteCache
+
+# --- CẤU HÌNH HỆ THỐNG ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 DATA_PATH = "medical_data"
 DB_PATH = "chroma_db"
+CACHE_DB_PATH = "llm_cache.db" # File lưu bộ nhớ đệm
 MAX_HISTORY_TURNS = 6
 FORCE_REBUILD_DB = False
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 
+# --- KÍCH HOẠT CACHING ---
+# Hệ thống sẽ lưu câu trả lời vào file .db.
+# Lần sau gặp câu hỏi y hệt, nó sẽ lấy từ đệm ra ngay lập tức.
+if not os.path.exists(CACHE_DB_PATH):
+    logging.info("Khởi tạo file cache mới.")
+set_llm_cache(SQLiteCache(database_path=CACHE_DB_PATH))
+
 def process_excel_file(file_path: str, filename: str) -> list[Document]:
-    """
-    Xử lý Excel thông minh: Biến mỗi dòng thành một Document riêng biệt
-    giúp tìm kiếm chính xác từng bản ghi thuốc/bệnh nhân.
-    """
+    """Xử lý Excel: Biến mỗi dòng thành một Document."""
     docs = []
     try:
         if file_path.endswith(".csv"):
@@ -108,35 +118,32 @@ def load_documents_from_folder(folder_path: str) -> list[Document]:
 
 def get_retriever_chain():
     logging.info("--- Tải Embedding Model ---")
-    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    # Chạy trên CPU để tiết kiệm resource, đổi 'cpu' thành 'cuda' nếu có GPU
+    embedding_model = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        model_kwargs={'device': 'cpu'}
+    )
     
     vectorstore = None
-    splits = []
-
     chroma_settings = Settings(anonymized_telemetry=False)
 
     if FORCE_REBUILD_DB and os.path.exists(DB_PATH):
-        logging.warning("Đang xóa DB cũ theo yêu cầu FORCE_REBUILD...")
         rmtree(DB_PATH, ignore_errors=True)
 
+    # 1. TỐI ƯU: Kiểm tra nhanh DB bằng count() thay vì load toàn bộ
     if os.path.exists(DB_PATH) and os.listdir(DB_PATH):
         try:
             vectorstore = Chroma(
                 persist_directory=DB_PATH, 
                 embedding_function=embedding_model,
-                client_settings=chroma_settings 
+                client_settings=chroma_settings
             )
-            
-            existing_data = vectorstore.get()
-            if existing_data['documents']:
-                for text, meta in zip(existing_data['documents'], existing_data['metadatas']):
-                    splits.append(Document(page_content=text, metadata=meta))
-                logging.info(f"Đã khôi phục {len(splits)} chunks từ DB.")
+            if vectorstore._collection.count() > 0:
+                logging.info(f"Đã kết nối DB cũ. Size: {vectorstore._collection.count()}")
             else:
-                logging.warning("DB rỗng, sẽ tạo mới.")
                 vectorstore = None
         except Exception as e:
-            logging.error(f"DB lỗi: {e}. Đang reset...")
+            logging.error(f"DB lỗi: {e}. Reset DB...")
             rmtree(DB_PATH, ignore_errors=True)
             vectorstore = None
 
@@ -158,25 +165,16 @@ def get_retriever_chain():
         )
         logging.info("Đã lưu VectorStore thành công.")
 
-    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
+    # 2. TỐI ƯU: Giảm k ban đầu xuống 6 để bớt tính toán
+    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 6})
     
-    if splits:
-        bm25_retriever = BM25Retriever.from_documents(splits)
-        bm25_retriever.k = 10
-        ensemble_retriever = EnsembleRetriever(
-            retrievers=[bm25_retriever, vector_retriever],
-            weights=[0.4, 0.6] 
-        )
-    else:
-        ensemble_retriever = vector_retriever
+    # 3. TỐI ƯU: Sử dụng FlashRank (Siêu nhẹ & Nhanh) thay vì CrossEncoder
+    logging.info("--- Tải Reranker Model (FlashRank) ---")
+    compressor = FlashrankRerank(model="ms-marco-MiniLM-L-12-v2") # Model ~40MB
 
-    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)
-    
     final_retriever = ContextualCompressionRetriever(
         base_compressor=compressor, 
-        base_retriever=ensemble_retriever
+        base_retriever=vector_retriever 
     )
     
     return final_retriever
@@ -187,140 +185,28 @@ class DeepMedBot:
         self.ready = False
         
         if not GOOGLE_API_KEY:
-            logging.error("⚠️ Thiếu GOOGLE_API_KEY! Vui lòng thiết lập biến môi trường.")
+            logging.error("⚠️ Thiếu GOOGLE_API_KEY!")
             return
 
         try:
-            self.retriever = get_retriever_chain()
-            if not self.retriever:
-                logging.warning("⚠️ Chưa có dữ liệu để Retreive. Bot sẽ chỉ trả lời bằng kiến thức nền.")
-            
-            self.llm = ChatGoogleGenerativeAI(
-                model="gemini-2.5-flash", 
+            self.retr2.5-flash", 
                 temperature=0.3,
                 google_api_key=GOOGLE_API_KEY
             )
             self._build_chains()
             self.ready = True
-            logging.info("✅ Bot DeepMed đã sẵn sàng phục vụ!")
+            logging.info("✅ Bot DeepMed đã sẵn sàng!")
         except Exception as e:
             logging.error(f"🔥 Lỗi khởi tạo bot: {e}")
             logging.debug(traceback.format_exc())
 
     def _build_chains(self):
         context_system_prompt = (
-            "Dựa trên lịch sử chat 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 đầy đủ ngữ cảnh để hệ thống có thể hiểu được. "
-            "KHÔNG trả lời câu hỏi, chỉ viết lại nó."
+            "Viết lại câu hỏi của người dùng thành câu đầy đủ ngữ cảnh. "
+            "KHÔNG trả lời, chỉ viết lại."
         )
         context_prompt = ChatPromptTemplate.from_messages([
-            ("system", context_system_prompt),
-            MessagesPlaceholder("chat_history"),
-            ("human", "{input}"),
-        ])
-        
-        if self.retriever:
-            history_aware_retriever = create_history_aware_retriever(
-                self.llm, self.retriever, context_prompt
-            )
-        
-        qa_system_prompt = (
-            "Bạn là 'DeepMed-AI' - Trợ lý Dược lâm sàng tại Trung Tâm Y Tế. "
-            "Sử dụng các thông tin được cung cấp trong phần Context dưới đây để trả lời câu hỏi về thuốc, bệnh học và y lệnh.\n"
-            "Nếu Context có dữ liệu từ Excel, hãy trình bày dạng bảng hoặc gạch đầu dòng rõ ràng.\n"
-            "Nếu không tìm thấy thông tin trong Context, hãy nói 'Tôi không tìm thấy thông tin trong dữ liệu nội bộ' và gợi ý dựa trên kiến thức y khoa chung của bạn.\n\n"
-            "Context:\n{context}"
-        )
-        
-        qa_prompt = ChatPromptTemplate.from_messages([
-            ("system", qa_system_prompt),
-            MessagesPlaceholder("chat_history"),
-            ("human", "{input}"),
-        ])
-        
-        question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt)
-        
-        if self.retriever:
-            self.rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
-        else:
-            self.rag_chain = qa_prompt | self.llm
-
-    def chat_stream(self, message: str, history: list):
-        if not self.ready:
-            yield "Hệ thống đang khởi động hoặc gặp lỗi cấu hình."
-            return
-
-        chat_history = []
-        for u, b in history[-MAX_HISTORY_TURNS:]:
-            chat_history.append(HumanMessage(content=str(u)))
-            chat_history.append(AIMessage(content=str(b)))
-
-        full_response = ""
-        retrieved_docs = []
-        
-        try:
-            stream_input = {"input": message, "chat_history": chat_history} if self.retriever else {"input": message, "chat_history": chat_history}
-            
-            if self.rag_chain:
-                for chunk in self.rag_chain.stream(stream_input):
-                    
-                    if isinstance(chunk, dict):
-                        if "answer" in chunk:
-                            full_response += chunk["answer"]
-                            yield full_response 
-                        
-                        if "context" in chunk:
-                            retrieved_docs = chunk["context"]
-
-                    elif hasattr(chunk, 'content'):
-                        full_response += chunk.content
-                        yield full_response
-                    
-                    elif isinstance(chunk, str):
-                        full_response += chunk
-                        yield full_response
-
-                if retrieved_docs:
-                    refs = self._build_references_text(retrieved_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 khi chat: {e}")
-            logging.debug(traceback.format_exc())
-            yield f"Đã xảy ra lỗi: {str(e)}"
-
-    @staticmethod
-    def _build_references_text(docs) -> str:
-        lines = []
-        seen = set()
-        for doc in docs:
-            src = doc.metadata.get("source", "Tài liệu")
-            row_info = ""
-            if "row" in doc.metadata:
-                row_info = f"(Dòng {doc.metadata['row']})"
-            
-            ref_str = f"- {src} {row_info}"
-            
-            if ref_str not in seen:
-                lines.append(ref_str)
-                seen.add(ref_str)
-        return "\n".join(lines)
-
-bot = DeepMedBot()
-
-def gradio_chat_stream(message, history):
-    yield from bot.chat_stream(message, history)
-
-css = """
-.gradio-container {min_height: 600px !important;}
-h1 {text-align: center; color: #2E86C1;}
-"""
-
-with gr.Blocks(css=css, title="DeepMed AI") as demo:
-    gr.Markdown("# 🏥 DeepMed AI - Trợ lý Lâm Sàng")
-    gr.Markdown("Hệ thống hỗ trợ lâm sàng tại Trung Tâm Y Tế Khu Vực Thanh Ba.")
+           Ba)")
     
     chat_interface = gr.ChatInterface(
         fn=gradio_chat_stream,