238 lines
7.8 KiB
Python
238 lines
7.8 KiB
Python
"""Prometheus export for bounded domain and request-path telemetry.
|
|
|
|
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
|
|
|
|
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,
|
|
)
|
|
|
|
_COUNTER_LABELS: dict[str, tuple[str, ...]] = {
|
|
ABSTENTION: ("reason",),
|
|
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 bounded domain reason.",
|
|
ANSWER_EXTRACTIVE: "Answers served as verbatim source text.",
|
|
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",
|
|
"/v1/rag/feedback",
|
|
"/v1/rag/history",
|
|
"/v1/rag/sections",
|
|
"/v1/rag/section-text",
|
|
"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 an isolated Prometheus registry."""
|
|
|
|
def __init__(self, registry: Any | None = None) -> None:
|
|
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 _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
|
|
def registry(self) -> Any:
|
|
return self._registry
|
|
|
|
def increment(self, name: str, **labels: str) -> None:
|
|
counter = self._counters.get(name)
|
|
expected = _COUNTER_LABELS.get(name)
|
|
if counter is None or expected is None or set(labels) != set(expected):
|
|
return
|
|
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
|
|
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]:
|
|
# 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
|