File size: 2,151 Bytes
74b76f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | import sys
import os
sys.path.append(r'D:\Storage\rag_project\src')
# os.environ["GOOGLE_API_KEY"] = "AIzaSyABvC8mPrwa0Kgy08mFFzkyeh2_N-Bb3lY" # Thay key thật
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from test_single_file_loader import test_single_file
def test_rag_single(filename):
print(f"\n FULL RAG TEST: {filename}")
docs = test_single_file(filename)
# Build FAISS
print(" Building FAISS...")
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
# LLM + Prompt
print(" Init Gemini...")
llm = ChatGoogleGenerativeAI(
model="gemini-2.5-flash",
temperature=0,
google_api_key="AIzaSyBKUfFRLphY4AgTY-j5sr-6s0SFWW0ATyg" # API KEY Ở ĐÂY
)
custom_prompt = PromptTemplate(
input_variables=["context", "question"],
template="""Bạn là bác sĩ nhi khoa. Dựa vào TÀI LIỆU Y KHOA sau:
CONTEXT: {context}
CÂU HỎI: {question}
TRẢ LỜI chính xác dựa trên CONTEXT, ngắn gọn, chuyên nghiệp."""
)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=retriever,
chain_type_kwargs={"prompt": custom_prompt},
return_source_documents=True
)
# CHỈ NHIKHOA2.json - 1 query đúng
query = "điều trị suy hô hấp"
print(f"\n Query: {query}")
result = qa_chain.invoke({"query": query})
print(f" Answer: {result['result'][:400]}...")
# FIX: Kiểm tra key tồn tại
if 'source_documents' in result:
print(f" Sources: {len(result['source_documents'])} docs")
else:
print(" Sources: Không có source_documents (Gemini 2.5 format)")
print("\n RAG SINGLE FILE OK!")
if __name__ == "__main__":
test_rag_single("NHIKHOA2.json") |