didi-lot1-ai/ai_platform/shared/observability.py

"""
Prometheus metrics + OpenTelemetry tracing setup for DiDi AI Platform modules.

Usage in FastAPI app.py:

    from shared.observability import init_observability, mount_metrics

    app = FastAPI(lifespan=...)
    init_observability(app, service_name="llm-inference", version="0.1.0")
    mount_metrics(app)

This:
  - Mounts /metrics endpoint (Prometheus scrape target)
  - Instruments FastAPI with OpenTelemetry (auto-traces requests)
  - Sets up resource attributes (service.name, service.version, cluster)
  - Configures OTLP exporter pointing to OTEL_EXPORTER_OTLP_ENDPOINT

Optional: import custom metrics:

    from shared.observability import (
        llm_calls_total, llm_tokens_total, inference_duration_seconds
    )
    llm_calls_total.labels(provider="vllm", model="qwen3.5", status="ok").inc()

Required pip deps (add to each module's pyproject.toml):
    prometheus-client>=0.20.0
    opentelemetry-api>=1.25.0
    opentelemetry-sdk>=1.25.0
    opentelemetry-exporter-otlp-proto-grpc>=1.25.0
    opentelemetry-instrumentation-fastapi>=0.46b0
    opentelemetry-instrumentation-httpx>=0.46b0
"""
from __future__ import annotations

import os
import logging
from typing import Optional

from prometheus_client import (
    Counter,
    Histogram,
    Gauge,
    CollectorRegistry,
    generate_latest,
    CONTENT_TYPE_LATEST,
    REGISTRY,
)

logger = logging.getLogger(__name__)

_initialized = False


def init_observability(
    app,
    service_name: str,
    version: str = "0.1.0",
    enable_tracing: bool = True,
) -> None:
    """Initialize OpenTelemetry tracing + auto-instrument FastAPI.

    Safe to call multiple times — no-op on subsequent calls.
    Fail-open: if OTel deps missing, logs warning and continues without tracing.
    """
    global _initialized
    if _initialized:
        return
    _initialized = True

    otel_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
    if not enable_tracing or not otel_endpoint:
        logger.info(
            "[observability] tracing disabled (OTEL_EXPORTER_OTLP_ENDPOINT not set or disabled)"
        )
        return

    try:
        from opentelemetry import trace
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor
        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

        resource = Resource(
            attributes={
                "service.name": service_name,
                "service.version": version,
                "cluster": os.environ.get("OTEL_CLUSTER", "didi-prod"),
                "environment": os.environ.get("OTEL_ENVIRONMENT", "production"),
            }
        )
        provider = TracerProvider(resource=resource)
        provider.add_span_processor(
            BatchSpanProcessor(OTLPSpanExporter(endpoint=otel_endpoint, insecure=True))
        )
        trace.set_tracer_provider(provider)

        FastAPIInstrumentor.instrument_app(app)

        # Optional: instrument httpx clients automatically
        try:
            from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor

            HTTPXClientInstrumentor().instrument()
        except ImportError:
            pass

        logger.info(f"[observability] OTel tracing initialized for {service_name}")
    except ImportError as e:
        logger.warning(
            f"[observability] OTel deps missing — tracing disabled. {e}. "
            f"Install: pip install opentelemetry-api opentelemetry-sdk "
            f"opentelemetry-exporter-otlp-proto-grpc opentelemetry-instrumentation-fastapi"
        )


def mount_metrics(app, registry: Optional[CollectorRegistry] = None) -> None:
    """Mount /metrics endpoint on FastAPI app for Prometheus scraping."""
    from fastapi import Response

    reg = registry or REGISTRY

    @app.get("/metrics", include_in_schema=False)
    async def _metrics():
        return Response(content=generate_latest(reg), media_type=CONTENT_TYPE_LATEST)


# ============================================================================
# Standard DiDi AI Platform metrics
# ============================================================================

llm_calls_total = Counter(
    "didi_llm_calls_total",
    "Total LLM calls made (per provider+model+status)",
    labelnames=["provider", "model", "status"],
)

llm_tokens_total = Counter(
    "didi_llm_tokens_total",
    "Total LLM tokens consumed",
    labelnames=["provider", "model", "type"],  # type=prompt|completion
)

inference_duration_seconds = Histogram(
    "didi_inference_duration_seconds",
    "Inference duration in seconds (per backend+model)",
    labelnames=["backend", "model"],
    buckets=(0.1, 0.5, 1, 2, 5, 10, 30, 60, 120),
)

http_requests_total = Counter(
    "didi_ai_http_requests_total",
    "HTTP requests received by AI platform service",
    labelnames=["method", "route", "status"],
)

active_backends = Gauge(
    "didi_active_backends",
    "Number of active/healthy backends per type",
    labelnames=["backend_type"],
)

cache_hits_total = Counter(
    "didi_ai_cache_hits_total",
    "Cache hits per cache type",
    labelnames=["cache"],
)

cache_misses_total = Counter(
    "didi_ai_cache_misses_total",
    "Cache misses per cache type",
    labelnames=["cache"],
)

# Brain-specific (only used by didi_brain module)
brain_atoms_total = Gauge(
    "didi_brain_atoms_total",
    "Total brain analysis atoms by tier",
    labelnames=["tier"],
)

brain_facts_total = Gauge(
    "didi_brain_facts_total",
    "Total brain facts by current_truth",
    labelnames=["truth"],
)

brain_judge_decisions_total = Counter(
    "didi_brain_judge_decisions_total",
    "NLI judge decisions (KEEP/INVALIDATE/RECHECK)",
    labelnames=["decision"],
)