app.py

"""
DeepMed AI - Trợ lý Dược lâm sàng (Phiên bản Tra cứu)
Hugging Face Spaces Ready
- Chức năng: Chỉ tra cứu (Read-only), không cho phép upload từ UI.
- Dữ liệu: Tự động tải từ Hugging Face Dataset khi khởi động.
"""

# ---------- 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

# ---------- THƯ VIỆN ----------
import os
import logging
import pickle
import traceback
from typing import List

import gradio as gr
import pandas as pd
import docx2txt
from huggingface_hub import login, snapshot_download, upload_file, upload_folder

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_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, 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")

# Lấy Key từ Environment Variables (Cài đặt trong Settings của Space)
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
HF_TOKEN = os.getenv("HF_TOKEN")

DATA_PATH = "medical_data"     # Thư mục chứa dữ liệu gốc (nếu cần rebuild)
DB_PATH = "chroma_db"          # Thư mục chứa Vector DB
SPLITS_CACHE = "splits_cache.pkl" # Cache file đã chunk
FORCE_REBUILD_DB = False       # Set True nếu muốn buộc tạo lại DB từ đầu
MAX_HISTORY_TURNS = 4          # Số lượt hội thoại được nhớ

# ---------- CẤU HÌNH DATASET PERSISTENT ----------
# 🔴 THAY ĐỔI DÒNG DƯỚI ĐÂY THÀNH USERNAME CỦA BẠN
HF_USERNAME = "PBThuong96"  
DATASET_NAME = "deepmed-db"
DATASET_REPO = f"{HF_USERNAME}/{DATASET_NAME}"

# ---------- 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.")

# ---------- HÀM PERSISTENT STORAGE (Sync Dữ Liệu) ----------
def download_persistent_data():
    """Tải Chroma DB và splits cache từ Dataset về local (gọi khi khởi động)"""
    
    # 1. Kiểm tra Username (Lỗi phổ biến nhất)
    if HF_USERNAME == "your-username":
        logging.error("❌ CẤU HÌNH LỖI: Bạn chưa đổi 'HF_USERNAME' trong code (dòng 55). Vui lòng sửa lại tên tài khoản HF của bạn.")
        return False

    if not HF_TOKEN:
        logging.warning("⚠️ CHƯA CÓ TOKEN: Biến môi trường 'HF_TOKEN' chưa được cài đặt. Không thể tải dữ liệu.")
        return False

    logging.info(f"🔄 Bắt đầu đồng bộ dữ liệu từ Dataset: {DATASET_REPO}")

    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,
                repo_type="dataset",
                # Sử dụng pattern ** để tìm đệ quy, tránh lỗi cấu trúc folder
                allow_patterns=f"{DB_PATH}/**", 
                local_dir=".",
                local_dir_use_symlinks=False,
                token=HF_TOKEN,
                ignore_patterns=["*.gitattributes", "README.md"]
            )
        
        # 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,
                repo_type="dataset",
                allow_patterns=SPLITS_CACHE,
                local_dir=".",
                local_dir_use_symlinks=False,
                token=HF_TOKEN,
                ignore_patterns=["*.gitattributes", "README.md"]
            )

        # Kiểm tra kết quả sau khi tải
        if os.path.exists(DB_PATH) and os.listdir(DB_PATH):
            logging.info("✅ Tải dữ liệu thành công!")
            return True
        else:
            logging.warning("⚠️ Đã chạy lệnh tải nhưng thư mục DB vẫn rỗng hoặc không tồn tại. Có thể Dataset của bạn đang trống?")
            return False

    except Exception as e:
        error_msg = str(e)
        logging.error(f"❌ LỖI TẢI DỮ LIỆU: {error_msg}")
        
        if "404" in error_msg:
            logging.error(f"👉 Không tìm thấy Dataset '{DATASET_REPO}'. Hãy kiểm tra:\n   1. Bạn đã tạo Dataset trên Hugging Face chưa?\n   2. Tên Username và Dataset Name trong code có đúng không?")
        elif "401" in error_msg or "403" in error_msg:
            logging.error("👉 Lỗi quyền truy cập (Auth). Hãy kiểm tra HF_TOKEN trong Settings của Space.")
            
        return False

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:
        # 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"❌ Upload thất bại: {e}")

# Gọi download ngay khi khởi chạy app
download_persistent_data()

# ---------- XỬ LÝ DOCUMENTS (Hỗ trợ Rebuild Local) ----------
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):
    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 & CORE LOGIC ----------
def get_retrievers():
    logging.info("--- Tải Embedding Model ---")
    embedding_model = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    )
    
    vectorstore = None
    splits = []
    
    # 1. Thử load DB từ Disk (đã được download từ Dataset)
    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,
            )
            if os.path.exists(SPLITS_CACHE):
                with open(SPLITS_CACHE, "rb") as f:
                    splits = pickle.load(f)
            else:
                # Nếu mất cache, load lại từ raw folder (nếu có)
                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
    
    # 2. Nếu không load được DB, tạo mới từ folder 'medical_data'
    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 & không tải được DB từ Dataset.")
            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 để lần sau dùng
        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 (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

# ---------- DEEPMED BOT ----------
class DeepMedBot:
    def __init__(self):
        self.fast_chain = None
        self.deep_chain = None
        self.ready = False
        
        if not GOOGLE_API_KEY:
            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,
                convert_system_message_to_human=True
            )
            
            if self.fast_retriever and self.deep_retriever:
                self._build_chains()
                self.ready = True
                logging.info("✅ Bot DeepMed đã sẵn sàng!")
            else:
                logging.warning("⚠️ Không có retriever (Chưa có dữ liệu).")
                self.ready = True
        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, 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."
        )
        context_prompt = ChatPromptTemplate.from_messages([
            ("system", context_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ])
        
        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}"
        )
        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 chưa có dữ liệu..."
            return
        
        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:
                        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 - Chưa có DB) {resp.content}"
                return
            except:
                yield "Lỗi kết nối AI hoặc thiếu dữ liệu."
                return
        
        full_response = ""
        retrieved_docs = []
        
        try:
            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"]
            
            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:
            logging.error(f"Lỗi khi chat: {e}")
            yield full_response + f"\n\n[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 = 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()

# ---------- 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}
"""

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"],
                value="⚡ Tốc độ (Nhanh)",
                label="Chế độ tra cứu",
            )
        with gr.Column(scale=1):
            clear_btn = gr.ClearButton(value="🗑️ Xoá chat", size="sm")
    
    chatbot = gr.Chatbot(
        avatar_images=("🧑‍⚕️", "🤖"),
        height=600,
        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)
    
    with gr.Accordion("📚 Nguồn tham khảo chi tiết", open=False):
        ref_markdown = gr.Markdown("_Chưa có nguồn trích dẫn._")
    
    # Xử lý sự kiện
    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))
        
        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)
    
    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()