GenAI Monitoring & Observability

from opentelemetry import trace, metrics
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.metrics import MeterProvider
from dataclasses import dataclass
from typing import Dict, Any, Optional
import time
import asyncio

# Initialize OpenTelemetry
tracer = trace.get_tracer("llm.service")
meter = metrics.get_meter("llm.metrics")

# Define custom metrics
request_counter = meter.create_counter(
    "llm_requests_total",
    description="Total number of LLM requests"
)

token_histogram = meter.create_histogram(
    "llm_tokens_used",
    description="Token usage distribution"
)

latency_histogram = meter.create_histogram(
    "llm_request_duration_seconds",
    description="LLM request duration"
)

@dataclass
class LLMMetrics:
    model: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    duration: float
    cost: float
    status: str
    user_id: str

class LLMObservability:
    def __init__(self):
        self.active_requests = {}
        
    async def trace_llm_request(self, request_id: str, model: str, prompt: str, **kwargs):
        """Trace LLM request with comprehensive metrics"""
        
        with tracer.start_as_current_span("llm_request") as span:
            # Set span attributes
            span.set_attribute("llm.model", model)
            span.set_attribute("llm.request_id", request_id)
            span.set_attribute("llm.prompt_length", len(prompt))
            
            start_time = time.time()
            
            try:
                # Make LLM request
                response = await self._make_llm_request(model, prompt, **kwargs)
                
                # Calculate metrics
                duration = time.time() - start_time
                metrics = LLMMetrics(
                    model=model,
                    prompt_tokens=response.get('prompt_tokens', 0),
                    completion_tokens=response.get('completion_tokens', 0),
                    total_tokens=response.get('total_tokens', 0),
                    duration=duration,
                    cost=self._calculate_cost(response),
                    status='success',
                    user_id=kwargs.get('user_id', 'anonymous')
                )
                
                # Record metrics
                self._record_metrics(metrics)
                span.set_attribute("llm.status", "success")
                
                return response
                
            except Exception as e:
                # Record error metrics
                duration = time.time() - start_time
                error_metrics = LLMMetrics(
                    model=model, prompt_tokens=0, completion_tokens=0,
                    total_tokens=0, duration=duration, cost=0.0,
                    status='error', user_id=kwargs.get('user_id', 'anonymous')
                )
                
                self._record_metrics(error_metrics)
                span.set_attribute("llm.status", "error")
                span.set_attribute("llm.error", str(e))
                raise
                
    def _record_metrics(self, metrics: LLMMetrics):
        """Record metrics to monitoring system"""
        # Increment request counter
        request_counter.add(1, {
            "model": metrics.model,
            "status": metrics.status,
            "user_id": metrics.user_id
        })
        
        # Record token usage
        token_histogram.record(metrics.total_tokens, {
            "model": metrics.model,
            "type": "total"
        })
        
        # Record latency
        latency_histogram.record(metrics.duration, {
            "model": metrics.model,
            "status": metrics.status
        })