Spaces:
Sleeping
Sleeping
| """Handles end-to-end query: preprocess → embed → retrieve → rerank → build context → LLM → format response.""" | |
| from __future__ import annotations | |
| from rag_engine.prompts.response_format import ResponseFormatter | |
| from rag_engine.utils.logger import get_logger | |
| logger = get_logger(__name__) | |
| class QueryService: | |
| """End-to-end query pipeline: retrieve → rerank → context → LLM → format.""" | |
| def __init__(self) -> None: | |
| from rag_engine.embeddings.embedding_factory import get_embedder | |
| from rag_engine.llm.llm_factory import get_llm | |
| from rag_engine.retrieval.context_builder import ContextBuilder | |
| from rag_engine.retrieval.reranker import CrossEncoderReranker | |
| from rag_engine.retrieval.retriever import PolicyRetriever | |
| from rag_engine.vector_store.store_factory import get_vector_store | |
| self._embedder = get_embedder() | |
| self._store = get_vector_store() | |
| self._retriever = PolicyRetriever(self._store, self._embedder) | |
| self._reranker = CrossEncoderReranker() | |
| self._context_builder = ContextBuilder() | |
| self._llm = get_llm() | |
| self._formatter = ResponseFormatter() | |
| logger.info("QueryService initialized — all components ready") | |
| # ------------------------------------------------------------------ # | |
| # single-policy query | |
| # ------------------------------------------------------------------ # | |
| def query( | |
| self, | |
| question: str, | |
| policy_id: str, | |
| k: int = 5, # reduced from 8 — fetch 5 candidates | |
| top_k_rerank: int = 3, # reduced from 5 — keep top 3 after rerank | |
| ) -> dict: | |
| """Answer *question* using chunks from *policy_id*.""" | |
| logger.info( | |
| "QueryService.query START | policy_id=%s | q='%s'", | |
| policy_id, | |
| question[:60], | |
| ) | |
| # Step 1 — retrieve | |
| raw_results = self._retriever.retrieve(question, policy_id, k=k) | |
| # Step 2 — rerank | |
| reranked = self._reranker.rerank(question, raw_results, top_k=top_k_rerank) | |
| # Step 3 — build context | |
| context = self._context_builder.build(reranked) | |
| # Step 4 — build prompt | |
| from rag_engine.prompts.context_template import build_query_prompt | |
| from rag_engine.prompts.system_prompt import SYSTEM_PROMPT | |
| prompt = build_query_prompt(question, context, policy_id) | |
| # Step 5 — LLM | |
| raw_answer = self._llm.complete(prompt, system=SYSTEM_PROMPT) | |
| # Step 6 — format | |
| result = self._formatter.format_answer(raw_answer, question, reranked) | |
| logger.info( | |
| "QueryService.query COMPLETE | sources=%d", | |
| result["source_count"], | |
| ) | |
| return result | |
| def stream_query( | |
| self, | |
| question: str, | |
| policy_id: str, | |
| k: int = 5, # reduced from 8 | |
| top_k_rerank: int = 3, # reduced from 5 | |
| ): | |
| """Answer *question* via stream using chunks from *policy_id*.""" | |
| logger.info( | |
| "QueryService.stream_query START | policy_id=%s | q='%s'", | |
| policy_id, | |
| question[:60], | |
| ) | |
| # Step 1 — retrieve | |
| raw_results = self._retriever.retrieve(question, policy_id, k=k) | |
| # Step 2 — rerank | |
| reranked = self._reranker.rerank(question, raw_results, top_k=top_k_rerank) | |
| # Step 3 — build context | |
| context = self._context_builder.build(reranked) | |
| # Step 4 — build prompt | |
| from rag_engine.prompts.context_template import build_query_prompt | |
| from rag_engine.prompts.system_prompt import SYSTEM_PROMPT | |
| prompt = build_query_prompt(question, context, policy_id) | |
| # Step 5 — LLM Stream | |
| for token in self._llm.stream(prompt, system=SYSTEM_PROMPT): | |
| yield token | |
| # ------------------------------------------------------------------ # | |
| # multi-policy query | |
| # ------------------------------------------------------------------ # | |
| def query_multi_policy( | |
| self, question: str, policy_ids: list[str] | |
| ) -> dict: | |
| """Answer *question* across multiple policies.""" | |
| raw_results = self._retriever.retrieve_multi_policy( | |
| question, policy_ids, k_per_policy=3 # reduced from 4 | |
| ) | |
| reranked = self._reranker.rerank(question, raw_results, top_k=3) # reduced from 5 | |
| context = self._context_builder.build(reranked) | |
| from rag_engine.prompts.context_template import build_query_prompt | |
| from rag_engine.prompts.system_prompt import SYSTEM_PROMPT | |
| prompt = build_query_prompt( | |
| question, context, f"Multi-policy: {', '.join(policy_ids)}" | |
| ) | |
| raw_answer = self._llm.complete(prompt, system=SYSTEM_PROMPT) | |
| return self._formatter.format_answer(raw_answer, question, reranked) | |