from __future__ import annotations import uuid from typing import Annotated, Any, Protocol from fastapi import APIRouter, Depends, HTTPException, Request from pydantic import BaseModel, Field from rag.answer import GroundedAnswerService from rag.metrics import TRACE_WRITE_FAILED, Metrics, NullMetrics from rag.models import QueryIntent, SubjectScope from rag.policy import resolve_subject_scope class TraceWriter(Protocol): def save(self, **fields: Any) -> str: ... class RagQueryRequest(BaseModel): query: str = Field(min_length=1, max_length=4000) subject_scope: SubjectScope intent: QueryIntent # Optional: when present, the turn is answered in conversation context # (follow-up inheritance, clarify, smalltalk). Absent → single-turn, exactly # as before, so existing callers are unchanged. conversation_id: str | None = Field(default=None, max_length=128) class CitationResponse(BaseModel): chunk_id: str printed_page_start: int printed_page_end: int physical_page: int block_id: str | None = None bbox: tuple[float, float, float, float] | None = None source_crop: str | None = None attachment: str | None = None class RagQueryResponse(BaseModel): trace_id: str decision: str reason: str answer: str | None resolved_drug_id: str | None citations: list[CitationResponse] def _answer_service(request: Request) -> GroundedAnswerService: service = getattr(request.app.state, "answer_service", None) if service is None: raise HTTPException(status_code=503, detail="RAG backend is not configured") return service def _trace_writer(request: Request) -> TraceWriter: writer = getattr(request.app.state, "trace_writer", None) if writer is None: raise HTTPException(status_code=503, detail="Trace database is not configured") return writer def _metrics(request: Request) -> Metrics: return getattr(request.app.state, "metrics", None) or NullMetrics() router = APIRouter(prefix="/v1/rag", tags=["rag"]) class SuggestResponse(BaseModel): suggestions: list[str] @router.get("/suggest", response_model=SuggestResponse) def suggest_drugs(q: str, request: Request) -> SuggestResponse: """As-you-type drug-name autocomplete, so a name is picked, not mistyped.""" agent = getattr(request.app.state, "conversational", None) if agent is None or not q.strip(): return SuggestResponse(suggestions=[]) return SuggestResponse(suggestions=agent.complete(q.strip())) def _map_citations(items) -> list[CitationResponse]: return [ CitationResponse( chunk_id=item.chunk_id, printed_page_start=item.printed_page_start, printed_page_end=item.printed_page_end, physical_page=item.physical_page, block_id=item.block_id, bbox=item.bbox, source_crop=item.source_crop, attachment=item.attachment, ) for item in items ] @router.post("/query", response_model=RagQueryResponse) def query_rag( payload: RagQueryRequest, request: Request, answers: Annotated[GroundedAnswerService, Depends(_answer_service)], traces: Annotated[TraceWriter, Depends(_trace_writer)], metrics: Annotated[Metrics, Depends(_metrics)], ) -> RagQueryResponse: agent = getattr(request.app.state, "conversational", None) # `payload.subject_scope`/`payload.intent` are what the CALLER claims — # logged below for audit, but the RagAgent path does not take them as an # input at all. It derives scope from the query text itself (the same # `rag.policy` check used here for the no-agent fallback) as part of its # own LLM understanding call, and does not gate on intent at all (this # product is for doctors/pharmacists; a client label must not be, and # here structurally cannot be, the safety decision — F-02). subject_scope = resolve_subject_scope(payload.query, payload.subject_scope) intent = payload.intent if agent is not None: # The live path (F-03): one LLM call understands the turn (drug # identity against the real catalog, turn type, population/weight), # then routes to the safety-verified retrieval + grounded-answer # engine. Replaces the old fuzzy resolver + keyword section router + # manual follow-up inheritance for both single- and multi-turn. reply = agent.handle(payload.query, payload.conversation_id) decision = reply.decision reason = reply.reason answer = reply.clarification if reply.clarification is not None else reply.answer resolved_drug_id = ", ".join(reply.drugs) if reply.drugs else None citations = _map_citations(reply.citations) else: # No generator configured (ANSWER_PROVIDER=disabled): there is no LLM # to understand a turn with, so this is retrieval-only, single-turn, # unchanged from before F-03. grounded = answers.answer(payload.query, subject_scope, intent) if grounded.clarification is not None: decision, reason = "clarify", "needs_more_info" answer = grounded.clarification resolved_drug_id = grounded.result.resolved_drug_id citations = [] else: decision = grounded.result.decision.value reason = grounded.result.reason answer = grounded.answer resolved_drug_id = grounded.result.resolved_drug_id citations = _map_citations(grounded.citations) # Trace persistence is fail-open (F-09): an already-computed, safe answer # must reach the caller even if Postgres is unreachable. `save()` opens a # fresh connection per call with no retry, so a DB outage previously # turned a good answer into a 500 for a reason that has nothing to do # with whether the answer was safe. `trace_id` degrades to a local, # unpersisted uuid — still a valid response field, just not one `GET # /v1/rag/trace/{id}` (if it existed) could later look up. try: trace_id = traces.save( query=payload.query, # The resolved (server-derived) values, not the caller's claim — # this is what actually gated the answer, so it's what the trace # must show. subject_scope=subject_scope.value, intent=intent.value, decision=decision, reason=reason, resolved_drug_id=resolved_drug_id, citations=tuple(item.model_dump() for item in citations), ) except Exception: metrics.increment(TRACE_WRITE_FAILED) trace_id = str(uuid.uuid4()) return RagQueryResponse( trace_id=trace_id, decision=decision, reason=reason, answer=answer, resolved_drug_id=resolved_drug_id, citations=citations, )