Update app.py

app.py

@@ -1,162 +1,545 @@
-import os
-import sys
-import logging
-import gradio as gr
+"""
+DeepMed AI - Trợ lý Dược lâm sàng
+Hugging Face Spaces Ready (Free Tier)
+- Persistent storage bằng Dataset (không mất DB khi restart)
+- Reranker BGE-M3 (giữ nguyên theo yêu cầu)
+- Giao diện Gradio nâng cao
+"""
 
-# --- 1. SỬA LỖI SQLITE TRÊN HUGGING FACE ---
+# ---------- FIX SQLITE3 CHO CHROMA (QUAN TRỌNG TRÊN HF) ----------
 try:
     __import__("pysqlite3")
+    import sys
     sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
 except ImportError:
-    pass 
+    pass
+
+# ---------- THƯ VIỆN ----------
+import os
+import logging
+import pickle
+import shutil
+import traceback
+from typing import List
+
+import gradio as gr
+import pandas as pd
+import docx2txt
+from huggingface_hub import HfApi, login, snapshot_download, upload_file, upload_folder
 
-import chromadb
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-
-# --- IMPORT ĐƠN GIẢN HÓA (LOẠI BỎ CÁC MODULE GÂY LỖI _type) ---
-# Chỉ sử dụng các thành phần cốt lõi ổn định nhất
-from langchain.chains import create_retrieval_chain
+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.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+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
+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
+
+# ---------- CẤU HÌNH ----------
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 
-# --- CẤU HÌNH ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+DATA_PATH = "medical_data"
 DB_PATH = "chroma_db"
+SPLITS_CACHE = "splits_cache.pkl"
+FORCE_REBUILD_DB = False
+MAX_HISTORY_TURNS = 4
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
+# ---------- CẤU HÌNH DATASET PERSISTENT (MIỄN PHÍ) ----------
+HF_USERNAME = "your-username"          # 🔴 THAY BẰNG TÊN CỦA BẠN
+DATASET_NAME = "deepmed-db"
+DATASET_REPO = f"datasets/{HF_USERNAME}/{DATASET_NAME}"
 
-def get_category_vn_name(cat_code):
-    return {
-        "drug_info": "💊 Thuốc Nội Bộ",
-        "local_regimen": "🏥 Phác Đồ Thanh Ba",
-        "moh_regimen": "🏛️ Bộ Y Tế",
-        "association": "🌐 Hiệp Hội"
-    }.get(cat_code, "Khác")
+# ---------- XÁC THỰC HUGGING FACE ----------
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+    logging.info("✅ Logged into Hugging Face Hub")
+else:
+    logging.warning("⚠️ HF_TOKEN not found. Dataset persistence disabled.")
 
-# --- 2. LOAD DB (VECTOR SEARCH THUẦN TÚY - ỔN ĐỊNH 100%) ---
-def get_retrievers():
-    if not os.path.exists(DB_PATH):
-        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?")
+# ---------- HÀM PERSISTENT STORAGE ----------
+def download_persistent_data():
+    """Tải Chroma DB và splits cache từ Dataset về local (gọi khi khởi động)"""
+    if not HF_TOKEN:
+        return False
+    try:
+        # Tải Chroma DB
+        if not os.path.exists(DB_PATH):
+            logging.info("📥 Đang tải Chroma DB từ Dataset...")
+            snapshot_download(
+                repo_id=DATASET_REPO,
+                allow_patterns=f"{DB_PATH}/*",
+                local_dir=".",
+                local_dir_use_symlinks=False,
+                token=HF_TOKEN,
+                ignore_patterns=["*.gitattributes", "README.md"]
+            )
+            # Nếu tải về thư mục tạm (tên là DB_PATH) thì move vào đúng
+            if os.path.exists(DB_PATH):
+                logging.info("✅ Đã tải Chroma DB.")
+        
+        # Tải splits cache
+        if not os.path.exists(SPLITS_CACHE):
+            logging.info("📥 Đang tải splits_cache.pkl...")
+            snapshot_download(
+                repo_id=DATASET_REPO,
+                allow_patterns=SPLITS_CACHE,
+                local_dir=".",
+                local_dir_use_symlinks=False,
+                token=HF_TOKEN,
+                ignore_patterns=["*.gitattributes", "README.md"]
+            )
+            if os.path.exists(SPLITS_CACHE):
+                logging.info("✅ Đã tải splits_cache.pkl.")
+        return True
+    except Exception as e:
+        logging.warning(f"⚠️ Không tải được dữ liệu từ Dataset (có thể lần đầu chạy): {e}")
+        return False
 
