179 lines
6.7 KiB
Python
179 lines
6.7 KiB
Python
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,
|
|
)
|