Create app.py

app.py

@@ -0,0 +1,244 @@
+__import__('pysqlite3')
+import sys
+sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
+
+import os
+import gradio as gr
+
+# --- IMPORT CÁC THƯ VIỆN ---
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_chroma import Chroma
+# Loaders cho nhiều định dạng
+from langchain_community.document_loaders import (
+    PyPDFLoader, 
+    DirectoryLoader, 
+    TextLoader, 
+    Docx2txtLoader, 
+    UnstructuredExcelLoader
+)
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers 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
+
+# ==========================================
+# CẤU HÌNH
+# ==========================================
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+DATA_PATH = "medical_data"
+DB_PATH = "chroma_db"
+
+if not GOOGLE_API_KEY:
+    print("CẢNH BÁO: Chưa thiết lập GOOGLE_API_KEY!")
+
+# ==========================================
+# XỬ LÝ DỮ LIỆU ĐA ĐỊNH DẠNG
+# ==========================================
+def load_documents_from_folder(folder_path):
+    """Hàm đọc nhiều loại file khác nhau từ thư mục"""
+    documents = []
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+        return []
+
+    for filename in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, filename)
+        loader = None
+        
+        try:
+            if filename.endswith(".pdf"):
+                loader = PyPDFLoader(file_path)
+            elif filename.endswith(".docx") or filename.endswith(".doc"):
+                loader = Docx2txtLoader(file_path)
+            elif filename.endswith(".txt"):
+                loader = TextLoader(file_path, encoding="utf-8")
+            elif filename.endswith(".xlsx") or filename.endswith(".xls"):
+                loader = UnstructuredExcelLoader(file_path)
+            
+            if loader:
+                print(f"-> Đang đọc file: {filename}")
+                docs = loader.load()
+                # Thêm tên file vào metadata nếu chưa có (để trích dẫn sau này)
+                for doc in docs:
+                    if "source" not in doc.metadata:
+                        doc.metadata["source"] = filename
+                documents.extend(docs)
+        except Exception as e:
+            print(f"Lỗi khi đọc file {filename}: {e}")
+            
+    return documents
+
+def get_retriever():
+    print("--- Đang tải model Embedding... ---")
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+
+    splits = []
+    vectorstore = None
+    
+    # 1. Kiểm tra DB cũ
+    if os.path.exists(DB_PATH) and os.listdir(DB_PATH):
+        try:
+            print("--- Tìm thấy Database cũ, đang tải lên... ---")
+            vectorstore = Chroma(persist_directory=DB_PATH, embedding_function=embedding_model)
+            
+            # Lấy dữ liệu để tái tạo BM25 (Keyword Search)
+            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))
+            else:
+                print("Cảnh báo: Database rỗng.")
+        except Exception as e:
+            print(f"Lỗi đọc DB cũ: {e}")
+
+    # 2. Nếu chưa có dữ liệu (splits rỗng), đọc từ file gốc
+    if not splits:
+        print("--- Chưa có dữ liệu index, bắt đầu đọc file nguồn... ---")
+        documents = load_documents_from_folder(DATA_PATH)
+        
+        if not documents:
+            print("Lỗi: Không tìm thấy tài liệu nào (PDF, DOCX, TXT...) trong thư mục medical_data.")
+            return None
+
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.split_documents(documents)
+
+        print(f"--- Đang mã hóa {len(splits)} đoạn văn bản vào ChromaDB... ---")
+        vectorstore = Chroma.from_documents(
+            documents=splits, 
+            embedding=embedding_model, 
+            persist_directory=DB_PATH
+        )
+
+    # 3. Tạo Hybrid Search (BM25 + Vector)
+    if not splits: return None
+
+    bm25_retriever = BM25Retriever.from_documents(splits)
+    bm25_retriever.k = 5
+    chroma_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
+
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[bm25_retriever, chroma_retriever],
+        weights=[0.4, 0.6]
+    )
+    return ensemble_retriever
+
+# ==========================================
+# LOGIC CHATBOT THÔNG MINH (CÓ NHỚ)
+# ==========================================
+rag_chain = None
+
+def init_chatbot():
+    global rag_chain
+    retriever = get_retriever()
+    if not retriever: return False
+
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.5-flash", 
+        temperature=0.3,
+        google_api_key=GOOGLE_API_KEY
+    )
+
+    # --- BƯỚC 1: Contextualize Question ---
+    # (Viết lại câu hỏi mới dựa trên lịch sử để AI hiểu context)
+    contextualize_q_system_prompt = (
+        "Dựa trên lịch sử trò chuyện và câu hỏi mới nhất của người dùng, "
+        "nếu câu hỏi liên quan đến ngữ cảnh trước đó, hãy viết lại nó thành một câu hỏi độc lập đầy đủ ý nghĩa. "
+        "Nếu không liên quan, giữ nguyên câu hỏi gốc. KHÔNG trả lời câu hỏi, chỉ viết lại thôi."
+    )
+    contextualize_q_prompt = ChatPromptTemplate.from_messages([
+        ("system", contextualize_q_system_prompt),
+        MessagesPlaceholder("chat_history"),
+        ("human", "{input}"),
+    ])
+    # Retriever biết nhớ lịch sử
+    history_aware_retriever = create_history_aware_retriever(
+        llm, retriever, contextualize_q_prompt
+    )
+
+    # --- BƯỚC 2: Answer Question ---
+    # (Trả lời dựa trên Documents tìm được)
+    qa_system_prompt = (
+        "Bạn là trợ lý y tế DeepMed. Sử dụng các đoạn văn bản được cung cấp (Context) để trả lời câu hỏi. "
+        "Nếu không biết, hãy nói không biết. Nếu tìm thấy nội dung trả lời hãy trích dẫn tài liệu. Giữ câu trả lời ngắn gọn, súc tích.\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(llm, qa_prompt)
+    
+    # Kết hợp lại thành chuỗi RAG hoàn chỉnh
+    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
+    
+    print("--- Chatbot đã sẵn sàng! ---")
+    return True
+
+init_success = init_chatbot()
+
+# ==========================================
+# HÀM CHAT (Xử lý Lịch sử & Trích dẫn)
+# ==========================================
+def chat(message, history):
+    if not init_success:
+        return "Hệ thống chưa sẵn sàng. Kiểm tra lại data và API Key."
+    
+    # 1. Chuyển đổi lịch sử Gradio sang format LangChain
+    chat_history = []
+    for user_msg, bot_msg in history:
+        chat_history.append(HumanMessage(content=user_msg))
+        chat_history.append(AIMessage(content=bot_msg))
+    
+    try:
+        # 2. Gọi Chain xử lý
+        response = rag_chain.invoke({
+            "input": message,
+            "chat_history": chat_history
+        })
+        
+        answer = response["answer"]
+        
+        # 3. Xử lý Trích dẫn nguồn (References)
+        sources = set()
+        if "context" in response:
+            for doc in response["context"]:
+                source_name = doc.metadata.get("source", "Tài liệu không tên")
+                page_num = doc.metadata.get("page", None)
+                
+                # Format tên file cho gọn (bỏ đường dẫn)
+                source_name = os.path.basename(source_name)
+                
+                if page_num is not None:
+                    sources.add(f"{source_name} (Trang {page_num + 1})")
+                else:
+                    sources.add(source_name)
+        
+        if sources:
+            answer += "\n\n---\n📚 **Tài liệu tham khảo:**\n" + "\n".join([f"- {s}" for s in sources])
+            
+        return answer
+        
+    except Exception as e:
+        return f"Lỗi hệ thống: {str(e)}"
+
+# ==========================================
+# GIAO DIỆN
+# ==========================================
+demo = gr.ChatInterface(
+    fn=chat,
+    title="🏥 DeepMed AI Pro - Trợ lý Y khoa Đa tài liệu",
+    description="Hỗ trợ PDF, DOCX, TXT, XLSX. Có khả năng nhớ ngữ cảnh hội thoại.",
+    theme="soft",
+    examples=["Bệnh cúm mùa là gì?", "Triệu chứng ra sao?", "Cách điều trị?"], 
+)
+
+if __name__ == "__main__":
+    demo.launch()