-    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)
-    
-    # Kiểm tra dữ liệu
+def upload_persistent_data():
+    """Upload Chroma DB và splits cache lên Dataset (gọi sau khi rebuild DB)"""
+    if not HF_TOKEN:
+        return
     try:
-        all_data = vectorstore.get()
-        if not all_data['documents']:
-            raise ValueError("Database rỗng")
-        logging.info(f"✅ Đã tải thành công {len(all_data['documents'])} tài liệu từ Database.")
+        # Upload Chroma DB
+        if os.path.exists(DB_PATH):
+            logging.info("📤 Đang upload Chroma DB lên Dataset...")
+            upload_folder(
+                folder_path=DB_PATH,
+                repo_id=DATASET_REPO,
+                repo_type="dataset",
+                path_in_repo=DB_PATH,
+                token=HF_TOKEN,
+                ignore_patterns=[".gitattributes", "README.md"]
+            )
+            logging.info("✅ Upload Chroma DB thành công.")
+        
+        # Upload splits cache
+        if os.path.exists(SPLITS_CACHE):
+            logging.info("📤 Đang upload splits_cache.pkl...")
+            upload_file(
+                path_or_fileobj=SPLITS_CACHE,
+                path_in_repo=SPLITS_CACHE,
+                repo_id=DATASET_REPO,
+                repo_type="dataset",
+                token=HF_TOKEN
+            )
+            logging.info("✅ Upload splits_cache.pkl thành công.")
     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}")
-
-    # --- TẠO RETRIEVER ĐƠN GIẢN ---
-    # Thay vì dùng Ensemble/Reranker (dễ lỗi), ta dùng Vector Search trực tiếp.
-    
-    # Mode 1: FAST (Tìm kiếm Thuốc - Lấy 5 kết quả sát nhất)
-    logging.info("--- Khởi tạo Fast Retriever (Vector Only) ---")
-    fast_retriever = vectorstore.as_retriever(
-        search_kwargs={
-            "k": 5, 
-            "filter": {"category": "drug_info"}
-        }
-    )
+        logging.error(f"❌ Upload thất bại: {e}")
 
-    # Mode 2: DEEP (Tìm kiếm Phác đồ - Lấy 15 kết quả sát nhất)
-    # Tăng k lên để bù đắp việc thiếu Reranker
-    logging.info("--- Khởi tạo Deep Retriever (Vector Only) ---")
-    cats = ["local_regimen", "moh_regimen", "association", "drug_info"]
-    deep_retriever = vectorstore.as_retriever(
-        search_kwargs={
-            "k": 15, 
-            "filter": {"category": {"$in": cats}}
-        }
-    )
+# Gọi download ngay khi khởi chạy app
+download_persistent_data()
+
+# ---------- XỬ LÝ DOCUMENTS ----------
+def process_excel_file(file_path: str, filename: str) -> List[Document]:
+    docs = []
+    try:
+        if file_path.endswith(".csv"):
+            df = pd.read_csv(file_path)
+        else:
+            df = pd.read_excel(file_path)
+        df.dropna(how='all', inplace=True)
+        df.fillna("Không có thông tin", inplace=True)
+        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"Dữ liệu từ file {filename} (Dòng {idx+1}):\n" + "\n".join(content_parts)
+                metadata = {"source": filename, "row": idx+1, "type": "excel_record"}
+                docs.append(Document(page_content=page_content, metadata=metadata))
+    except Exception as e:
+        logging.error(f"Lỗi xử lý Excel {filename}: {e}")
+    return docs
+
+def load_documents_from_folder(folder_path: str) -> List[Document]:
+    logging.info(f"--- Quét thư mục: {folder_path} ---")
+    documents = []
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path, exist_ok=True)
+        return []
+    for root, _, files in os.walk(folder_path):
+        for filename in files:
+            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")):
+                    excel_docs = process_excel_file(file_path, filename)
+                    documents.extend(excel_docs)
+                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 đọc file {filename}: {e}")
+    logging.info(f"Tổng số tài liệu gốc: {len(documents)}")
+    return documents
 
