Wire the guarded conversational RAG answer layer end-to-end

apps/ai-service/routers/rag.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+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.models import QueryIntent, SubjectScope
+
+
+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
+
+
+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."""
+    conversational = getattr(request.app.state, "conversational", None)
+    if conversational is None or not q.strip():
+        return SuggestResponse(suggestions=[])
+    return SuggestResponse(suggestions=conversational.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)],
+) -> RagQueryResponse:
+    conversational = getattr(request.app.state, "conversational", None)
+
+    # Single-turn path (no conversation id, or conversational layer disabled):
+    # unchanged behaviour so existing callers keep working.
+    if payload.conversation_id is None or conversational is None:
+        grounded = answers.answer(payload.query, payload.subject_scope, payload.intent)
+        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)
+    else:
+        turn = conversational.answer(
+            payload.conversation_id,
+            payload.query,
+            payload.subject_scope,
+            payload.intent,
+        )
+        if turn.clarification is not None:
+            decision, reason = "clarify", turn.clarification.reason
+            answer, resolved_drug_id, citations = turn.clarification.question, None, []
+        elif turn.grounded is not None:
+            decision = turn.grounded.result.decision.value
+            reason = turn.grounded.result.reason
+            answer = turn.answer
+            resolved_drug_id = turn.grounded.result.resolved_drug_id
+            citations = _map_citations(turn.grounded.citations)
+        else:  # smalltalk
+            decision, reason = "answerable", turn.reason
+            answer, resolved_drug_id, citations = turn.answer, None, []
+
+    trace_id = traces.save(
+        query=payload.query,
+        subject_scope=payload.subject_scope.value,
+        intent=payload.intent.value,
+        decision=decision,
+        reason=reason,
+        resolved_drug_id=resolved_drug_id,
+        citations=tuple(item.model_dump() for item in citations),
+    )
+    return RagQueryResponse(
+        trace_id=trace_id,
+        decision=decision,
+        reason=reason,
+        answer=answer,
+        resolved_drug_id=resolved_drug_id,
+        citations=citations,
+    )