Production AI Monitoring & Observability

# BEFORE (FATAL: High-cardinality explosion in metrics)
# ai_requests_counter.labels(model="gpt-4o", user_id=user.id, session_id=session.id).inc()

# AFTER (CORRECT: Keep low-cardinality dimensions in Prometheus)
from prometheus_client import Counter
from opentelemetry import trace

ai_requests_counter = Counter(
    "ai_inference_requests_total",
    "Total AI inference calls",
    ["model_name", "status_code", "provider"]
)

# Record metric with bounded label values
ai_requests_counter.labels(
    model_name="gpt-4o",
    status_code="200",
    provider="openai"
).inc()

# Store user_id and session_id inside OpenTelemetry Trace Span attributes instead!
current_span = trace.get_current_span()
if current_span.is_recording():
    current_span.set_attribute("user.id", user.id)
    current_span.set_attribute("session.id", session.id)

# Production Resilient LLM Client with Telemetry & Circuit Breaking
import time
from opentelemetry import trace
from prometheus_client import Counter, Histogram

LLM_UPSTREAM_REQUESTS = Counter(
    "ai_upstream_requests_total",
    "Upstream LLM requests by provider and status",
    ["provider", "model", "http_status"]
)
LLM_RETRY_COUNT = Counter(
    "ai_upstream_retries_total",
    "Number of retry attempts triggered",
    ["provider", "reason"]
)

async def call_llm_with_telemetry(prompt: str, model: str, attempt: int = 1):
    tracer = trace.get_tracer("ai.llm.client")
    with tracer.start_as_current_span("upstream_llm_call") as span:
        span.set_attribute("gen_ai.request.model", model)
        span.set_attribute("retry.attempt", attempt)
        
        try:
            resp = await client.chat.completions.create(model=model, messages=[...])
            LLM_UPSTREAM_REQUESTS.labels(provider="openai", model=model, http_status="200").inc()
            return resp
        except RateLimitError as e:
            LLM_UPSTREAM_REQUESTS.labels(provider="openai", model=model, http_status="429").inc()
            LLM_RETRY_COUNT.labels(provider="openai", reason="rate_limit_429").inc()
            span.record_exception(e)
            span.set_status(trace.Status(trace.StatusCode.ERROR, "Rate limited"))
            
            # Fail fast to fallback model if attempt > 2 instead of 45s backoff stall
            if attempt >= 2:
                return await call_fallback_local_vllm(prompt)

# ANTI-PATTERN: Manual span without exception safety
# span = tracer.start_span("streaming_chat")
# for chunk in stream: yield chunk
# span.end() # NEVER REACHED IF CLIENT DISCONNECTS!

# PRODUCTION FIX: Context manager guarantees span.end() execution
from fastapi import Request
from fastapi.responses import StreamingResponse
from opentelemetry import trace

tracer = trace.get_tracer("ai.chat.service")

@app.post("/v1/chat/stream")
async def chat_stream(request: Request, payload: ChatPayload):
    async def token_generator():
        # Context manager handles span lifecycle across async cancellation
        with tracer.start_as_current_span("llm_stream_generation") as span:
            span.set_attribute("gen_ai.system", "openai")
            span.set_attribute("gen_ai.request.model", payload.model)
            token_count = 0
            try:
                async for chunk in model_stream(payload.prompt):
                    token_count += 1
                    yield f"data: {chunk}\n\n"
            except Exception as exc:
                span.record_exception(exc)
                span.set_status(trace.Status(trace.StatusCode.ERROR))
                raise
            finally:
                span.set_attribute("gen_ai.usage.output_tokens", token_count)
                # span context automatically closes here, preventing memory leaks!
                
    return StreamingResponse(token_generator(), media_type="text/event-stream")

# Prometheus Histogram bucket configuration in Python
from prometheus_client import Histogram

VECTOR_SEARCH_DURATION = Histogram(
    "rag_vector_search_duration_seconds",
    "Latency of Qdrant HNSW similarity search",
    ["collection_name", "index_type"],
    # Explicit custom exponential buckets up to 25 seconds
    buckets=[0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 25.0]
)

# PromQL Alert Rule for Prometheus Alertmanager (alerts.yaml)
"""
groups:
  - name: rag_alerts
    rules:
      - alert: RAGVectorSearchTailLatencySpike
        expr: histogram_quantile(0.99, sum(rate(rag_vector_search_duration_seconds_bucket[5m])) by (le)) > 3.0
        for: 3m
        labels:
          severity: warning
        annotations:
          summary: "Vector DB p99 search latency > 3s (currently {{ $value }}s)"
          runbook_url: "https://wiki.corp/ops/rag-tuning"
"""

# Production Structured Logging Sanitizer Middleware
import re
import logging
import json