+def load_or_create_splits(raw_docs):
+    """Load splits từ cache nếu có, nếu không thì chunk và cache lại"""
+    if os.path.exists(SPLITS_CACHE) and not FORCE_REBUILD_DB:
+        logging.info("--- Load splits từ cache ---")
+        with open(SPLITS_CACHE, "rb") as f:
+            return pickle.load(f)
+    logging.info("--- Tạo splits mới (chunk_size=800, overlap=150) ---")
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+    splits = text_splitter.split_documents(raw_docs)
+    with open(SPLITS_CACHE, "wb") as f:
+        pickle.dump(splits, f)
+    return splits
+
+# ---------- RETRIEVERS ----------
+def get_retrievers():
+    logging.info("--- Tải Embedding Model ---")
+    embedding_model = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+    )
+    
+    # Khởi tạo Chroma (load nếu có, tạo mới nếu không)
+    vectorstore = None
+    splits = []
+    
+    if os.path.exists(DB_PATH) and os.listdir(DB_PATH) and not FORCE_REBUILD_DB:
+        logging.info("--- Phát hiện Chroma DB cũ, đang tải... ---")
+        try:
+            vectorstore = Chroma(
+                persist_directory=DB_PATH,
+                embedding_function=embedding_model,
+            )
+            # Load splits từ cache
+            if os.path.exists(SPLITS_CACHE):
+                with open(SPLITS_CACHE, "rb") as f:
+                    splits = pickle.load(f)
+            else:
+                # Trường hợp hiếm: có DB nhưng mất cache, phải load lại raw
+                raw_docs = load_documents_from_folder(DATA_PATH)
+                if raw_docs:
+                    splits = load_or_create_splits(raw_docs)
+        except Exception as e:
+            logging.error(f"Lỗi load Chroma DB: {e}. Tiến hành tạo mới.")
+            vectorstore = None
+    
+    if vectorstore is None:
+        logging.info("--- Tạo Index dữ liệu mới ---")
+        raw_docs = load_documents_from_folder(DATA_PATH)
+        if not raw_docs:
+            logging.warning("⚠️ Không có tài liệu nào trong thư mục data.")
+            return None, None
+        splits = load_or_create_splits(raw_docs)
+        vectorstore = Chroma.from_documents(
+            documents=splits,
+            embedding=embedding_model,
+            persist_directory=DB_PATH
+        )
+        # Upload dữ liệu mới lên Dataset
+        upload_persistent_data()
+    
+    # === FAST RETRIEVER (Ensemble, k=8) ===
+    bm25_fast = BM25Retriever.from_documents(splits)
+    bm25_fast.k = 8
+    vector_fast = vectorstore.as_retriever(search_kwargs={"k": 8})
+    fast_retriever = EnsembleRetriever(
+        retrievers=[bm25_fast, vector_fast],
+        weights=[0.5, 0.5]
+    )
+    
+    # === DEEP RETRIEVER (GIỮ NGUYÊN bge-reranker-v2-m3) ===
+    bm25_deep = BM25Retriever.from_documents(splits)
+    bm25_deep.k = 12
+    vector_deep = vectorstore.as_retriever(search_kwargs={"k": 12})
+    ensemble_deep = EnsembleRetriever(
+        retrievers=[bm25_deep, vector_deep],
+        weights=[0.5, 0.5]
+    )
+    
+    logging.info("--- Tải CrossEncoderReranker (bge-reranker-v2-m3) ---")
+    reranker_model = HuggingFaceCrossEncoder(
+        model_name="BAAI/bge-reranker-v2-m3",
+        model_kwargs={'device': 'cpu', 'low_cpu_mem_usage': True}
+    )
+    compressor = CrossEncoderReranker(model=reranker_model, top_n=3)
+    deep_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor,
+        base_retriever=ensemble_deep
+    )
+    
     return fast_retriever, deep_retriever
 
