| 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 |
|
|
| |
| _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) |
|
|
