app/middleware/metrics.py

"""Author RAG Chatbot SaaS — Prometheus Metrics Middleware.

R-099: Prometheus /metrics endpoint MUST be configured.
R-102: Per-step pipeline latency MUST be tracked.

Exposes request counts, latencies, and business metrics via /metrics endpoint.
"""

import time

import structlog
from starlette.types import ASGIApp, Receive, Scope, Send

logger = structlog.get_logger(__name__)

try:
    from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST

    # ─── HTTP Metrics ─────────────────────────────────────────────────────────
    REQUEST_COUNT = Counter(
        "http_requests_total",
        "Total HTTP requests",
        ["method", "path", "status"],
    )
    REQUEST_LATENCY = Histogram(
        "http_request_duration_seconds",
        "HTTP request latency in seconds",
        ["method", "path"],
        buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0],
    )

    # ─── Business Metrics ─────────────────────────────────────────────────────
    CHAT_TURNS = Counter(
        "chat_turns_total",
        "Total chat turns processed",
        ["author_id"],
    )
    TOKEN_USAGE = Counter(
        "tokens_used_total",
        "Total LLM tokens consumed",
        ["author_id"],
    )
    PIPELINE_STEP_LATENCY = Histogram(
        "rag_pipeline_step_duration_ms",
        "RAG pipeline per-step latency in milliseconds",
        ["step"],
        buckets=[1, 5, 10, 25, 50, 100, 250, 500, 1000, 5000],
    )
    INGESTION_COUNT = Counter(
        "document_ingestions_total",
        "Total document ingestions",
        ["status"],
    )

    PROMETHEUS_AVAILABLE = True

except ImportError:
    PROMETHEUS_AVAILABLE = False
    logger.warning("prometheus_client not installed — /metrics endpoint disabled")


class PrometheusMiddleware:
    """ASGI middleware that tracks request counts and latencies."""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        if not PROMETHEUS_AVAILABLE:
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        method = scope.get("method", "GET")

        # Skip metrics endpoint itself to avoid recursion
        if path == "/metrics":
            await self.app(scope, receive, send)
            return

        # Normalize paths to avoid cardinality explosion
        normalized = _normalize_path(path)

        start = time.perf_counter()
        status_code = 500  # Default in case of crash

        async def send_wrapper(message):
            nonlocal status_code
            if message["type"] == "http.response.start":
                status_code = message.get("status", 200)
            await send(message)

        try:
            await self.app(scope, receive, send_wrapper)
        finally:
            duration = time.perf_counter() - start
            REQUEST_COUNT.labels(method=method, path=normalized, status=str(status_code)).inc()
            REQUEST_LATENCY.labels(method=method, path=normalized).observe(duration)


def _normalize_path(path: str) -> str:
    """Normalize URL paths to prevent label cardinality explosion.

    Replaces UUIDs and dynamic segments with placeholders.
    """
    import re
    # Replace UUIDs
    path = re.sub(r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}', '{id}', path)
    # Replace plain hex IDs (32+ chars)
    path = re.sub(r'/[0-9a-f]{32,}', '/{id}', path)
    return path


def get_metrics_response():
    """Generate Prometheus metrics response for /metrics endpoint."""
    if not PROMETHEUS_AVAILABLE:
        return b"# prometheus_client not installed\n", "text/plain"
    return generate_latest(), CONTENT_TYPE_LATEST