-# --- 3. BOT LOGIC ---
+# ---------- DEEPMED BOT ----------
 class DeepMedBot:
     def __init__(self):
+        self.fast_chain = None
+        self.deep_chain = None
         self.ready = False
-        self.init_error = "Đang khởi động..." 
+        self.fallback_llm = None
         
         if not GOOGLE_API_KEY:
-            self.init_error = "❌ LỖI: Chưa cấu hình GOOGLE_API_KEY trong Settings."
+            logging.error("⚠️ Thiếu 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
-            self.init_error = ""
-            logging.info("✅ BOT KHỞI ĐỘNG THÀNH CÔNG (CHẾ ĐỘ VECTOR STABLE)!")
+            self.llm = ChatGoogleGenerativeAI(
+                model="gemini-1.5-flash",   # ổn định và nhanh
+                temperature=0.2,
+                google_api_key=GOOGLE_API_KEY,
+                convert_system_message_to_human=True
+            )
+            self.fallback_llm = self.llm
+            
+            if self.fast_retriever and self.deep_retriever:
+                self._build_chains()
+                self.ready = True
+                logging.info("✅ Bot DeepMed đã sẵn sàng với reranker v2-m3!")
+            else:
+                logging.warning("⚠️ Không có retriever, chỉ dùng kiến thức nền.")
+                self.ready = True
         except Exception as e:
-            self.init_error = f"❌ LỖI KHỞI TẠO: {str(e)}"
-            logging.error(self.init_error)
-
+            logging.error(f"🔥 Lỗi khởi tạo bot: {e}")
+            logging.debug(traceback.format_exc())
+    
     def _build_chains(self):
-        # 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"
-            "| 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 tìm thấy trong kho'."
-            "Context:\n{context}"
+        # Prompt viết lại câu hỏi
+        context_system_prompt = (
+            "Dựa trên lịch sử chat 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 để tìm kiếm thông tin. "
+            "CHỈ TRẢ VỀ CÂU HỎI ĐÃ VIẾT LẠI, KHÔNG TRẢ LỜI."
         )
-        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)
+        context_prompt = ChatPromptTemplate.from_messages([
+            ("system", context_system_prompt),
+            MessagesPlaceholder("chat_history"),
+            ("human", "{input}"),
+        ])
         
-        # 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"
-            "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"
+        # Prompt trả lời
+        qa_system_prompt = (
+            "Bạn là 'DeepMed-AI' - Trợ lý Dược lâm sàng chuyên nghiệp.\n"
+            "Nhiệm vụ: Tư vấn điều trị CHỈ DỰA TRÊN Dữ liệu nội bộ (Context) được cung cấp.\n\n"
+            "QUY TẮC AN TOÀN:\n"
+            "1. Nếu thông tin không có trong Context, trả lời: 'Xin lỗi, tôi không tìm thấy thông tin này trong dữ liệu nội bộ'.\n"
+            "2. Chỉ đề xuất thuốc có trong danh sách Context.\n"
+            "3. Mọi khẳng định phải trích dẫn từ Context.\n\n"
             "Context:\n{context}"
         )
