Verify exact production traces and record rollout

apps/ai-service/routers/rag.py

@@ -7,9 +7,15 @@ 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.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,
+)
 
 
 class TraceWriter(Protocol):
@@ -61,6 +67,8 @@ class AnswerPlanResponse(BaseModel):
 
 class RagQueryResponse(BaseModel):
     trace_id: str
+    correlation_id: str
+    otel_trace_id: str | None = None
     decision: str
     reason: str
     answer: str | None
@@ -166,7 +174,10 @@ def query_rag(
     traces: Annotated[TraceWriter, Depends(_trace_writer)],
     metrics: Annotated[Metrics, Depends(_metrics)],
 ) -> RagQueryResponse:
-    agent = getattr(request.app.state, "conversational", None)
+    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
@@ -175,9 +186,6 @@ def query_rag(
     # 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),
@@ -199,7 +207,8 @@ def query_rag(
         # 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)
+        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
@@ -229,32 +238,51 @@ def query_rag(
     # 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:
-        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),
-        )
+        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())
-    return RagQueryResponse(
-        trace_id=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,
+    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,
+        )
+    return response