duocthu/apps/ai-service/adapters/prometheus.py

"""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",
            "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