Verify exact production traces and record rollout

apps/ai-service/bootstrap.py

@@ -6,14 +6,22 @@ from adapters.embedding import BedrockCohereQueryEmbedder
 from adapters.postgres import PostgresConversationStore, PostgresTraceRepository
 from adapters.qdrant import QdrantParentStore, QdrantRetriever
 from config import Settings
-from rag.agent import RagAgent
-from rag.answer import GroundedAnswerService
 from rag.artifacts import load_aliases
 from rag.manifest import MANIFEST_POINT_ID, check_manifest, manifest_collection
+from rag.instrumentation import (
+    InstrumentedEmbedder,
+    InstrumentedGenerator,
+    InstrumentedGroundedAnswerService,
+    InstrumentedQueryUnderstander,
+    InstrumentedRagAgent,
+    InstrumentedReranker,
+    InstrumentedRetrievalService,
+)
 from rag.metrics import NullMetrics
 from rag.routing import CatalogDrugResolver, QueryRoutingService
 from rag.sections import SectionResolver
-from rag.service import EvidencePolicy, RetrievalService
+from rag.service import EvidencePolicy
+from rag.telemetry import configure_telemetry
 from rag.understanding import LlmQueryUnderstander
 
 # How many aliases to show per candidate drug (F-04 bounds *which* drugs are
@@ -126,6 +134,8 @@ def _verify_corpus_manifest(client, collection: str, embedder, settings: Setting
 
 def build_runtime(settings: Settings):
     metrics = _build_metrics(settings)
+    effective_metrics = metrics or NullMetrics()
+    configure_telemetry(settings, effective_metrics)
     if settings.embedding_provider == "disabled":
         return None, None, PostgresTraceRepository(settings.postgres_dsn), metrics
     if settings.embedding_provider != "cohere-v4":
@@ -140,8 +150,11 @@ def build_runtime(settings: Settings):
         api_key=settings.qdrant_api_key,
         timeout=30,
     )
-    embedder = BedrockCohereQueryEmbedder(
-        settings.embedding_dimensions, region=settings.aws_region
+    embedder = InstrumentedEmbedder(
+        BedrockCohereQueryEmbedder(
+            settings.embedding_dimensions, region=settings.aws_region
+        ),
+        effective_metrics,
     )
     # F-05: a collection built with one model and queried with another
     # returns hits and raises nothing — the results are just meaningless,
@@ -155,17 +168,25 @@ def build_runtime(settings: Settings):
     # `.resolve()` is no longer on the live query path; `RagAgent` resolves
     # drug identity through `LlmQueryUnderstander` against the same catalog.
     resolver = CatalogDrugResolver(aliases)
-    retrieval = RetrievalService(
+    reranker = _build_reranker(settings)
+    retrieval = InstrumentedRetrievalService(
         QdrantRetriever(client, settings.qdrant_collection, embedder),
         QdrantParentStore(client, settings.qdrant_collection),
         EvidencePolicy(minimum_score=settings.evidence_minimum_score),
         section_resolver=section_resolver,
-        reranker=_build_reranker(settings),
+        reranker=(
+            InstrumentedReranker(reranker, effective_metrics)
+            if reranker is not None
+            else None
+        ),
+        metrics=effective_metrics,
     )
     routing = QueryRoutingService(retrieval, resolver)
     generator = _build_generator(settings)
-    answers = GroundedAnswerService(
-        routing, generator=generator, metrics=metrics or NullMetrics()
+    if generator is not None:
+        generator = InstrumentedGenerator(generator, effective_metrics)
+    answers = InstrumentedGroundedAnswerService(
+        routing, generator=generator, metrics=effective_metrics
     )
     trace_writer = PostgresTraceRepository(settings.postgres_dsn)
     if generator is None:
@@ -175,13 +196,16 @@ def build_runtime(settings: Settings):
         # capability to offer. Answer-only (retrieval-verified, no
         # generation) still works through `answers` directly.
         return answers, None, trace_writer, metrics
-    agent = RagAgent(
-        understander=LlmQueryUnderstander(generator, _catalog_names(aliases), resolver),
+    agent = InstrumentedRagAgent(
+        understander=InstrumentedQueryUnderstander(
+            LlmQueryUnderstander(generator, _catalog_names(aliases), resolver)
+        ),
         retrieval=retrieval,
         answers=answers,
         autocomplete=resolver,
         max_wall_clock_ms=settings.max_wall_clock_ms,
         max_llm_calls_per_turn=settings.max_llm_calls_per_turn,
         store=PostgresConversationStore(settings.postgres_dsn),
+        metrics=effective_metrics,
     )
     return answers, agent, trace_writer, metrics