class PIIMaskingFormatter(logging.Formatter):
    CC_REGEX = re.compile(r'\b(?:\d[ -]*?){13,16}\b')
    SSN_REGEX = re.compile(r'\b\d{3}-\d{2}-\d{4}\b')
    BEARER_REGEX = re.compile(r'Bearer\s+[A-Za-z0-9\-\._~\+\/]+=*', re.IGNORECASE)
    
    def format(self, record: logging.LogRecord) -> str:
        original = super().format(record)
        # Redact credit cards
        sanitized = self.CC_REGEX.sub("[REDACTED_CC]", original)
        # Redact SSNs
        sanitized = self.SSN_REGEX.sub("[REDACTED_SSN]", sanitized)
        # Redact Authorization header tokens
        sanitized = self.BEARER_REGEX.sub("Bearer [REDACTED_TOKEN]", sanitized)
        return sanitized

# Setup safe logger
logger = logging.getLogger("ai.safe.service")
handler = logging.StreamHandler()
handler.setFormatter(PIIMaskingFormatter('{"time":"%(asctime)s", "level":"%(levelname)s", "msg":"%(message)s"}'))
logger.addHandler(handler)
logger.setLevel(logging.INFO)

# Agent Observability & Circuit Breaking Loop Guard
from opentelemetry import trace
from prometheus_client import Counter, Histogram

AGENT_TOOL_CALLS = Counter(
    "ai_agent_tool_calls_total", 
    "Total tool invocations by agent", 
    ["agent_name", "tool_name", "status"]
)
AGENT_ITERATION_DEPTH = Histogram(
    "ai_agent_iterations_per_request",
    "Number of iterations before agent resolution",
    buckets=[1, 2, 3, 4, 5, 8, 10, 15]
)

MAX_ALLOWED_ITERATIONS = 6

async def run_agent_loop(user_query: str):
    tracer = trace.get_tracer("ai.agent")
    with tracer.start_as_current_span("agent_execution_cycle") as span:
        iteration = 0
        while iteration < MAX_ALLOWED_ITERATIONS:
            iteration += 1
            span.add_event(f"Iteration_{iteration}_start")
            
            tool_decision = await llm_decide_action(user_query)
            if tool_decision.is_final_answer:
                AGENT_ITERATION_DEPTH.observe(iteration)
                span.set_attribute("agent.total_iterations", iteration)
                return tool_decision.answer
                
            # Track tool execution
            AGENT_TOOL_CALLS.labels(
                agent_name="support_v1", 
                tool_name=tool_decision.tool_name, 
                status="invoked"
            ).inc()
            await execute_tool(tool_decision.tool_name, tool_decision.args)
            
        # Circuit breaker triggered!
        span.set_status(trace.Status(trace.StatusCode.ERROR, "Max iterations exceeded"))
        span.set_attribute("agent.loop_broken", True)
        return "I apologize, but I was unable to resolve your request in a timely manner."

# otel-collector-config.yaml (CORRECT Sampling Pipeline)
processors:
  tail_sampling:
    decision_wait: 10s
    num_traces: 10000
    expected_new_traces_per_sec: 2000
    policies:
      # Policy 1: Always keep 100% of traces that contain errors
      - name: keep_all_errors
        type: status_code
        status_code: { status_codes: [ ERROR ] }
        
      # Policy 2: Keep 100% of slow traces (> 3 seconds)
      - name: keep_slow_traces
        type: latency
        latency: { threshold_ms: 3000 }
        
      # Policy 3: Sample only 5% of healthy, fast traces to save disk space
      - name: sample_normal_traffic
        type: probabilistic
        probabilistic: { sampling_percentage: 5.0 }

service:
  pipelines:
    traces:
      receivers: [otlp]
      processors: [memory_limiter, tail_sampling, batch]
      exporters: [otlp/tempo]

# Prometheus Scrape Config for AI Model Serving (prometheus.yml)
scrape_configs:
  # 1. Scrape vLLM Native Metrics (KV cache, queue depths)
  - job_name: 'vllm-model-server'
    metrics_path: '/metrics'
    static_configs:
      - targets: ['vllm-service.inference.svc:8000']
    metric_relabel_configs:
      - source_labels: [__name__]
        regex: '(vllm:.*)'
        action: keep

  # 2. Scrape NVIDIA DCGM Exporter for Hardware Accelerators
  - job_name: 'nvidia-dcgm'
    static_configs:
      - targets: ['dcgm-exporter.monitoring.svc:9400']

# Critical Alert in Alertmanager:
# expr: vllm:gpu_cache_usage_factor > 0.95 and vllm:num_requests_waiting > 10
# for: 2m
# summary: "vLLM KV Cache exhausted! Requests are stalling in scheduler queue."