File size: 2,886 Bytes
b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 8c4590b b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c 74b76f3 b59fc2c | 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 68 69 70 71 72 73 74 75 76 | from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception
from langchain_core.prompts import PromptTemplate
from config import Config
from key_manager import GroqKeyManager
from hybrid_retriever import HybridRetriever
from vector_store import VectorStoreManager
# Shared key manager -- single instance reused across all RAGChain objects
_KEY_MANAGER = None
def get_key_manager():
global _KEY_MANAGER
if _KEY_MANAGER is None:
_KEY_MANAGER = GroqKeyManager(
keys=[Config.GROQ_API_KEY_1, Config.GROQ_API_KEY_2, Config.GROQ_API_KEY_3],
model=Config.GROQ_MODEL,
)
return _KEY_MANAGER
def _is_rate_limit(exc):
msg = str(exc).lower()
return "429" in msg or "quota" in msg or "rate limit" in msg or "ratelimit" in msg
class RAGChain:
def __init__(self, vector_store_manager):
self._km = get_key_manager()
self.vectorstore = vector_store_manager.vector_store
self.retriever = HybridRetriever(self.vectorstore)
self.prompt_template = PromptTemplate(
input_variables=["context", "question"],
template="Tài liệu y khoa:\n{context}\n\nCâu hỏi: {question}\n\nTrả lời ngắn gọn, chọn lọc thông tin quan trọng nhất từ tài liệu (tối đa 200 từ):"
)
def query(self, question):
sources = self.retriever.hybrid_search(question, k=3)
ranked = self.rerank_sources(sources, question)
context = self.build_context(ranked)
prompt = self.prompt_template.format(context=context, question=question)
@retry(
retry=retry_if_exception(_is_rate_limit),
wait=wait_exponential(multiplier=1, min=5, max=30),
stop=stop_after_attempt(4),
reraise=True,
)
def _invoke():
try:
llm = self._km.build_llm(temperature=0)
return llm.invoke([prompt])
except Exception as exc:
if _is_rate_limit(exc):
self._km.mark_rate_limited(self._km.current())
self._km.rotate()
raise
result = _invoke()
return result.content, ranked
def rerank_sources(self, sources, question):
keywords = question.lower().split()
def score(doc):
text = doc.page_content.lower() + doc.metadata.get("chunk_title", "").lower()
return sum(1 for kw in keywords if kw in text)
return sorted(sources, key=score, reverse=True)
def build_context(self, sources):
parts = []
for i, doc in enumerate(sources[:3]):
meta = f"[{i+1}] {doc.metadata.get('source_file','?')} | {doc.metadata.get('chunk_title','?')}"
content = doc.page_content[:600]
parts.append(f"{meta}\n{content}")
return "\n\n".join(parts)
|