| from __future__ import annotations |
|
|
| import time |
| from collections.abc import AsyncIterator |
| from typing import Any |
|
|
| from secrag.analytics.planner import execute_plan, plan_numeric |
| from secrag.cache import SemanticCache, partition_key |
| from secrag.core.config import Settings, get_settings |
| from secrag.core.logging import get_logger |
| from secrag.core.types import ( |
| Answer, |
| AnswerStatus, |
| NumericResult, |
| QueryIntent, |
| QueryRequest, |
| QueryResponse, |
| RouteDecision, |
| ScoredChunk, |
| ) |
| from secrag.generation.synthesize import Generator |
| from secrag.ingest.xbrl import FactStore |
| from secrag.observability.tracing import Trace, span, start_trace |
| from secrag.retrieval.embedder import Embedder |
| from secrag.retrieval.pipeline import HybridRetriever |
| from secrag.retrieval.store import SearchFilter |
| from secrag.routing.router import QueryRouter |
|
|
| log = get_logger(__name__) |
|
|
|
|
| class QueryEngine: |
| def __init__(self, settings: Settings | None = None) -> None: |
| self.settings = settings or get_settings() |
| self.settings.ensure_dirs() |
|
|
| self.embedder = Embedder(self.settings) |
| self.retriever = HybridRetriever(self.settings, embedder=self.embedder) |
| self.generator = Generator(self.settings, embedder=self.embedder) |
| self.router = QueryRouter(self.settings, embedder=self.embedder) |
| self.cache = SemanticCache(self.settings, embedder=self.embedder) |
| self.facts = FactStore.load(self.settings) |
|
|
| def warmup(self) -> None: |
| self.retriever.warmup() |
| log.info( |
| "engine_ready", |
| corpus=self.retriever.corpus_size, |
| facts=len(self.facts.df), |
| router=self.router.is_available, |
| tickers=self.facts.tickers(), |
| ) |
|
|
| async def aclose(self) -> None: |
| closer = getattr(self.generator.provider, "aclose", None) |
| if closer is not None: |
| await closer() |
| self.retriever.store.close() |
|
|
| def stats(self) -> dict[str, Any]: |
| return { |
| "corpus_chunks": self.retriever.corpus_size, |
| "bm25_documents": self.retriever.bm25.size, |
| "xbrl_rows": len(self.facts.df), |
| "tickers": self.facts.tickers(), |
| "router_available": self.router.is_available, |
| "providers": getattr(self.generator.provider, "describe", lambda: [])(), |
| "cache": self.cache.snapshot(), |
| "settings": { |
| "dense_model": self.settings.dense_model, |
| "sparse_model": self.settings.sparse_model, |
| "rerank_model": self.settings.rerank_model, |
| "splade_enabled": self.settings.enable_splade, |
| "min_groundedness": self.settings.min_groundedness, |
| }, |
| } |
|
|
| def _filter(self, request: QueryRequest) -> SearchFilter: |
| return SearchFilter( |
| tickers=request.companies, |
| fiscal_years=request.fiscal_years, |
| sections=request.sections, |
| ) |
|
|
| def _partition(self, request: QueryRequest) -> str: |
| return partition_key( |
| request.companies, request.fiscal_years, [s.value for s in request.sections] |
| ) |
|
|
| def _numeric(self, question: str, route: RouteDecision) -> list[NumericResult]: |
| if route.intent not in (QueryIntent.NUMERIC, QueryIntent.COMPARATIVE): |
| return [] |
| if self.facts.is_empty: |
| return [] |
| with span("numeric_plan", intent=route.intent.value): |
| plan = plan_numeric(question, self.facts) |
| if plan is None: |
| log.info("numeric_plan_unresolved", question=question[:80]) |
| return [] |
| return execute_plan(plan, self.facts) |
|
|
| def _retrieve( |
| self, request: QueryRequest, route: RouteDecision |
| ) -> tuple[list[ScoredChunk], str]: |
| top_k = request.top_k |
| if route.intent is QueryIntent.COMPARATIVE: |
| top_k = min(top_k * 2, 20) |
|
|
| reranker = request.reranker if request.use_reranker else "none" |
| result = self.retriever.retrieve( |
| request.question, top_k=top_k, flt=self._filter(request), reranker=reranker |
| ) |
| return result.chunks, result.reranker |
|
|
| async def answer(self, request: QueryRequest) -> QueryResponse: |
| started = time.perf_counter() |
| partition = self._partition(request) |
|
|
| with start_trace(question=request.question[:120]) as trace: |
| if request.use_cache and (hit := self.cache.get(request.question, partition)): |
| hit.trace_id = trace.trace_id |
| hit.latency_ms = round((time.perf_counter() - started) * 1000, 2) |
| return hit |
|
|
| route = self.router.route(request.question) |
| numeric = self._numeric(request.question, route) |
| contexts, reranker = self._retrieve(request, route) |
|
|
| if not contexts and not numeric: |
| return self._empty_response(request, route, trace, started) |
|
|
| generation = await self.generator.generate(request.question, contexts, numeric) |
|
|
| response = QueryResponse( |
| question=request.question, |
| answer=generation.answer, |
| route=route, |
| contexts=generation.contexts, |
| numeric_results=numeric, |
| trace_id=trace.trace_id, |
| cached=False, |
| latency_ms=round((time.perf_counter() - started) * 1000, 2), |
| ) |
| trace.metadata.update( |
| { |
| "intent": route.intent.value, |
| "reranker": reranker, |
| "contexts": len(generation.contexts), |
| "numeric_results": len(numeric), |
| "groundedness": generation.answer.groundedness, |
| "status": generation.answer.status.value, |
| } |
| ) |
|
|
| if request.use_cache and response.answer.status is AnswerStatus.OK: |
| self.cache.put(request.question, response, partition) |
|
|
| log.info( |
| "query_answered", |
| intent=route.intent.value, |
| status=response.answer.status.value, |
| groundedness=response.answer.groundedness, |
| latency_ms=response.latency_ms, |
| ) |
| return response |
|
|
| def _empty_response( |
| self, request: QueryRequest, route: RouteDecision, trace: Trace, started: float |
| ) -> QueryResponse: |
| from secrag.generation.prompts import REFUSAL_NO_CONTEXT |
|
|
| return QueryResponse( |
| question=request.question, |
| answer=Answer( |
| text=REFUSAL_NO_CONTEXT, |
| status=AnswerStatus.REFUSED_NO_CONTEXT, |
| refusal_reason="Retrieval returned no passages and no figure could be resolved.", |
| ), |
| route=route, |
| trace_id=trace.trace_id, |
| latency_ms=round((time.perf_counter() - started) * 1000, 2), |
| ) |
|
|
| async def stream(self, request: QueryRequest) -> AsyncIterator[dict[str, Any]]: |
| started = time.perf_counter() |
| partition = self._partition(request) |
|
|
| with start_trace(question=request.question[:120]) as trace: |
| if request.use_cache and (hit := self.cache.get(request.question, partition)): |
| yield { |
| "event": "meta", |
| "data": { |
| "trace_id": trace.trace_id, |
| "route": hit.route.model_dump(mode="json") if hit.route else None, |
| "reranker": "cached", |
| "contexts": [_context_payload(c) for c in hit.contexts], |
| "numeric_results": [n.model_dump(mode="json") for n in hit.numeric_results], |
| "cached": True, |
| }, |
| } |
| yield {"event": "token", "data": {"text": hit.answer.text}} |
| yield { |
| "event": "done", |
| "data": { |
| "status": hit.answer.status.value, |
| "groundedness": hit.answer.groundedness, |
| "refusal_reason": hit.answer.refusal_reason, |
| "citations": [c.model_dump(mode="json") for c in hit.answer.citations], |
| "cached": True, |
| "latency_ms": round((time.perf_counter() - started) * 1000, 2), |
| "trace": trace.to_dict(), |
| }, |
| } |
| return |
|
|
| route = self.router.route(request.question) |
| numeric = self._numeric(request.question, route) |
| contexts, reranker = self._retrieve(request, route) |
|
|
| yield { |
| "event": "meta", |
| "data": { |
| "trace_id": trace.trace_id, |
| "route": route.model_dump(mode="json"), |
| "reranker": reranker, |
| "contexts": [_context_payload(c) for c in contexts], |
| "numeric_results": [n.model_dump(mode="json") for n in numeric], |
| }, |
| } |
|
|
| if not contexts and not numeric: |
| from secrag.generation.prompts import REFUSAL_NO_CONTEXT |
|
|
| yield {"event": "token", "data": {"text": REFUSAL_NO_CONTEXT}} |
| yield { |
| "event": "done", |
| "data": { |
| "status": AnswerStatus.REFUSED_NO_CONTEXT.value, |
| "groundedness": 0.0, |
| "citations": [], |
| "cached": False, |
| "latency_ms": round((time.perf_counter() - started) * 1000, 2), |
| "trace": trace.to_dict(), |
| }, |
| } |
| return |
|
|
| collected: list[str] = [] |
| async for piece in self.generator.stream(request.question, contexts, numeric): |
| collected.append(piece) |
| yield {"event": "token", "data": {"text": piece}} |
|
|
| final = self.generator.finalise_streamed( |
| "".join(collected), contexts, has_verified_figures=bool(numeric) |
| ) |
| latency_ms = round((time.perf_counter() - started) * 1000, 2) |
|
|
| if request.use_cache and final.answer.status is AnswerStatus.OK: |
| self.cache.put( |
| request.question, |
| QueryResponse( |
| question=request.question, |
| answer=final.answer, |
| route=route, |
| contexts=final.contexts, |
| numeric_results=list(numeric), |
| trace_id=trace.trace_id, |
| latency_ms=latency_ms, |
| ), |
| partition, |
| ) |
|
|
| yield { |
| "event": "done", |
| "data": { |
| "status": final.answer.status.value, |
| "groundedness": final.answer.groundedness, |
| "refusal_reason": final.answer.refusal_reason, |
| "citations": [c.model_dump(mode="json") for c in final.answer.citations], |
| "cached": False, |
| "latency_ms": latency_ms, |
| "trace": trace.to_dict(), |
| }, |
| } |
|
|
|
|
| def _context_payload(scored: ScoredChunk) -> dict[str, Any]: |
| return { |
| "chunk_id": scored.chunk.chunk_id, |
| "label": scored.chunk.citation_label(), |
| "text": scored.chunk.text, |
| "kind": scored.chunk.kind.value, |
| "section": scored.chunk.section.value, |
| "company": scored.chunk.company, |
| "ticker": scored.chunk.ticker, |
| "fiscal_year": scored.chunk.fiscal_year, |
| "source_url": scored.chunk.source_url, |
| "score": round(scored.score, 6), |
| "stage": scored.stage, |
| "rank": scored.rank, |
| "component_scores": {k: round(v, 6) for k, v in scored.component_scores.items()}, |
| } |
|
|
|
|
| def build_engine(settings: Settings | None = None) -> QueryEngine: |
| return QueryEngine(settings) |
|
|
