app.py

import os
try:
    __import__("pysqlite3")
    import sys
    sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
except ImportError:
    pass

import logging
import traceback
import pandas as pd
import docx2txt
import chromadb
from chromadb.config import Settings
import chainlit as cl # Thư viện giao diện mới

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.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
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

# === CẤU HÌNH ===
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
DATA_PATH = "medical_data"
DB_PATH = "chroma_db"
MAX_HISTORY_TURNS = 4

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# ... (GIỮ NGUYÊN CÁC HÀM: process_excel_file, load_documents_from_folder, get_retrievers) ...
# Để tiết kiệm chỗ hiển thị, tôi không paste lại đoạn code xử lý file ở đây
# Bạn hãy copy-paste lại các hàm process_excel_file, load_documents_from_folder, get_retrievers Y HỆT NHƯ CŨ vào đây.

# === LOGIC BOT (Sửa nhẹ để tách biệt khởi tạo) ===
class DeepMedBot:
    def __init__(self):
        self.fast_chain = None
        self.deep_chain = None
        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,
                convert_system_message_to_human=True
            )
            
            if self.fast_retriever and self.deep_retriever:
                self._build_chains()
                self.ready = True
            else:
                 self.ready = True 

        except Exception as e:
            logging.error(f"Lỗi khởi tạo bot: {e}")

    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).\n"
            "QUY TẮC: Trung thực tuyệt đối, Kiểm tra thuốc có trong kho, Trích dẫn nguồ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 _build_references_text(self, 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 ""
            ref_str = f"- {src} {row_info}"
            if ref_str not in seen:
                lines.append(ref_str)
                seen.add(ref_str)
        return "\n".join(lines)

# === PHẦN GIAO DIỆN CHAINLIT ===

@cl.on_chat_start
async def start():
    # Khởi tạo bot khi phiên chat bắt đầu
    msg = cl.Message(content="Đang khởi động hệ thống DeepMed AI...")
    await msg.send()
    
    bot = DeepMedBot()
    cl.user_session.set("bot", bot) # Lưu bot vào session của người dùng
    
    # Gửi tin nhắn chào mừng và các nút chọn chế độ
    actions = [
        cl.Action(name="fast_mode", value="fast", label="⚡ Tốc độ (Nhanh)", description="Trả lời nhanh"),
        cl.Action(name="deep_mode", value="deep", label="🧠 Chuyên sâu (Kỹ)", description="Phân tích kỹ dữ liệu")
    ]
    cl.user_session.set("mode", "fast") # Mặc định là fast
    
    await msg.update(content="👋 Xin chào! Tôi là **DeepMed AI**. \n\nTôi có thể hỗ trợ tra cứu thuốc, phác đồ điều trị và thông tin dược lâm sàng. \n\n*Mặc định đang ở chế độ: Tốc độ.*", actions=actions)

@cl.action_callback("fast_mode")
async def on_fast_mode(action):
    cl.user_session.set("mode", "fast")
    await cl.Message(content="✅ Đã chuyển sang chế độ: **Tốc độ (Nhanh)**").send()

@cl.action_callback("deep_mode")
async def on_deep_mode(action):
    cl.user_session.set("mode", "deep")
    await cl.Message(content="✅ Đã chuyển sang chế độ: **Chuyên sâu (Chậm & Chính xác)**").send()

@cl.on_message
async def main(message: cl.Message):
    bot = cl.user_session.get("bot")
    mode_setting = cl.user_session.get("mode")
    
    if not bot or not bot.ready:
        await cl.Message(content="⚠️ Hệ thống chưa sẵn sàng hoặc gặp lỗi khởi tạo.").send()
        return

    # Lấy lịch sử chat từ Chainlit
    # Chainlit tự quản lý history, nhưng để tương thích code cũ, ta có thể lấy memory
    # Ở đây ta dùng bộ nhớ session đơn giản hoặc để LangChain tự lo
    
    # Xác định chain cần dùng
    active_chain = bot.deep_chain if mode_setting == "deep" else bot.fast_chain
    
    msg = cl.Message(content="")
    await msg.send()
    
    # Lấy lịch sử chat (đơn giản hóa cho demo)
    history = [] 
    
    full_response = ""
    retrieved_docs = []

    # Gọi Chain Async
    try:
        async for chunk in active_chain.astream({"input": message.content, "chat_history": history}):
            if "answer" in chunk:
                token = chunk["answer"]
                full_response += token
                await msg.stream_token(token)
            elif "context" in chunk:
                retrieved_docs = chunk["context"]
        
        # Hiển thị nguồn tham khảo đẹp mắt hơn
        if retrieved_docs:
            refs = bot._build_references_text(retrieved_docs)
            if refs:
                ref_block = f"\n\n---\n📚 **Nguồn tham khảo:**\n{refs}"
                await msg.stream_token(ref_block)
                
        await msg.update()
        
    except Exception as e:
        await cl.Message(content=f"⚠️ Lỗi: {str(e)}").send()