Files
duocthu/apps/ai-service/routers/rag.py
T

378 lines
14 KiB
Python

from __future__ import annotations
import uuid
from typing import Annotated, Any, Literal, Protocol
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, Field
from rag.answer import DISCLAIMER, GroundedAnswerService
from rag.metrics import DECISION, TRACE_WRITE_FAILED, Metrics, NullMetrics
from rag.models import QueryIntent, SubjectScope
from rag.policy import resolve_subject_scope
from rag.telemetry import (
annotate_current_span,
current_correlation_id,
current_trace_id,
stage,
)
from adapters.postgres import FeedbackTraceNotFound
class TraceWriter(Protocol):
def save(self, **fields: Any) -> str: ...
def save_feedback(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
# The exact retrieved chunk text this citation stands for — lets the UI
# show precisely what was retrieved, not a client-side guess at it.
evidence_text: str = ""
drug_id: str | None = None
drug_name: str | None = None
section_key: str | None = None
section_title: str | None = None
source_document: str = "Dược thư Quốc gia Việt Nam 2018"
class AnswerClaimResponse(BaseModel):
text: str
source_ids: list[str]
class AnswerBlockResponse(BaseModel):
title: str
kind: str
claims: list[AnswerClaimResponse]
class AnswerPlanResponse(BaseModel):
verbosity: str
layout: str
reasoning_mode: str
show_heading: bool
needs_warning: bool
class MedicationCandidateAssessmentResponse(BaseModel):
drug_id: str
drug_name: str
indication_supported: bool
status: str
indication_source_ids: list[str]
safety_source_ids: list[str]
class RagFeedbackRequest(BaseModel):
trace_id: uuid.UUID
rating: Literal["helpful", "not_helpful"]
comment: str | None = Field(default=None, max_length=2000)
conversation_id: str | None = Field(default=None, max_length=128)
class RagFeedbackResponse(BaseModel):
feedback_id: str
status: Literal["saved"] = "saved"
class RagQueryResponse(BaseModel):
trace_id: str
correlation_id: str
otel_trace_id: str | None = None
decision: str
reason: str
answer: str | None
resolved_drug_id: str | None
citations: list[CitationResponse]
# Whether `answer` is an LLM paraphrase (verified by grounding +
# entailment) or a verbatim extractive quote of the retrieved source —
# the UI shows these differently so a clinician knows which they're
# reading.
generated: bool = False
# Short suggested replies for a `decision == "clarify"` turn (e.g.
# ["Người lớn", "Trẻ em"]) — only populated by the sufficiency-check
# clarify path today; other clarify sources (no_drug, dosing_calc's
# needs_clarify) leave this empty rather than fabricate options.
quick_replies: list[str] = []
blocks: list[AnswerBlockResponse] = []
answer_mode: str = "concise"
answer_plan: AnswerPlanResponse | None = None
candidate_assessments: list[MedicationCandidateAssessmentResponse] = []
# `docs/architecture.md` specifies the disclaimer at several layers. The
# web banner was the only one in place, so any other consumer received
# medical content with nothing attached. It is a fixed string owned by
# `rag/answer.py`, never produced by the model, and it is present on
# every decision — an abstain or a clarification is still a clinical
# response. Defaulted here as well so a response constructed in a test or
# a future code path cannot accidentally omit it.
disclaimer: str = DISCLAIMER
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"])
@router.post("/feedback", response_model=RagFeedbackResponse)
def save_feedback(
payload: RagFeedbackRequest,
traces: Annotated[TraceWriter, Depends(_trace_writer)],
) -> RagFeedbackResponse:
comment = payload.comment.strip() if payload.comment else None
try:
feedback_id = traces.save_feedback(
trace_id=str(payload.trace_id),
rating=payload.rating,
comment=comment or None,
conversation_id=payload.conversation_id,
)
except FeedbackTraceNotFound as exc:
raise HTTPException(status_code=404, detail="trace_not_found") from exc
except Exception as exc:
raise HTTPException(status_code=503, detail="feedback_store_unavailable") from exc
return RagFeedbackResponse(feedback_id=feedback_id)
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,
evidence_text=item.evidence_text,
drug_id=item.drug_id,
drug_name=item.drug_name,
section_key=item.section_key,
section_title=item.section_title,
source_document=item.source_document,
)
for item in items
]
def _map_blocks(items) -> list[AnswerBlockResponse]:
return [
AnswerBlockResponse(
title=item.title,
kind=item.kind,
claims=[
AnswerClaimResponse(text=claim.text, source_ids=list(claim.source_ids))
for claim in item.claims
],
)
for item in items
]
def _map_plan(item) -> AnswerPlanResponse | None:
if item is None:
return None
return AnswerPlanResponse(
verbosity=item.verbosity,
layout=item.layout,
reasoning_mode=item.reasoning_mode,
show_heading=item.show_heading,
needs_warning=item.needs_warning,
)
def _map_candidate_assessments(items) -> list[MedicationCandidateAssessmentResponse]:
return [
MedicationCandidateAssessmentResponse(
drug_id=item.drug_id,
drug_name=item.drug_name,
indication_supported=item.indication_supported,
status=item.status.value,
indication_source_ids=[
evidence.matched_doc_id for evidence in item.indication_evidence
],
safety_source_ids=[
evidence.matched_doc_id
for evidence in item.evidence
if evidence not in item.indication_evidence
],
)
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:
with stage("receive"):
agent = getattr(request.app.state, "conversational", None)
subject_scope = resolve_subject_scope(payload.query, payload.subject_scope)
intent = payload.intent
# `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).
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)
generated = reply.generated
quick_replies = list(reply.quick_replies)
blocks = _map_blocks(reply.blocks)
answer_mode = reply.answer_mode
answer_plan = _map_plan(reply.plan)
candidate_assessments = _map_candidate_assessments(
reply.candidate_assessments
)
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.
with stage("routing"):
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 = []
generated = False
quick_replies = list(grounded.quick_replies)
blocks = []
answer_mode = "concise"
answer_plan = None
candidate_assessments = []
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)
generated = grounded.generated
quick_replies = []
blocks = _map_blocks(grounded.blocks)
answer_mode = grounded.answer_mode
answer_plan = _map_plan(grounded.plan)
candidate_assessments = []
# 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.
correlation_id = current_correlation_id()
otel_trace_id = current_trace_id()
try:
with stage("persistence"):
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),
correlation_id=correlation_id,
otel_trace_id=otel_trace_id,
)
except Exception:
metrics.increment(TRACE_WRITE_FAILED)
trace_id = str(uuid.uuid4())
metrics.increment(DECISION, decision=decision, reason=reason)
annotate_current_span(
**{
"duocthu.decision": decision,
"duocthu.reason": reason,
"duocthu.citation_count": len(citations),
"duocthu.generated": generated,
"duocthu.persisted_trace_id": trace_id,
}
)
with stage("response"):
response = RagQueryResponse(
trace_id=trace_id,
correlation_id=correlation_id,
otel_trace_id=otel_trace_id,
decision=decision,
reason=reason,
answer=answer,
resolved_drug_id=resolved_drug_id,
citations=citations,
generated=generated,
quick_replies=quick_replies,
blocks=blocks,
answer_mode=answer_mode,
answer_plan=answer_plan,
candidate_assessments=candidate_assessments,
)
return response