50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import logging
|
|
from typing import List, Dict
|
|
from search.retriever import SearchRetriever
|
|
from .llm_factory import LLMFactory
|
|
|
|
logger = logging.getLogger("RAGEngine")
|
|
|
|
class RAGEngine:
|
|
def __init__(self):
|
|
self.retriever = SearchRetriever()
|
|
self.llm = LLMFactory.get_provider()
|
|
logger.info(f"RAG Engine đã sẵn sàng với LLM Provider: {type(self.llm).__name__}")
|
|
|
|
def chat(self, user_query: str, history: List[Dict[str, str]] = None) -> Dict:
|
|
"""
|
|
Quy trình RAG hoàn chỉnh: Search -> Augment -> Generate
|
|
"""
|
|
# 1. RETRIEVAL: Tìm kiếm ngữ cảnh liên quan
|
|
relevant_chunks = self.retriever.retrieve(user_query, top_k=5)
|
|
|
|
if not relevant_chunks:
|
|
context_text = "Không tìm thấy thông tin liên quan trong cơ sở dữ liệu nội bộ."
|
|
else:
|
|
# Gộp text từ các chunks lại thành 1 khối context
|
|
context_text = "\n---\n".join([
|
|
f"[Nguồn: {c.file_name}, Trang: {c.page_from}]\nNội dung: {c.text}"
|
|
for c in relevant_chunks
|
|
])
|
|
|
|
# 2. GENERATION: Gửi sang LLM để trả lời
|
|
logger.info("Đang yêu cầu LLM tổng hợp câu trả lời...")
|
|
answer = self.llm.generate_response(
|
|
prompt=user_query,
|
|
context=context_text,
|
|
history=history
|
|
)
|
|
|
|
# 3. Trả về kết quả kèm theo nguồn trích dẫn (Citations)
|
|
return {
|
|
"answer": answer,
|
|
"context_used": context_text,
|
|
"sources": [
|
|
{
|
|
"file_name": c.file_name,
|
|
"page": c.page_from,
|
|
"url": c.source_url
|
|
} for c in relevant_chunks
|
|
]
|
|
}
|