-        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 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
+        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)
+        
+        history_aware_fast = create_history_aware_retriever(self.llm, self.fast_retriever, context_prompt)
+        self.fast_chain = create_retrieval_chain(history_aware_fast, question_answer_chain)
+        
+        history_aware_deep = create_history_aware_retriever(self.llm, self.deep_retriever, context_prompt)
+        self.deep_chain = create_retrieval_chain(history_aware_deep, question_answer_chain)
+    
+    def chat_stream(self, message: str, history: list, mode: str):
+        if not self.ready:
+            yield "Hệ thống đang khởi động hoặc gặp lỗi cấu hình..."
+            return
+        
+        # Xử lý lịch sử
+        chat_history = []
+        if history:
+            for turn in history[-MAX_HISTORY_TURNS:]:
+                if isinstance(turn, (list, tuple)) and len(turn) == 2:
+                    u, b = turn
+                    if u and b and str(u).strip() and str(b).strip():
+                        chat_history.append(HumanMessage(content=str(u)))
+                        chat_history.append(AIMessage(content=str(b)))
+        
+        active_chain = self.deep_chain if "Chuyên sâu" in mode else self.fast_chain
+        
+        if not active_chain:
+            try:
+                resp = self.llm.invoke([HumanMessage(content=message)])
+                yield f"⚠️ (Chế độ kiến thức chung) {resp.content}"
+                return
+            except:
+                yield "Lỗi: Không thể kết nối với AI. Vui lòng kiểm tra API Key."
+                return
+        
+        full_response = ""
+        retrieved_docs = []
+        
         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
+            for chunk in active_chain.stream({"input": message, "chat_history": chat_history}):
+                if "answer" in chunk:
+                    full_response += chunk["answer"]
+                    yield full_response
+                elif "context" in chunk:
+                    retrieved_docs = chunk["context"]
+            
+            # Thêm phần nguồn tham khảo
+            if retrieved_docs:
+                refs = self._build_references_text(retrieved_docs)
+                if refs:
+                    full_response += f"\n\n---\n📚 **Nguồn tham khảo ({mode}):**\n{refs}"
+                    yield full_response
         except Exception as e:
-            return f"❌ Lỗi khi trả lời: {str(e)}"
+            logging.error(f"Lỗi khi chat: {e}")
+            logging.error(traceback.format_exc())
+            if not full_response:
+                try:
+                    yield "⚠️ Gặp lỗi truy xuất dữ liệu. Đang chuyển sang chế độ trả lời nhanh...\n\n"
+                    fallback_resp = self.llm.invoke([HumanMessage(content=message)])
+                    yield fallback_resp.content
+                except:
+                    yield f"Đã xảy ra lỗi hệ thống. Vui lòng thử lại. (Lỗi: {str(e)})"
+            else:
+                yield full_response + f"\n\n[Lỗi ngắt kết nố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 = f"(Dòng {doc.metadata['row']})" if "row" in doc.metadata else ""
+            type_info = " [Kho thuốc]" if doc.metadata.get("type") == "excel_record" else ""
+            ref_str = f"- {src}{type_info} {row_info}"
+            if ref_str not in seen:
+                lines.append(ref_str)
+                seen.add(ref_str)
+        return "\n".join(lines)
 
+# Khởi tạo bot
 bot = DeepMedBot()
 
-def respond(message, history, mode):
-    return bot.chat(message, history, mode)
-
-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}"
+# ---------- GRADIO UI ----------
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="emerald",
+    neutral_hue="gray",
+    font=gr.themes.GoogleFont("Inter")
 )
 
