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

637 lines
24 KiB
Python

from __future__ import annotations
import uuid
from typing import Annotated, Any, Literal, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query, 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.ports import SectionListRetriever, SectionRetriever
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: ...
def list_by_conversation(self, conversation_id: str, limit: int) -> list[Any]: ...
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)
response_mode: Literal["ai", "monograph"] = "ai"
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()
def _section_retriever(request: Request) -> SectionListRetriever:
retriever = getattr(request.app.state, "section_retriever", None)
if retriever is None:
raise HTTPException(
status_code=503, detail="Section retrieval backend is not configured"
)
return retriever
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 HistoryItem(BaseModel):
trace_id: str
query: str
decision: str
reason: str
resolved_drug_id: str | None = None
created_at: str
class HistoryResponse(BaseModel):
items: list[HistoryItem]
_HISTORY_LIMIT = 50
@router.get("/history", response_model=HistoryResponse)
def list_history(
conversation_id: Annotated[str, Query(max_length=128)],
traces: Annotated[TraceWriter, Depends(_trace_writer)],
) -> HistoryResponse:
"""Feature-List #25: past queries for one session, most recent first, so
the UI can list them and let the user click one to re-run — not to
replay the old answer verbatim, which isn't persisted (see
`PostgresTraceRepository.list_by_conversation`'s docstring). Scoped
strictly to `conversation_id`: this system has no auth anywhere, so an
unscoped listing would leak every session's queries to every caller.
An empty/missing `conversation_id` returns no rows rather than every
trace in the database."""
trimmed = conversation_id.strip()
if not trimmed:
return HistoryResponse(items=[])
try:
rows = traces.list_by_conversation(trimmed, limit=_HISTORY_LIMIT)
except Exception as exc:
raise HTTPException(status_code=503, detail="trace_store_unavailable") from exc
return HistoryResponse(
items=[
HistoryItem(
trace_id=row.trace_id,
query=row.query,
decision=row.decision,
reason=row.reason,
resolved_drug_id=row.resolved_drug_id,
created_at=row.created_at.isoformat() if row.created_at else "",
)
for row in rows
]
)
class TranscriptMessage(BaseModel):
role: Literal["user", "assistant"]
trace_id: str
content: str | None = None
decision: str | None = None
reason: str | None = None
resolved_drug_id: str | None = None
citations: list[dict[str, Any]] = []
generated: bool = False
quick_replies: list[str] = []
blocks: list[dict[str, Any]] = []
answer_mode: str | None = None
answer_plan: dict[str, Any] | None = None
candidate_assessments: list[dict[str, Any]] = []
disclaimer: str | None = None
created_at: str
class TranscriptResponse(BaseModel):
messages: list[TranscriptMessage]
@router.get("/transcript", response_model=TranscriptResponse)
def get_transcript(
conversation_id: Annotated[str, Query(max_length=128)],
traces: Annotated[TraceWriter, Depends(_trace_writer)],
) -> TranscriptResponse:
"""The actual conversation for one session — both the question AND the
answer for each turn, oldest first, so a resumed session can be redrawn
and continued rather than starting blank (the gap `/history` deliberately
leaves open, see its docstring). Same no-auth scoping as `/history`: an
empty/missing `conversation_id` returns nothing.
A row saved before `response_payload` existed (or where persistence
raced a DB outage) has no replayable answer — its turn contributes only
the user message, not a silently-wrong assistant one.
"""
trimmed = conversation_id.strip()
if not trimmed:
return TranscriptResponse(messages=[])
try:
rows = traces.list_by_conversation(trimmed, limit=_HISTORY_LIMIT)
except Exception as exc:
raise HTTPException(status_code=503, detail="trace_store_unavailable") from exc
messages: list[TranscriptMessage] = []
for row in reversed(rows): # list_by_conversation is newest-first
created_at = row.created_at.isoformat() if row.created_at else ""
messages.append(
TranscriptMessage(
role="user",
trace_id=row.trace_id,
content=row.query,
created_at=created_at,
)
)
payload = row.response_payload
if payload is None:
continue
messages.append(
TranscriptMessage(
role="assistant",
trace_id=row.trace_id,
content=payload.get("answer"),
decision=row.decision,
reason=row.reason,
resolved_drug_id=payload.get("resolved_drug_id"),
citations=payload.get("citations") or [],
generated=bool(payload.get("generated", False)),
quick_replies=payload.get("quick_replies") or [],
blocks=payload.get("blocks") or [],
answer_mode=payload.get("answer_mode"),
answer_plan=payload.get("answer_plan"),
candidate_assessments=payload.get("candidate_assessments") or [],
disclaimer=payload.get("disclaimer"),
created_at=created_at,
)
)
return TranscriptResponse(messages=messages)
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()))
class SectionListItem(BaseModel):
section_key: str
section_title: str
class SectionListResponse(BaseModel):
sections: list[SectionListItem]
@router.get("/sections", response_model=SectionListResponse)
def list_drug_sections(
drug_id: str,
retriever: Annotated[SectionListRetriever, Depends(_section_retriever)],
) -> SectionListResponse:
"""Feature-List #4: the real per-drug section checklist, not a generic
fixed list — coverage genuinely varies (measured corpus-wide: 7 to 19
sections per drug). No LLM/generation involved, a plain indexed-payload
read, so an unknown or unresolved `drug_id` returns an empty list rather
than a 404 — the caller (UI attribute picker) already knows which
`drug_id` it resolved before calling this."""
sections = retriever.list_sections(drug_id.strip())
return SectionListResponse(
sections=[
SectionListItem(section_key=key, section_title=title)
for key, title in sections
]
)
class SectionTextPart(BaseModel):
part_index: int | None = None
text: str
# True for a quarantined table/formula chunk — `text` is then the
# chunker's own descriptor sentence ("bảng, trang N..."), not the
# table's content; never a paraphrase, per the quarantine contract
# (docs-legacy/adr/0006). The caller must not present this the same way
# as a real verbatim excerpt.
is_quarantined: bool
printed_page_start: int | None = None
printed_page_end: int | None = None
physical_page: int | None = None
class SectionTextResponse(BaseModel):
drug_id: str
section_key: str
section_title: str | None = None
parts: list[SectionTextPart]
def _section_text_part(hit) -> SectionTextPart:
doc = hit.document
ref = doc.source_refs[0] if doc.source_refs else None
printed_range = ref.printed_page_range if ref else None
printed_start = printed_range[0] if printed_range else (ref.printed_page if ref else None)
printed_end = printed_range[1] if printed_range else (ref.printed_page if ref else None)
return SectionTextPart(
part_index=doc.part_index,
text=doc.text,
is_quarantined=doc.requires_visual_check,
printed_page_start=printed_start,
printed_page_end=printed_end,
physical_page=ref.physical_page if ref else None,
)
@router.get("/section-text", response_model=SectionTextResponse)
def get_section_text(
drug_id: str,
section_key: str,
retriever: Annotated[SectionRetriever, Depends(_section_retriever)],
) -> SectionTextResponse:
"""Feature-List #23: the verbatim source of one section, on demand — no
LLM/generation/entailment involved, so there is nothing to verify;
`evidence_text` on a `/query` citation is the same underlying text but
only for chunks the model actually cited, never a guaranteed whole
section. `find_by_section` already returns every part in book order
(never truncated), which this just joins into an ordered part list —
curation of WHICH sections a UI offers this for (e.g. a "6 mục an
toàn" default) is a client concern; this endpoint is generic to any
real `section_key`, same as `find_by_section` itself."""
hits = retriever.find_by_section(drug_id.strip(), section_key.strip())
parts = [_section_text_part(hit) for hit in hits]
section_title = hits[0].document.section_title if hits else None
return SectionTextResponse(
drug_id=drug_id,
section_key=section_key,
section_title=section_title,
parts=parts,
)
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,
response_mode=payload.response_mode,
)
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()
# The full assistant turn, persisted alongside the trace so a later
# `/transcript` read can redraw this exact answer without re-running the
# query — the shape mirrors `RagQueryResponse` minus the ids (trace_id is
# the row's own primary key; correlation/otel ids are separate columns).
response_payload = {
"answer": answer,
"resolved_drug_id": resolved_drug_id,
"citations": [item.model_dump() for item in citations],
"generated": generated,
"quick_replies": quick_replies,
"blocks": [item.model_dump() for item in blocks],
"answer_mode": answer_mode,
"answer_plan": answer_plan.model_dump() if answer_plan else None,
"candidate_assessments": [item.model_dump() for item in candidate_assessments],
"disclaimer": DISCLAIMER,
}
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,
conversation_id=payload.conversation_id,
response_payload=response_payload,
)
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