Verify exact production traces and record rollout
This commit is contained in:
@@ -26,6 +26,9 @@ RUN pip install --no-cache-dir \
|
||||
"qdrant-client>=1.7,<2" \
|
||||
"uvicorn[standard]>=0.30,<1" \
|
||||
"prometheus-client>=0.20,<1" \
|
||||
"opentelemetry-api>=1.27,<2" \
|
||||
"opentelemetry-sdk>=1.27,<2" \
|
||||
"opentelemetry-exporter-otlp-proto-http>=1.27,<2" \
|
||||
"anthropic>=0.112,<1" \
|
||||
"boto3"
|
||||
|
||||
|
||||
@@ -18,6 +18,8 @@ class RetrievalTrace:
|
||||
reason: str
|
||||
resolved_drug_id: str | None
|
||||
citations: tuple[dict[str, Any], ...]
|
||||
correlation_id: str | None = None
|
||||
otel_trace_id: str | None = None
|
||||
created_at: datetime | None = None
|
||||
|
||||
|
||||
@@ -55,6 +57,8 @@ class PostgresTraceRepository:
|
||||
reason: str,
|
||||
resolved_drug_id: str | None,
|
||||
citations: tuple[dict[str, Any], ...],
|
||||
correlation_id: str | None = None,
|
||||
otel_trace_id: str | None = None,
|
||||
) -> str:
|
||||
import psycopg
|
||||
|
||||
@@ -64,12 +68,14 @@ class PostgresTraceRepository:
|
||||
"""
|
||||
INSERT INTO rag_retrieval_trace (
|
||||
trace_id, query_text, subject_scope, query_intent,
|
||||
decision, reason, resolved_drug_id, citations
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb)
|
||||
decision, reason, resolved_drug_id, citations,
|
||||
correlation_id, otel_trace_id
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s)
|
||||
""",
|
||||
(
|
||||
trace_id, query, subject_scope, intent, decision, reason,
|
||||
resolved_drug_id, json.dumps(citations, ensure_ascii=False),
|
||||
correlation_id, otel_trace_id,
|
||||
),
|
||||
)
|
||||
return trace_id
|
||||
@@ -81,7 +87,8 @@ class PostgresTraceRepository:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT trace_id::text, query_text, subject_scope, query_intent,
|
||||
decision, reason, resolved_drug_id, citations, created_at
|
||||
decision, reason, resolved_drug_id, citations,
|
||||
correlation_id, otel_trace_id, created_at
|
||||
FROM rag_retrieval_trace WHERE trace_id = %s
|
||||
""",
|
||||
(trace_id,),
|
||||
@@ -91,7 +98,8 @@ class PostgresTraceRepository:
|
||||
return RetrievalTrace(
|
||||
trace_id=row[0], query=row[1], subject_scope=row[2], intent=row[3],
|
||||
decision=row[4], reason=row[5], resolved_drug_id=row[6],
|
||||
citations=tuple(row[7]), created_at=row[8],
|
||||
citations=tuple(row[7]), correlation_id=row[8], otel_trace_id=row[9],
|
||||
created_at=row[10],
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
"""Exports the domain's counters; the only module that names prometheus_client.
|
||||
"""Prometheus export for bounded domain and request-path telemetry.
|
||||
|
||||
`rag/metrics.py` defines what is counted and why. This decides how it leaves
|
||||
the process, and is imported lazily so the service runs — and the suite passes
|
||||
— with no metrics stack installed.
|
||||
|
||||
Counter names carry a `duocthu_` prefix and a `_total` suffix because that is
|
||||
what Prometheus expects of a counter; the dashboard queries them by name.
|
||||
The domain owns metric names in :mod:`rag.metrics`; this adapter owns label
|
||||
vocabularies, buckets and OpenMetrics exposition. Unknown label values are
|
||||
collapsed to ``other`` so a query, drug id, exception message or raw URL can
|
||||
never accidentally create an unbounded time series.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -14,44 +12,192 @@ from typing import Any
|
||||
from rag.metrics import (
|
||||
ABSTENTION,
|
||||
ANSWER_EXTRACTIVE,
|
||||
CLARIFY_ASKED,
|
||||
DECISION,
|
||||
FOLLOWUP_INHERITED,
|
||||
GENERATION_REJECTED,
|
||||
GENERATION_SERVED,
|
||||
LOOP_REFINED,
|
||||
LOOP_REPAIRED,
|
||||
LOOP_ROUNDS,
|
||||
PROVIDER_FAILURE,
|
||||
REQUEST_DURATION,
|
||||
REQUESTS,
|
||||
RETRIEVAL_ROUTE,
|
||||
STAGE_DURATION,
|
||||
TRACE_WRITE_FAILED,
|
||||
)
|
||||
|
||||
_LABELS: dict[str, tuple[str, ...]] = {
|
||||
_COUNTER_LABELS: dict[str, tuple[str, ...]] = {
|
||||
ABSTENTION: ("reason",),
|
||||
GENERATION_REJECTED: ("reason",),
|
||||
RETRIEVAL_ROUTE: ("route",),
|
||||
GENERATION_SERVED: (),
|
||||
ANSWER_EXTRACTIVE: (),
|
||||
CLARIFY_ASKED: ("reason",),
|
||||
DECISION: ("decision", "reason"),
|
||||
FOLLOWUP_INHERITED: (),
|
||||
GENERATION_REJECTED: ("reason",),
|
||||
GENERATION_SERVED: (),
|
||||
LOOP_REFINED: (),
|
||||
LOOP_REPAIRED: (),
|
||||
LOOP_ROUNDS: (),
|
||||
PROVIDER_FAILURE: ("provider", "operation", "reason"),
|
||||
REQUESTS: ("method", "route", "status"),
|
||||
RETRIEVAL_ROUTE: ("route",),
|
||||
TRACE_WRITE_FAILED: (),
|
||||
}
|
||||
|
||||
_HISTOGRAM_LABELS: dict[str, tuple[str, ...]] = {
|
||||
REQUEST_DURATION: ("method", "route", "status"),
|
||||
STAGE_DURATION: ("stage", "outcome"),
|
||||
}
|
||||
|
||||
_HELP = {
|
||||
ABSTENTION: "Answers refused, by the reason retrieval gave.",
|
||||
GENERATION_REJECTED: (
|
||||
"Generations discarded before reaching the caller. `reason=\"ungrounded_number\"` "
|
||||
"counts answers that stated a figure absent from the cited source; "
|
||||
"`reason=\"uncited_claim\"` counts claims with no valid citation at all; "
|
||||
"`reason=\"unsupported_claim\"` counts claims the entailment pass judged "
|
||||
"not actually stated by the block they cite."
|
||||
),
|
||||
GENERATION_SERVED: "Generations that passed grounding verification and were served.",
|
||||
ABSTENTION: "Answers refused, by bounded domain reason.",
|
||||
ANSWER_EXTRACTIVE: "Answers served as verbatim source text.",
|
||||
RETRIEVAL_ROUTE: "Retrievals by route: section filter, or similarity fallback.",
|
||||
CLARIFY_ASKED: "Clarifying questions returned instead of guessing.",
|
||||
DECISION: "Final request decisions by bounded domain reason.",
|
||||
FOLLOWUP_INHERITED: "Follow-up turns that inherited prior context.",
|
||||
GENERATION_REJECTED: "Generated answers discarded by a safety or availability gate.",
|
||||
GENERATION_SERVED: "Generated answers that passed grounding and entailment.",
|
||||
LOOP_REFINED: "Conversational retrieval loops that refined a query.",
|
||||
LOOP_REPAIRED: "Conversational retrieval loops that repaired an answer.",
|
||||
LOOP_ROUNDS: "Retrieval loop rounds completed.",
|
||||
PROVIDER_FAILURE: "Failures from bounded external-provider operations.",
|
||||
REQUESTS: "HTTP requests by route template and status class.",
|
||||
REQUEST_DURATION: "End-to-end HTTP request latency in seconds.",
|
||||
RETRIEVAL_ROUTE: "Retrievals by bounded route.",
|
||||
STAGE_DURATION: "RAG stage latency in seconds.",
|
||||
TRACE_WRITE_FAILED: "Final trace rows that could not be persisted to PostgreSQL.",
|
||||
}
|
||||
|
||||
_REQUEST_BUCKETS = (0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 20, 40, 60)
|
||||
_STAGE_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 20, 40)
|
||||
|
||||
_ALLOWED: dict[str, frozenset[str]] = {
|
||||
"decision": frozenset({"answerable", "clarify", "verify_pdf", "abstain", "error"}),
|
||||
"method": frozenset({"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"}),
|
||||
"operation": frozenset(
|
||||
{"understand", "embed", "retrieve", "rerank", "generate", "entailment", "persist"}
|
||||
),
|
||||
"outcome": frozenset({"ok", "error", "cancelled"}),
|
||||
"provider": frozenset(
|
||||
{"bedrock_converse", "bedrock_claude", "bedrock_cohere", "qdrant", "postgres"}
|
||||
),
|
||||
"route": frozenset(
|
||||
{
|
||||
"/health",
|
||||
"/ready",
|
||||
"/metrics",
|
||||
"/v1/rag/query",
|
||||
"/v1/rag/suggest",
|
||||
"section",
|
||||
"overview",
|
||||
"similarity",
|
||||
"indication",
|
||||
"other",
|
||||
}
|
||||
),
|
||||
"stage": frozenset(
|
||||
{
|
||||
"receive",
|
||||
"context",
|
||||
"understanding",
|
||||
"routing",
|
||||
"retrieval",
|
||||
"rerank",
|
||||
"evidence",
|
||||
"generation",
|
||||
"grounding",
|
||||
"entailment",
|
||||
"persistence",
|
||||
"response",
|
||||
}
|
||||
),
|
||||
"status": frozenset({"1xx", "2xx", "3xx", "4xx", "5xx"}),
|
||||
}
|
||||
|
||||
_REASONS = frozenset(
|
||||
{
|
||||
"clarify_loop_exhausted",
|
||||
"drug_not_in_formulary",
|
||||
"drug_not_resolved",
|
||||
"drug_resolution_ambiguous",
|
||||
"drug_resolution_invalid_state",
|
||||
"evidence_insufficient",
|
||||
"error",
|
||||
"generation_unavailable",
|
||||
"grounded_evidence_available",
|
||||
"incomplete_answer",
|
||||
"insufficient_retrieval_score",
|
||||
"invalid_citation",
|
||||
"malformed_output",
|
||||
"missing_attribute",
|
||||
"missing_indication",
|
||||
"missing_pediatric_age_or_weight",
|
||||
"missing_population",
|
||||
"missing_printed_page_provenance",
|
||||
"missing_provenance",
|
||||
"missing_query_or_drug",
|
||||
"needs_more_info",
|
||||
"no_drug",
|
||||
"no_indication",
|
||||
"no_indication_match",
|
||||
"no_interaction_evidence",
|
||||
"out_of_scope",
|
||||
"out_of_scope_non_human",
|
||||
"parent_hydration_failed",
|
||||
"provider_unavailable",
|
||||
"query_embedding_unavailable",
|
||||
"query_intent_unknown",
|
||||
"request_budget_exhausted",
|
||||
"smalltalk",
|
||||
"subject_scope_unknown",
|
||||
"timeout",
|
||||
"uncited_claim",
|
||||
"ungrounded_number",
|
||||
"unsupported_claim",
|
||||
"visual_verification_required",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _bounded(label: str, value: str) -> str:
|
||||
allowed = _REASONS if label == "reason" else _ALLOWED.get(label)
|
||||
if allowed is None:
|
||||
return value
|
||||
return value if value in allowed else "other"
|
||||
|
||||
|
||||
def _exemplar() -> dict[str, str] | None:
|
||||
# Lazy import avoids making OpenTelemetry a hard requirement for metrics.
|
||||
try:
|
||||
from rag.telemetry import current_trace_id
|
||||
|
||||
trace_id = current_trace_id()
|
||||
except ImportError:
|
||||
trace_id = None
|
||||
return {"trace_id": trace_id} if trace_id else None
|
||||
|
||||
|
||||
class PrometheusMetrics:
|
||||
"""Domain `Metrics` backed by a Prometheus registry."""
|
||||
"""Domain ``Metrics`` backed by an isolated Prometheus registry."""
|
||||
|
||||
def __init__(self, registry: Any | None = None) -> None:
|
||||
from prometheus_client import CollectorRegistry, Counter
|
||||
from prometheus_client import CollectorRegistry, Counter, Histogram
|
||||
|
||||
self._registry = registry or CollectorRegistry()
|
||||
self._counters = {
|
||||
name: Counter(name, _HELP[name], labels, registry=self._registry)
|
||||
for name, labels in _LABELS.items()
|
||||
for name, labels in _COUNTER_LABELS.items()
|
||||
}
|
||||
self._histograms = {
|
||||
name: Histogram(
|
||||
name,
|
||||
_HELP[name],
|
||||
labels,
|
||||
buckets=_REQUEST_BUCKETS if name == REQUEST_DURATION else _STAGE_BUCKETS,
|
||||
registry=self._registry,
|
||||
)
|
||||
for name, labels in _HISTOGRAM_LABELS.items()
|
||||
}
|
||||
|
||||
@property
|
||||
@@ -60,17 +206,28 @@ class PrometheusMetrics:
|
||||
|
||||
def increment(self, name: str, **labels: str) -> None:
|
||||
counter = self._counters.get(name)
|
||||
if counter is None:
|
||||
expected = _COUNTER_LABELS.get(name)
|
||||
if counter is None or expected is None or set(labels) != set(expected):
|
||||
return
|
||||
# An unexpected label would raise at scrape time, far from its cause.
|
||||
# Metrics must not be able to break a clinical answer, so a mismatch
|
||||
# drops the sample rather than the request.
|
||||
expected = set(_LABELS[name])
|
||||
if set(labels) != expected:
|
||||
safe = {key: _bounded(key, labels[key]) for key in expected}
|
||||
child = counter.labels(**safe) if safe else counter
|
||||
child.inc(exemplar=_exemplar())
|
||||
|
||||
def observe(self, name: str, value: float, **labels: str) -> None:
|
||||
histogram = self._histograms.get(name)
|
||||
expected = _HISTOGRAM_LABELS.get(name)
|
||||
if histogram is None or expected is None or set(labels) != set(expected):
|
||||
return
|
||||
(counter.labels(**labels) if labels else counter).inc()
|
||||
safe = {key: _bounded(key, labels[key]) for key in expected}
|
||||
child = histogram.labels(**safe) if safe else histogram
|
||||
child.observe(value, exemplar=_exemplar())
|
||||
|
||||
def render(self) -> tuple[bytes, str]:
|
||||
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
|
||||
# OpenMetrics exposition preserves histogram exemplars. Grafana uses
|
||||
# their trace_id label to jump from a slow aggregate bucket to Tempo.
|
||||
from prometheus_client.openmetrics.exposition import (
|
||||
CONTENT_TYPE_LATEST,
|
||||
generate_latest,
|
||||
)
|
||||
|
||||
return generate_latest(self._registry), CONTENT_TYPE_LATEST
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -25,6 +25,7 @@ class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(env_file=".env", extra="ignore")
|
||||
|
||||
app_name: str = "vsf-duoc-thu-ai-service"
|
||||
environment: str = "local"
|
||||
qdrant_url: str = "http://localhost:6333"
|
||||
qdrant_collection: str = "duocthu_v1"
|
||||
qdrant_api_key: str | None = None
|
||||
@@ -50,6 +51,13 @@ class Settings(BaseSettings):
|
||||
# never uses it.
|
||||
rerank_enabled: bool = False
|
||||
metrics_enabled: bool = True
|
||||
# OpenTelemetry is opt-in so the existing EC2 Compose deployment keeps
|
||||
# answering when no collector is present. Docker/Kubernetes observability
|
||||
# profiles enable it and point OTLP/HTTP at their local collector Service.
|
||||
otel_enabled: bool = False
|
||||
otel_service_name: str = "ai-service"
|
||||
otel_exporter_otlp_endpoint: str = "http://localhost:4318/v1/traces"
|
||||
otel_sample_ratio: float = Field(default=1.0, ge=0.0, le=1.0)
|
||||
entities_path: Path = _default_entities_path()
|
||||
# F-08: a per-turn budget across RagAgent's sequential Bedrock calls
|
||||
# (understand, generate, one entailment check on the live agent path).
|
||||
|
||||
+64
-1
@@ -1,13 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from time import perf_counter
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, Response
|
||||
from fastapi import FastAPI, Request, Response
|
||||
|
||||
from adapters.postgres import PostgresTraceRepository
|
||||
from bootstrap import build_runtime
|
||||
from config import Settings, get_settings
|
||||
from rag.answer import GroundedAnswerService
|
||||
from rag.metrics import REQUEST_DURATION, REQUESTS, NullMetrics
|
||||
from rag.telemetry import (
|
||||
annotate_current_span,
|
||||
configure_telemetry,
|
||||
correlation_context,
|
||||
current_trace_id,
|
||||
request_span,
|
||||
)
|
||||
from routers.rag import router as rag_router
|
||||
|
||||
|
||||
@@ -20,16 +29,65 @@ def create_app(
|
||||
metrics: Any | None = None,
|
||||
) -> FastAPI:
|
||||
configured = settings or get_settings()
|
||||
effective_metrics = metrics or NullMetrics()
|
||||
configure_telemetry(configured, effective_metrics)
|
||||
app = FastAPI(title=configured.app_name, version="0.1.0")
|
||||
app.state.answer_service = answer_service
|
||||
app.state.conversational = conversational
|
||||
app.state.trace_writer = trace_writer
|
||||
app.state.metrics = metrics
|
||||
|
||||
@app.middleware("http")
|
||||
async def correlate_and_trace(request: Request, call_next):
|
||||
route = _route_label(request.url.path)
|
||||
method = request.method.upper()
|
||||
started = perf_counter()
|
||||
status_code = 500
|
||||
with correlation_context(request.headers.get("x-correlation-id")) as correlation_id:
|
||||
with request_span(method, route, request.headers) as span:
|
||||
try:
|
||||
response = await call_next(request)
|
||||
status_code = response.status_code
|
||||
response.headers["X-Correlation-ID"] = correlation_id
|
||||
trace_id = current_trace_id()
|
||||
if trace_id:
|
||||
response.headers["X-Trace-ID"] = trace_id
|
||||
return response
|
||||
finally:
|
||||
status = f"{status_code // 100}xx"
|
||||
elapsed = perf_counter() - started
|
||||
effective_metrics.increment(
|
||||
REQUESTS, method=method, route=route, status=status
|
||||
)
|
||||
effective_metrics.observe(
|
||||
REQUEST_DURATION,
|
||||
elapsed,
|
||||
method=method,
|
||||
route=route,
|
||||
status=status,
|
||||
)
|
||||
if span is not None:
|
||||
span.set_attribute("http.response.status_code", status_code)
|
||||
annotate_current_span(
|
||||
**{
|
||||
"duocthu.http.status_class": status,
|
||||
"duocthu.duration_ms": elapsed * 1000,
|
||||
}
|
||||
)
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.get("/ready")
|
||||
def ready() -> Response:
|
||||
# Runtime construction already verifies the Qdrant corpus manifest and
|
||||
# fails startup on mismatch. PostgreSQL trace/history are intentionally
|
||||
# fail-open and therefore must not make readiness flap.
|
||||
if app.state.answer_service is None and configured.embedding_provider != "disabled":
|
||||
return Response(status_code=503)
|
||||
return Response(content='{"status":"ready"}', media_type="application/json")
|
||||
|
||||
@app.get("/metrics")
|
||||
def prometheus_metrics() -> Response:
|
||||
exporter = getattr(app.state, "metrics", None)
|
||||
@@ -44,6 +102,11 @@ def create_app(
|
||||
return app
|
||||
|
||||
|
||||
def _route_label(path: str) -> str:
|
||||
known = {"/health", "/ready", "/metrics", "/v1/rag/query", "/v1/rag/suggest"}
|
||||
return path if path in known else "other"
|
||||
|
||||
|
||||
_settings = get_settings()
|
||||
_answer_service, _conversational, _trace_writer, _metrics = build_runtime(_settings)
|
||||
app = create_app(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from pathlib import Path
|
||||
|
||||
from adapters.postgres import PostgresConversationStore, PostgresTraceRepository
|
||||
from adapters.postgres import PostgresTraceRepository
|
||||
from config import get_settings
|
||||
|
||||
|
||||
@@ -8,13 +8,10 @@ def main() -> None:
|
||||
dsn = get_settings().postgres_dsn
|
||||
migrations_dir = Path(__file__).parent / "migrations"
|
||||
|
||||
trace_migration = migrations_dir / "001_rag_retrieval_trace.sql"
|
||||
PostgresTraceRepository(dsn).migrate(trace_migration)
|
||||
print(f"Applied {trace_migration.name}")
|
||||
|
||||
conversation_migration = migrations_dir / "002_rag_conversation_turn.sql"
|
||||
PostgresConversationStore(dsn).migrate(conversation_migration)
|
||||
print(f"Applied {conversation_migration.name}")
|
||||
repository = PostgresTraceRepository(dsn)
|
||||
for migration in sorted(migrations_dir.glob("*.sql")):
|
||||
repository.migrate(migration)
|
||||
print(f"Applied {migration.name}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
ALTER TABLE rag_retrieval_trace
|
||||
ADD COLUMN IF NOT EXISTS correlation_id text,
|
||||
ADD COLUMN IF NOT EXISTS otel_trace_id varchar(32);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS rag_retrieval_trace_correlation_idx
|
||||
ON rag_retrieval_trace (correlation_id)
|
||||
WHERE correlation_id IS NOT NULL;
|
||||
|
||||
CREATE INDEX IF NOT EXISTS rag_retrieval_trace_otel_idx
|
||||
ON rag_retrieval_trace (otel_trace_id)
|
||||
WHERE otel_trace_id IS NOT NULL;
|
||||
@@ -18,6 +18,11 @@ test = ["pytest>=7.4,<9"]
|
||||
# without a cloud generator. Neither is a precondition for a grounded answer.
|
||||
metrics = ["prometheus-client>=0.20,<1"]
|
||||
generation = ["anthropic>=0.112,<1"]
|
||||
observability = [
|
||||
"opentelemetry-api>=1.27,<2",
|
||||
"opentelemetry-sdk>=1.27,<2",
|
||||
"opentelemetry-exporter-otlp-proto-http>=1.27,<2",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
|
||||
@@ -99,12 +99,30 @@ class AnswerPlan:
|
||||
needs_warning: bool = False
|
||||
|
||||
|
||||
# A fixed, non-LLM string. `docs/architecture.md`'s guardrail section
|
||||
# specifies the disclaimer at several layers, and the web banner
|
||||
# (`packages/ui/src/DisclaimerBanner.tsx`) was the only one in place: the
|
||||
# `disclaimer` field declared in `packages/shared-types/src/dto/chat.ts` was
|
||||
# never populated, so any consumer other than this one web UI received medical
|
||||
# content with nothing attached. Keeping it out of the prompt is deliberate —
|
||||
# a disclaimer the model writes is one the model can also reword, shorten or
|
||||
# omit, and it would then have to be verified like any other generated claim.
|
||||
DISCLAIMER = (
|
||||
"Nội dung được trích từ Dược thư Quốc gia Việt Nam 2018, phục vụ tra cứu "
|
||||
"chuyên môn và không thay thế chỉ định của bác sĩ hoặc dược sĩ lâm sàng."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GroundedAnswer:
|
||||
result: RetrievalResult
|
||||
answer: str | None
|
||||
citations: tuple[Citation, ...] = ()
|
||||
generated: bool = False
|
||||
# Carried on the payload rather than added by the caller, so no response
|
||||
# path can be built that omits it — including abstains and clarifications,
|
||||
# which are also medical content in the sense that matters here.
|
||||
disclaimer: str = DISCLAIMER
|
||||
# Set when the model decided the turn is under-specified and asked back
|
||||
# (e.g. a dose question with no age/weight). The answer field carries the
|
||||
# question; the caller renders it as a clarification, not a final answer.
|
||||
|
||||
@@ -41,6 +41,8 @@ from __future__ import annotations
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .telemetry import traced_stage
|
||||
|
||||
# A digit run with internal separators kept: "500", "7,5", "1.000".
|
||||
# Ranges ("4 - 6 giờ") yield two tokens, and each is checked on its own.
|
||||
_NUMBER = re.compile(r"\d+(?:[.,]\d+)*")
|
||||
@@ -127,6 +129,7 @@ def split_claims(answer: str, evidence_count: int) -> tuple[Claim, ...]:
|
||||
return tuple(claims)
|
||||
|
||||
|
||||
@traced_stage("grounding")
|
||||
def verify(answer: str, evidence_texts: tuple[str, ...]) -> GroundingReport:
|
||||
"""Whether `answer` states only figures and sources traceable to the
|
||||
specific evidence block(s) cited immediately after each claim.
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
"""Non-invasive instrumentation wrappers for the live RAG object graph.
|
||||
|
||||
Claude owns ``rag/agent.py`` and ``rag/answer.py`` in the shared worktree.
|
||||
Subclasses here add spans at their stable method boundaries without changing
|
||||
those files or duplicating their domain decisions.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from .agent import RagAgent
|
||||
from .answer import GroundedAnswerService
|
||||
from .metrics import CLARIFY_ASKED, Metrics, PROVIDER_FAILURE, RETRIEVAL_ROUTE
|
||||
from .service import RetrievalService
|
||||
from .telemetry import (
|
||||
annotate_current_span,
|
||||
current_stage,
|
||||
dependency_span,
|
||||
stage,
|
||||
)
|
||||
|
||||
|
||||
def _provider_name(delegate: Any) -> str:
|
||||
name = delegate.__class__.__name__.casefold()
|
||||
if "cohere" in name:
|
||||
return "bedrock_cohere"
|
||||
if "converse" in name:
|
||||
return "bedrock_converse"
|
||||
if "claude" in name:
|
||||
return "bedrock_claude"
|
||||
return "other"
|
||||
|
||||
|
||||
def _failure_reason(exc: BaseException) -> str:
|
||||
name = exc.__class__.__name__.casefold()
|
||||
if "timeout" in name:
|
||||
return "timeout"
|
||||
if "budget" in name:
|
||||
return "request_budget_exhausted"
|
||||
if "unavailable" in name or "connection" in name:
|
||||
return "provider_unavailable"
|
||||
return "error"
|
||||
|
||||
|
||||
class InstrumentedGenerator:
|
||||
def __init__(self, delegate: Any, metrics: Metrics) -> None:
|
||||
self._delegate = delegate
|
||||
self._metrics = metrics
|
||||
self._provider = _provider_name(delegate)
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._delegate.model_id
|
||||
|
||||
def generate(self, system: str, user: str, schema: dict) -> str:
|
||||
operation = {
|
||||
"understanding": "understand",
|
||||
"generation": "generate",
|
||||
"entailment": "entailment",
|
||||
}.get(current_stage(), "generate")
|
||||
with dependency_span(self._provider, operation):
|
||||
try:
|
||||
return self._delegate.generate(system, user, schema)
|
||||
except Exception as exc:
|
||||
self._metrics.increment(
|
||||
PROVIDER_FAILURE,
|
||||
provider=self._provider,
|
||||
operation=operation,
|
||||
reason=_failure_reason(exc),
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class InstrumentedEmbedder:
|
||||
def __init__(self, delegate: Any, metrics: Metrics) -> None:
|
||||
self._delegate = delegate
|
||||
self._metrics = metrics
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._delegate.dimensions
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._delegate.model_id
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
with dependency_span("bedrock_cohere", "embed"):
|
||||
try:
|
||||
return self._delegate.embed_query(text)
|
||||
except Exception as exc:
|
||||
self._metrics.increment(
|
||||
PROVIDER_FAILURE,
|
||||
provider="bedrock_cohere",
|
||||
operation="embed",
|
||||
reason=_failure_reason(exc),
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class InstrumentedReranker:
|
||||
def __init__(self, delegate: Any, metrics: Metrics) -> None:
|
||||
self._delegate = delegate
|
||||
self._metrics = metrics
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], top_n: int | None = None
|
||||
) -> list[int]:
|
||||
with dependency_span("bedrock_cohere", "rerank"):
|
||||
try:
|
||||
return self._delegate.rerank(query, documents, top_n=top_n)
|
||||
except Exception as exc:
|
||||
self._metrics.increment(
|
||||
PROVIDER_FAILURE,
|
||||
provider="bedrock_cohere",
|
||||
operation="rerank",
|
||||
reason=_failure_reason(exc),
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class InstrumentedRetrievalService(RetrievalService):
|
||||
def __init__(self, *args, metrics: Metrics, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._observability_metrics = metrics
|
||||
|
||||
def retrieve_framed(self, *args, **kwargs):
|
||||
with stage("retrieval"):
|
||||
section_key = args[1] if len(args) > 1 else kwargs.get("section_key")
|
||||
is_overview = args[3] if len(args) > 3 else kwargs.get("is_overview", False)
|
||||
route = "section" if section_key else (
|
||||
"overview" if is_overview else "similarity"
|
||||
)
|
||||
self._observability_metrics.increment(RETRIEVAL_ROUTE, route=route)
|
||||
try:
|
||||
result = super().retrieve_framed(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
self._record_retrieval_failure(exc)
|
||||
raise
|
||||
self._annotate_result(result)
|
||||
return result
|
||||
|
||||
def retrieve(self, *args, **kwargs):
|
||||
with stage("retrieval"):
|
||||
self._observability_metrics.increment(RETRIEVAL_ROUTE, route="other")
|
||||
try:
|
||||
result = super().retrieve(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
self._record_retrieval_failure(exc)
|
||||
raise
|
||||
self._annotate_result(result)
|
||||
return result
|
||||
|
||||
def retrieve_by_indication(self, *args, **kwargs):
|
||||
with stage("retrieval"):
|
||||
self._observability_metrics.increment(RETRIEVAL_ROUTE, route="indication")
|
||||
try:
|
||||
result = super().retrieve_by_indication(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
self._record_retrieval_failure(exc)
|
||||
raise
|
||||
self._annotate_result(result)
|
||||
return result
|
||||
|
||||
def _record_retrieval_failure(self, exc: BaseException) -> None:
|
||||
self._observability_metrics.increment(
|
||||
PROVIDER_FAILURE,
|
||||
provider="qdrant",
|
||||
operation="retrieve",
|
||||
reason=_failure_reason(exc),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _annotate_result(result) -> None:
|
||||
annotate_current_span(
|
||||
**{
|
||||
"duocthu.decision": result.decision.value,
|
||||
"duocthu.reason": result.reason,
|
||||
"duocthu.evidence_count": len(result.evidence),
|
||||
}
|
||||
)
|
||||
|
||||
def _rerank(self, query, hits):
|
||||
with stage("rerank", configured=self._reranker is not None):
|
||||
return super()._rerank(query, hits)
|
||||
|
||||
def _decide(self, evidence, is_drug_overview=False):
|
||||
with stage("evidence", evidence_count=len(evidence)):
|
||||
return super()._decide(evidence, is_drug_overview=is_drug_overview)
|
||||
|
||||
|
||||
class InstrumentedGroundedAnswerService(GroundedAnswerService):
|
||||
def _generate(self, *args, **kwargs):
|
||||
with stage("generation"):
|
||||
return super()._generate(*args, **kwargs)
|
||||
|
||||
def _verify_entailment(self, *args, **kwargs):
|
||||
with stage("entailment"):
|
||||
return super()._verify_entailment(*args, **kwargs)
|
||||
|
||||
|
||||
class InstrumentedQueryUnderstander:
|
||||
def __init__(self, delegate: Any) -> None:
|
||||
self._delegate = delegate
|
||||
|
||||
def understand(self, *args, **kwargs):
|
||||
with stage("understanding"):
|
||||
frame = self._delegate.understand(*args, **kwargs)
|
||||
annotate_current_span(
|
||||
**{
|
||||
"duocthu.turn_type": frame.turn_type,
|
||||
"duocthu.needs_clarify": frame.needs_clarify,
|
||||
"duocthu.system_error": frame.system_error,
|
||||
}
|
||||
)
|
||||
return frame
|
||||
|
||||
|
||||
class InstrumentedRagAgent(RagAgent):
|
||||
def __init__(self, *args, metrics: Metrics, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._observability_metrics = metrics
|
||||
|
||||
def handle(self, *args, **kwargs):
|
||||
reply = super().handle(*args, **kwargs)
|
||||
if reply.decision == "clarify":
|
||||
self._observability_metrics.increment(CLARIFY_ASKED, reason=reply.reason)
|
||||
return reply
|
||||
|
||||
def _get_history(self, *args, **kwargs):
|
||||
with stage("context"):
|
||||
return super()._get_history(*args, **kwargs)
|
||||
|
||||
def _route(self, *args, **kwargs):
|
||||
with stage("routing"):
|
||||
return super()._route(*args, **kwargs)
|
||||
|
||||
def _remember(self, *args, **kwargs):
|
||||
with stage("persistence"):
|
||||
return super()._remember(*args, **kwargs)
|
||||
@@ -18,18 +18,28 @@ from typing import Protocol
|
||||
class Metrics(Protocol):
|
||||
def increment(self, name: str, **labels: str) -> None: ...
|
||||
|
||||
def observe(self, name: str, value: float, **labels: str) -> None: ...
|
||||
|
||||
|
||||
class NullMetrics:
|
||||
def increment(self, name: str, **labels: str) -> None: # noqa: ARG002
|
||||
# Deliberately inert: the default when no metrics stack is configured.
|
||||
return None
|
||||
|
||||
def observe(
|
||||
self, name: str, value: float, **labels: str # noqa: ARG002
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class InMemoryMetrics:
|
||||
"""Reference implementation of the contract; also what tests assert on."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counts: dict[tuple[str, tuple[tuple[str, str], ...]], int] = {}
|
||||
self.observations: dict[
|
||||
tuple[str, tuple[tuple[str, str], ...]], list[float]
|
||||
] = {}
|
||||
|
||||
def increment(self, name: str, **labels: str) -> None:
|
||||
key = (name, tuple(sorted(labels.items())))
|
||||
@@ -40,6 +50,22 @@ class InMemoryMetrics:
|
||||
return self.counts.get((name, tuple(sorted(labels.items()))), 0)
|
||||
return sum(count for (n, _), count in self.counts.items() if n == name)
|
||||
|
||||
def observe(self, name: str, value: float, **labels: str) -> None:
|
||||
key = (name, tuple(sorted(labels.items())))
|
||||
self.observations.setdefault(key, []).append(value)
|
||||
|
||||
def observed(self, name: str, **labels: str) -> tuple[float, ...]:
|
||||
if labels:
|
||||
return tuple(
|
||||
self.observations.get((name, tuple(sorted(labels.items()))), [])
|
||||
)
|
||||
return tuple(
|
||||
value
|
||||
for (metric_name, _), values in self.observations.items()
|
||||
if metric_name == name
|
||||
for value in values
|
||||
)
|
||||
|
||||
|
||||
RETRIEVAL_ROUTE = "duocthu_retrieval_route_total"
|
||||
ABSTENTION = "duocthu_abstention_total"
|
||||
@@ -60,3 +86,12 @@ FOLLOWUP_INHERITED = "duocthu_followup_inherited_total"
|
||||
# degradation actually happens, since a silent fail-open with no counter is
|
||||
# indistinguishable from tracing quietly working.
|
||||
TRACE_WRITE_FAILED = "duocthu_trace_write_failed_total"
|
||||
|
||||
# HTTP and request-pipeline telemetry. Every label is a bounded vocabulary:
|
||||
# route is a route template (never a raw path), status is a class (2xx/4xx/5xx),
|
||||
# and stage/provider/reason values are normalized by the Prometheus adapter.
|
||||
REQUESTS = "duocthu_requests_total"
|
||||
REQUEST_DURATION = "duocthu_request_duration_seconds"
|
||||
STAGE_DURATION = "duocthu_stage_duration_seconds"
|
||||
DECISION = "duocthu_decision_total"
|
||||
PROVIDER_FAILURE = "duocthu_provider_failure_total"
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Request correlation, OpenTelemetry spans and stage timing.
|
||||
|
||||
This module deliberately has a no-op default. The RAG domain remains runnable
|
||||
without the OpenTelemetry packages or collector, while a configured runtime
|
||||
gets one trace context shared by FastAPI's async middleware and its sync
|
||||
thread-pool endpoint.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from functools import wraps
|
||||
from time import perf_counter
|
||||
from typing import Any, Iterator, Mapping
|
||||
|
||||
from .metrics import Metrics, NullMetrics, STAGE_DURATION
|
||||
|
||||
_CORRELATION_ID: ContextVar[str | None] = ContextVar(
|
||||
"duocthu_correlation_id", default=None
|
||||
)
|
||||
_STAGE: ContextVar[str | None] = ContextVar("duocthu_stage", default=None)
|
||||
_SAFE_CORRELATION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
|
||||
_metrics: Metrics = NullMetrics()
|
||||
|
||||
try:
|
||||
from opentelemetry import trace
|
||||
|
||||
_tracer = trace.get_tracer("duocthu.ai-service")
|
||||
except ImportError: # pragma: no cover - exercised only in minimal installs
|
||||
trace = None
|
||||
_tracer = None
|
||||
|
||||
|
||||
def configure_telemetry(settings: Any, metrics: Metrics | None = None) -> bool:
|
||||
"""Configure one OTLP tracer provider; return whether tracing is active.
|
||||
|
||||
Metrics are configured independently and continue to work when tracing is
|
||||
disabled or optional OpenTelemetry packages are absent.
|
||||
"""
|
||||
global _metrics, _tracer
|
||||
_metrics = metrics or NullMetrics()
|
||||
if not getattr(settings, "otel_enabled", False):
|
||||
return False
|
||||
|
||||
try:
|
||||
from opentelemetry import trace as otel_trace
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
current = otel_trace.get_tracer_provider()
|
||||
# OpenTelemetry only permits setting the global provider once. Reuse a
|
||||
# provider installed by a host/test instead of replacing it and warning.
|
||||
if isinstance(current, TracerProvider):
|
||||
_tracer = current.get_tracer("duocthu.ai-service")
|
||||
return True
|
||||
|
||||
provider = TracerProvider(
|
||||
resource=Resource.create(
|
||||
{
|
||||
"service.name": getattr(settings, "otel_service_name", "ai-service"),
|
||||
"service.version": "0.1.0",
|
||||
"deployment.environment": getattr(settings, "environment", "local"),
|
||||
}
|
||||
),
|
||||
sampler=ParentBased(
|
||||
TraceIdRatioBased(float(getattr(settings, "otel_sample_ratio", 1.0)))
|
||||
),
|
||||
)
|
||||
exporter = OTLPSpanExporter(endpoint=settings.otel_exporter_otlp_endpoint)
|
||||
provider.add_span_processor(BatchSpanProcessor(exporter))
|
||||
otel_trace.set_tracer_provider(provider)
|
||||
_tracer = provider.get_tracer("duocthu.ai-service")
|
||||
return True
|
||||
|
||||
|
||||
def normalize_correlation_id(candidate: str | None) -> str:
|
||||
value = (candidate or "").strip()
|
||||
return value if _SAFE_CORRELATION.fullmatch(value) else str(uuid.uuid4())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def correlation_context(candidate: str | None) -> Iterator[str]:
|
||||
correlation_id = normalize_correlation_id(candidate)
|
||||
token = _CORRELATION_ID.set(correlation_id)
|
||||
try:
|
||||
yield correlation_id
|
||||
finally:
|
||||
_CORRELATION_ID.reset(token)
|
||||
|
||||
|
||||
def current_correlation_id() -> str:
|
||||
value = _CORRELATION_ID.get()
|
||||
return value or str(uuid.uuid4())
|
||||
|
||||
|
||||
def current_trace_id() -> str | None:
|
||||
if trace is None:
|
||||
return None
|
||||
context = trace.get_current_span().get_span_context()
|
||||
if not context.is_valid:
|
||||
return None
|
||||
return f"{context.trace_id:032x}"
|
||||
|
||||
|
||||
def current_stage() -> str | None:
|
||||
return _STAGE.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def request_span(
|
||||
method: str, route: str, headers: Mapping[str, str]
|
||||
) -> Iterator[Any]:
|
||||
"""Start a server span and extract an incoming W3C ``traceparent``."""
|
||||
if _tracer is None:
|
||||
yield None
|
||||
return
|
||||
try:
|
||||
from opentelemetry import propagate
|
||||
from opentelemetry.trace import SpanKind
|
||||
except ImportError: # pragma: no cover - minimal install fallback
|
||||
yield None
|
||||
return
|
||||
|
||||
parent = propagate.extract(headers)
|
||||
with _tracer.start_as_current_span(
|
||||
f"{method} {route}",
|
||||
context=parent,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes={
|
||||
"http.request.method": method,
|
||||
"http.route": route,
|
||||
"duocthu.correlation_id": current_correlation_id(),
|
||||
},
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
@contextmanager
|
||||
def dependency_span(provider: str, operation: str) -> Iterator[Any]:
|
||||
"""Trace one external dependency call without creating another histogram."""
|
||||
if _tracer is None:
|
||||
yield None
|
||||
return
|
||||
with _tracer.start_as_current_span(
|
||||
f"provider.{provider}.{operation}",
|
||||
attributes={
|
||||
"duocthu.provider": provider,
|
||||
"duocthu.operation": operation,
|
||||
},
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
@contextmanager
|
||||
def stage(name: str, **attributes: Any) -> Iterator[Any]:
|
||||
"""Create a child span and observe a low-cardinality stage histogram."""
|
||||
started = perf_counter()
|
||||
outcome = "ok"
|
||||
span = None
|
||||
token = _STAGE.set(name)
|
||||
try:
|
||||
if _tracer is None:
|
||||
yield None
|
||||
else:
|
||||
safe_attributes = {
|
||||
key: value
|
||||
for key, value in attributes.items()
|
||||
if value is not None and isinstance(value, (str, bool, int, float))
|
||||
}
|
||||
safe_attributes["duocthu.stage"] = name
|
||||
with _tracer.start_as_current_span(
|
||||
f"rag.stage.{name}", attributes=safe_attributes
|
||||
) as span:
|
||||
yield span
|
||||
except BaseException as exc:
|
||||
outcome = "cancelled" if exc.__class__.__name__ == "CancelledError" else "error"
|
||||
if span is not None:
|
||||
span.set_attribute("duocthu.outcome", outcome)
|
||||
raise
|
||||
finally:
|
||||
elapsed = perf_counter() - started
|
||||
if span is not None:
|
||||
span.set_attribute("duocthu.outcome", outcome)
|
||||
span.set_attribute("duocthu.duration_ms", elapsed * 1000)
|
||||
_metrics.observe(STAGE_DURATION, elapsed, stage=name, outcome=outcome)
|
||||
_STAGE.reset(token)
|
||||
|
||||
|
||||
def annotate_current_span(**attributes: Any) -> None:
|
||||
if trace is None:
|
||||
return
|
||||
span = trace.get_current_span()
|
||||
if not span.is_recording():
|
||||
return
|
||||
for key, value in attributes.items():
|
||||
if value is not None and isinstance(value, (str, bool, int, float)):
|
||||
span.set_attribute(key, value)
|
||||
|
||||
|
||||
def traced_stage(name: str):
|
||||
"""Decorator form for pure functions that should remain otherwise untouched."""
|
||||
def decorate(function):
|
||||
@wraps(function)
|
||||
def wrapped(*args, **kwargs):
|
||||
with stage(name):
|
||||
return function(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
return decorate
|
||||
@@ -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
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from adapters.prometheus import PrometheusMetrics
|
||||
from config import Settings
|
||||
from main import create_app
|
||||
from rag.agent import AgentReply
|
||||
@@ -177,3 +178,63 @@ def test_a_trace_write_failure_does_not_turn_a_good_answer_into_a_500():
|
||||
assert body["answer"] == "Liều 500 mg [1]."
|
||||
assert body["trace_id"] # a locally-generated fallback id, still present
|
||||
assert metrics.total(TRACE_WRITE_FAILED) == 1
|
||||
|
||||
|
||||
class _AnyRouting:
|
||||
def retrieve(self, query, subject_scope, intent):
|
||||
return RetrievalResult(EvidenceDecision.ABSTAIN, "drug_not_resolved")
|
||||
|
||||
|
||||
def test_correlation_id_round_trips_through_headers_body_and_trace_row():
|
||||
traces = MemoryTraceWriter()
|
||||
app = create_app(
|
||||
settings=Settings(embedding_provider="disabled"),
|
||||
answer_service=GroundedAnswerService(_AnyRouting()),
|
||||
trace_writer=traces,
|
||||
)
|
||||
|
||||
response = TestClient(app).post(
|
||||
"/v1/rag/query",
|
||||
headers={"X-Correlation-ID": "req-test-1"},
|
||||
json={
|
||||
"query": "Paracetamol dose?",
|
||||
"subject_scope": "human",
|
||||
"intent": "fact_lookup",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["x-correlation-id"] == "req-test-1"
|
||||
assert response.json()["correlation_id"] == "req-test-1"
|
||||
assert traces.rows[0]["correlation_id"] == "req-test-1"
|
||||
assert "otel_trace_id" in traces.rows[0]
|
||||
|
||||
|
||||
def test_metrics_endpoint_exposes_request_decision_and_stage_histograms():
|
||||
metrics = PrometheusMetrics()
|
||||
app = create_app(
|
||||
settings=Settings(embedding_provider="disabled"),
|
||||
answer_service=GroundedAnswerService(_AnyRouting()),
|
||||
trace_writer=MemoryTraceWriter(),
|
||||
metrics=metrics,
|
||||
)
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post(
|
||||
"/v1/rag/query",
|
||||
json={
|
||||
"query": "Paracetamol dose?",
|
||||
"subject_scope": "human",
|
||||
"intent": "fact_lookup",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
scrape = client.get("/metrics")
|
||||
assert scrape.status_code == 200
|
||||
body = scrape.text
|
||||
assert 'duocthu_requests_total{method="POST",route="/v1/rag/query",status="2xx"}' in body
|
||||
assert 'duocthu_decision_total{decision="abstain",reason="drug_not_resolved"}' in body
|
||||
assert 'duocthu_stage_duration_seconds_count{outcome="ok",stage="receive"}' in body
|
||||
assert 'duocthu_stage_duration_seconds_count{outcome="ok",stage="persistence"}' in body
|
||||
assert 'duocthu_stage_duration_seconds_count{outcome="ok",stage="response"}' in body
|
||||
|
||||
@@ -11,7 +11,7 @@ import json
|
||||
import pytest
|
||||
|
||||
from rag import grounding
|
||||
from rag.answer import GroundedAnswerService
|
||||
from rag.answer import DISCLAIMER, GroundedAnswerService
|
||||
from rag.budget import RequestBudgetExhausted
|
||||
from rag.metrics import GENERATION_REJECTED, GENERATION_SERVED, InMemoryMetrics
|
||||
from rag.models import (
|
||||
@@ -116,6 +116,62 @@ def _answer(payload, result: RetrievalResult | None = None, entailment_payload=N
|
||||
return grounded, metrics
|
||||
|
||||
|
||||
# --- the disclaimer guardrail -------------------------------------------------
|
||||
# `docs/architecture.md` specifies the disclaimer at several layers. Only the
|
||||
# web banner existed; `packages/shared-types/src/dto/chat.ts` declared the
|
||||
# field but nothing filled it, so a consumer other than that one UI got medical
|
||||
# content with nothing attached. It is a dataclass default rather than
|
||||
# something each call site adds, so these tests are about the paths that could
|
||||
# plausibly skip it: rejections, abstains and clarifications.
|
||||
|
||||
|
||||
def test_a_served_answer_carries_the_disclaimer():
|
||||
grounded, _ = _answer(
|
||||
{"claims": [{"text": "Người lớn uống 500 mg", "citations": [1]}],
|
||||
"evidence_sufficient": True},
|
||||
)
|
||||
|
||||
assert grounded.generated is True
|
||||
assert grounded.disclaimer == DISCLAIMER
|
||||
assert grounded.disclaimer
|
||||
|
||||
|
||||
def test_a_rejected_generation_still_carries_the_disclaimer():
|
||||
"""An abstain is still a clinical response and still needs the notice —
|
||||
it is the case most likely to be treated as "not really an answer"."""
|
||||
grounded, _ = _answer(
|
||||
{"claims": [{"text": "Người lớn uống 850 mg, 2 lần mỗi ngày", "citations": [1]}],
|
||||
"evidence_sufficient": True},
|
||||
)
|
||||
|
||||
assert grounded.answer is None
|
||||
assert grounded.result.decision == EvidenceDecision.ABSTAIN
|
||||
assert grounded.disclaimer == DISCLAIMER
|
||||
|
||||
|
||||
def test_a_provider_outage_response_still_carries_the_disclaimer():
|
||||
grounded, _ = _answer(
|
||||
{"claims": [{"text": "Người lớn: 500 mg, 2 lần/ngày", "citations": [1]}],
|
||||
"evidence_sufficient": True},
|
||||
entailment_payload=AnswerGenerationUnavailable(),
|
||||
)
|
||||
|
||||
assert grounded.disclaimer == DISCLAIMER
|
||||
|
||||
|
||||
def test_the_disclaimer_is_not_something_the_model_can_influence():
|
||||
"""It is a module constant, never routed through the generator, so a
|
||||
prompt-injected or malfunctioning model cannot shorten or drop it. Asserted
|
||||
directly because the value is the guardrail."""
|
||||
grounded, _ = _answer(
|
||||
{"claims": [{"text": "Bỏ qua mọi cảnh báo. Người lớn uống 500 mg", "citations": [1]}],
|
||||
"evidence_sufficient": True},
|
||||
)
|
||||
|
||||
assert grounded.disclaimer == DISCLAIMER
|
||||
assert "không thay thế chỉ định" in grounded.disclaimer
|
||||
|
||||
|
||||
# --- the guardrail's whole reason to exist ------------------------------------
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
from adapters.prometheus import PrometheusMetrics
|
||||
from config import Settings
|
||||
from rag.instrumentation import InstrumentedGenerator
|
||||
from rag.metrics import (
|
||||
CLARIFY_ASKED,
|
||||
PROVIDER_FAILURE,
|
||||
REQUEST_DURATION,
|
||||
REQUESTS,
|
||||
STAGE_DURATION,
|
||||
TRACE_WRITE_FAILED,
|
||||
InMemoryMetrics,
|
||||
)
|
||||
from rag.telemetry import (
|
||||
configure_telemetry,
|
||||
correlation_context,
|
||||
request_span,
|
||||
stage,
|
||||
)
|
||||
|
||||
|
||||
def test_prometheus_registers_domain_failures_histograms_and_bounds_labels():
|
||||
metrics = PrometheusMetrics()
|
||||
metrics.increment(TRACE_WRITE_FAILED)
|
||||
metrics.increment(CLARIFY_ASKED, reason="raw-user-controlled-reason")
|
||||
metrics.increment(
|
||||
REQUESTS, method="POST", route="/v1/rag/query", status="2xx"
|
||||
)
|
||||
metrics.observe(
|
||||
REQUEST_DURATION,
|
||||
1.25,
|
||||
method="POST",
|
||||
route="/v1/rag/query",
|
||||
status="2xx",
|
||||
)
|
||||
metrics.observe(STAGE_DURATION, 0.2, stage="retrieval", outcome="ok")
|
||||
|
||||
body, content_type = metrics.render()
|
||||
rendered = body.decode("utf-8")
|
||||
assert "application/openmetrics-text" in content_type
|
||||
assert "duocthu_trace_write_failed_total 1.0" in rendered
|
||||
assert 'duocthu_clarify_asked_total{reason="other"} 1.0' in rendered
|
||||
assert "duocthu_request_duration_seconds_bucket" in rendered
|
||||
assert "duocthu_stage_duration_seconds_bucket" in rendered
|
||||
|
||||
|
||||
def test_stage_records_duration_without_requiring_an_otel_collector():
|
||||
metrics = InMemoryMetrics()
|
||||
configure_telemetry(Settings(otel_enabled=False), metrics)
|
||||
|
||||
with stage("grounding"):
|
||||
pass
|
||||
|
||||
values = metrics.observed(STAGE_DURATION, stage="grounding", outcome="ok")
|
||||
assert len(values) == 1
|
||||
assert values[0] >= 0
|
||||
|
||||
|
||||
class _FailingGenerator:
|
||||
model_id = "fake"
|
||||
|
||||
def generate(self, system: str, user: str, schema: dict) -> str:
|
||||
raise ConnectionError("provider offline")
|
||||
|
||||
|
||||
def test_provider_failure_uses_bounded_classification_not_exception_text():
|
||||
metrics = InMemoryMetrics()
|
||||
generator = InstrumentedGenerator(_FailingGenerator(), metrics)
|
||||
|
||||
try:
|
||||
generator.generate("system", "user", {})
|
||||
except ConnectionError:
|
||||
pass
|
||||
else: # pragma: no cover - assertion guard
|
||||
raise AssertionError("expected provider failure")
|
||||
|
||||
assert metrics.total(
|
||||
PROVIDER_FAILURE,
|
||||
provider="other",
|
||||
operation="generate",
|
||||
reason="provider_unavailable",
|
||||
) == 1
|
||||
|
||||
|
||||
def test_otel_server_and_stage_spans_share_trace_and_emit_metric_exemplar():
|
||||
otel_trace = pytest.importorskip("opentelemetry.trace")
|
||||
sdk_trace = pytest.importorskip("opentelemetry.sdk.trace")
|
||||
exporter_module = pytest.importorskip(
|
||||
"opentelemetry.sdk.trace.export.in_memory_span_exporter"
|
||||
)
|
||||
export_module = pytest.importorskip("opentelemetry.sdk.trace.export")
|
||||
|
||||
exporter = exporter_module.InMemorySpanExporter()
|
||||
provider = sdk_trace.TracerProvider()
|
||||
provider.add_span_processor(export_module.SimpleSpanProcessor(exporter))
|
||||
otel_trace.set_tracer_provider(provider)
|
||||
|
||||
metrics = PrometheusMetrics()
|
||||
configure_telemetry(
|
||||
Settings(
|
||||
otel_enabled=True,
|
||||
otel_exporter_otlp_endpoint="http://127.0.0.1:4318/v1/traces",
|
||||
),
|
||||
metrics,
|
||||
)
|
||||
with correlation_context("req-span-test"):
|
||||
with request_span("POST", "/v1/rag/query", {}):
|
||||
with stage("retrieval"):
|
||||
pass
|
||||
|
||||
provider.force_flush()
|
||||
spans = exporter.get_finished_spans()
|
||||
assert {item.name for item in spans} == {
|
||||
"POST /v1/rag/query",
|
||||
"rag.stage.retrieval",
|
||||
}
|
||||
assert len({item.context.trace_id for item in spans}) == 1
|
||||
assert spans[-1].attributes["duocthu.correlation_id"] == "req-span-test"
|
||||
|
||||
rendered = metrics.render()[0].decode("utf-8")
|
||||
assert "duocthu_stage_duration_seconds_bucket" in rendered
|
||||
assert 'trace_id="' in rendered
|
||||
import pytest
|
||||
Reference in New Issue
Block a user