+css = """
+footer {visibility: hidden}
+.gr-chatbot .user-message {background-color: #e6f7ff}
+.gr-chatbot .bot-message {background-color: #f0f2f6}
+"""
+
+def chat_handler(message, history, mode):
+    """Generator cho Gradio ChatInterface"""
+    yield from bot.chat_stream(message, history, mode)
+
+with gr.Blocks(theme=theme, css=css, title="DeepMed AI") as demo:
+    gr.Markdown("""
+    # 🏥 DeepMed AI - Trợ lý Dược lâm sàng
+    **Hệ thống tra cứu phác đồ, thuốc, bệnh án nội bộ**  
+    Chạy trên nền tảng Gemini + RAG với reranker BGE-M3.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            mode_select = gr.Radio(
+                choices=["⚡ Tốc độ (Nhanh)", "🔍 Chuyên sâu (Chính xác)"],
+                value="⚡ Tốc độ (Nhanh)",
+                label="Chế độ tra cứu",
+                info="'Tốc độ' dùng ensemble, 'Chuyên sâu' dùng thêm reranker BGE-M3."
+            )
+        with gr.Column(scale=1):
+            clear_btn = gr.ClearButton(value="🗑️ Xoá chat", size="sm")
+    
+    # Chatbot với avatar
+    chatbot = gr.Chatbot(
+        avatar_images=("🧑‍⚕️", "🤖"),
+        height=550,
+        show_copy_button=True,
+        bubble_full_width=False,
+        layout="panel"
+    )
+    
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Nhập câu hỏi (VD: 'Phác đồ điều trị tăng huyết áp?', 'Thuốc Paracetamol giá bao nhiêu?')",
+            scale=9,
+            container=False
+        )
+        submit = gr.Button("📨 Gửi", variant="primary", scale=1, min_width=100)
+    
+    # Accordion hiển thị nguồn tham khảo
+    with gr.Accordion("📚 Nguồn tham khảo", open=False):
+        ref_markdown = gr.Markdown("_Chưa có nguồn trích dẫn._")
+    
+    # Tab quản lý dữ liệu (upload file)
+    with gr.Tab("📂 Quản lý dữ liệu"):
+        gr.Markdown("Tải lên tài liệu PDF, Excel, Word, TXT để cập nhật cơ sở tri thức.")
+        file_upload = gr.File(
+            file_count="multiple",
+            label="Chọn file",
+            file_types=[".pdf", ".xlsx", ".xls", ".csv", ".docx", ".txt", ".md"]
+        )
+        upload_btn = gr.Button("⬆️ Cập nhật dữ liệu", variant="secondary")
+        upload_status = gr.Textbox(label="Trạng thái", interactive=False)
+        
+        def upload_files(files):
+            if not files:
+                return "❌ Chưa chọn file nào."
+            os.makedirs(DATA_PATH, exist_ok=True)
+            for f in files:
+                # Lưu file vào thư mục dữ liệu
+                f.save(os.path.join(DATA_PATH, f.name))
+            # Cảnh báo: cần rebuild lại DB để áp dụng dữ liệu mới
+            return "✅ Đã tải lên. Vui lòng **khởi động lại Space** (Settings → Restart this Space) để rebuild database."
+        
+        upload_btn.click(upload_files, inputs=file_upload, outputs=upload_status)
+    
+    # Xử lý chat
+    def respond(message, chat_history, mode):
+        bot_response = ""
+        for chunk in bot.chat_stream(message, chat_history, mode):
+            bot_response = chunk
+        chat_history.append((message, bot_response))
+        # Trích xuất phần nguồn tham khảo từ bot_response để hiển thị riêng
+        refs = ""
+        if "📚 **Nguồn tham khảo**" in bot_response:
+            parts = bot_response.split("---\n📚 **Nguồn tham khảo**")
+            bot_response = parts[0].strip()
+            refs = "📚 **Nguồn tham khảo**" + parts[1]
+        return chat_history, bot_response, refs
+    
+    submit.click(
+        respond,
+        inputs=[msg, chatbot, mode_select],
+        outputs=[chatbot, msg, ref_markdown]
+    ).then(lambda: "", None, msg)  # Xóa textbox sau khi gửi
+    
+    msg.submit(
+        respond,
+        inputs=[msg, chatbot, mode_select],
+        outputs=[chatbot, msg, ref_markdown]
+    ).then(lambda: "", None, msg)
+    
+    clear_btn.click(lambda: ([], "", "_Chưa có nguồn trích dẫn._"), None, [chatbot, msg, ref_markdown])
+
+# ---------- KHỞI CHẠY APP ----------
 if __name__ == "__main__":
     demo.launch()