Observability for AI Agents: Monitoring the Black Box

from dataclasses import dataclass
from typing import Optional, Dict, List
import time
import asyncio

@dataclass
class AgentMetrics:
    request_id: str
    user_id: str
    agent_type: str
    
    # Timing metrics
    total_duration: float
    llm_duration: float
    context_retrieval_duration: float
    tool_execution_duration: float
    
    # Token usage
    input_tokens: int
    output_tokens: int
    context_tokens: int
    
    # Cost tracking
    total_cost: float
    model_cost: float
    
    # Quality indicators
    success: bool
    error_type: Optional[str]
    confidence_score: Optional[float]

class AgentMetricsCollector:
    def __init__(self):
        self.metrics_buffer = []
        self.flush_interval = 60  # seconds
        self.start_flush_loop()
    
    async def track_agent_execution(self, agent_func, **kwargs):
        """Decorator-style metrics collection"""
        start_time = time.time()
        request_id = kwargs.get('request_id', generate_id())
        
        try:
            # Execute agent function with context tracking
            with TokenUsageTracker() as token_tracker:
                result = await agent_func(**kwargs)
            
            # Collect success metrics
            metrics = AgentMetrics(
                request_id=request_id,
                user_id=kwargs.get('user_id'),
                agent_type=agent_func.__name__,
                total_duration=time.time() - start_time,
                llm_duration=token_tracker.llm_duration,
                input_tokens=token_tracker.input_tokens,
                output_tokens=token_tracker.output_tokens,
                total_cost=token_tracker.total_cost,
                success=True,
                confidence_score=result.get('confidence')
            )
            
        except Exception as e:
            # Collect error metrics
            metrics = AgentMetrics(
                request_id=request_id,
                user_id=kwargs.get('user_id'),
                agent_type=agent_func.__name__,
                total_duration=time.time() - start_time,
                success=False,
                error_type=type(e).__name__
            )
            raise
        
        finally:
            await self.record_metrics(metrics)
        
        return result

class QualityMetricsCollector:
    def __init__(self):
        self.quality_evaluator = ResponseQualityEvaluator()
        self.feedback_collector = UserFeedbackCollector()
    
    async def evaluate_response_quality(self, query: str, response: str, 
                                      context: str) -> Dict[str, float]:
        """Multi-dimensional quality assessment"""
        
        # Automated quality checks
        relevance_score = await self._measure_relevance(query, response)
        coherence_score = await self._measure_coherence(response)
        factual_score = await self._check_factual_accuracy(response, context)
        
        # Semantic similarity to ground truth (if available)
        similarity_score = await self._compare_to_ground_truth(
            query, response
        )
        
        quality_metrics = {
            'relevance': relevance_score,
            'coherence': coherence_score,
            'factual_accuracy': factual_score,
            'similarity_to_ground_truth': similarity_score,
            'overall_quality': self._calculate_overall_score([
                relevance_score, coherence_score, factual_score
            ])
        }
        
        return quality_metrics
    
    async def _measure_relevance(self, query: str, response: str) -> float:
        """Use embedding similarity to measure relevance"""
        query_embedding = self.embedder.encode([query])[0]
        response_embedding = self.embedder.encode([response])[0]
        
        similarity = cosine_similarity([query_embedding], [response_embedding])
        return float(similarity[0][0])
    
    async def _measure_coherence(self, response: str) -> float:
        """Evaluate response coherence using specialized model"""
        coherence_prompt = f"""
        Rate the coherence of this response on a scale of 0.0 to 1.0:
        
        Response: {response}
        
        Consider:
        - Logical flow of ideas
        - Consistency of information
        - Clarity of expression
        
        Score (0.0-1.0):
        """
        
        score_text = await self.quality_evaluator.evaluate(coherence_prompt)
        return float(score_text.strip())

class CostTracker:
    def __init__(self):
        self.model_pricing = {
            'gpt-4': {'input': 0.03, 'output': 0.06},  # per 1K tokens
            'gpt-3.5-turbo': {'input': 0.0015, 'output': 0.002},
            'claude-3': {'input': 0.025, 'output': 0.075}
        }
        
    def calculate_request_cost(self, model: str, input_tokens: int, 
                              output_tokens: int) -> float:
        """Calculate cost for a single request"""
        if model not in self.model_pricing:
            return 0.0
        
        pricing = self.model_pricing[model]
        input_cost = (input_tokens / 1000) * pricing['input']
        output_cost = (output_tokens / 1000) * pricing['output']
        
        return input_cost + output_cost
    
    async def track_daily_costs(self) -> Dict[str, float]:
        """Aggregate daily cost metrics"""
        today = datetime.now().date()
        
        # Query metrics database for today's usage
        daily_usage = await self.metrics_db.get_daily_usage(today)
        
        cost_breakdown = {
            'total_cost': sum(usage.cost for usage in daily_usage),
            'cost_by_model': self._group_costs_by_model(daily_usage),
            'cost_by_user': self._group_costs_by_user(daily_usage),
            'cost_by_agent_type': self._group_costs_by_agent_type(daily_usage)
        }
        
        # Set up alerts for cost thresholds
        if cost_breakdown['total_cost'] > self.daily_cost_threshold:
            await self.alert_manager.send_cost_alert(cost_breakdown)
        
        return cost_breakdown

from opentelemetry import trace
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

class DistributedTracing:
    def __init__(self):
        # Set up OpenTelemetry
        trace.set_tracer_provider(TracerProvider())
        tracer = trace.get_tracer(__name__)
        
        jaeger_exporter = JaegerExporter(
            agent_host_name="jaeger",
            agent_port=6831,
        )
        
        span_processor = BatchSpanProcessor(jaeger_exporter)
        trace.get_tracer_provider().add_span_processor(span_processor)
        
        self.tracer = tracer
    
    async def trace_agent_interaction(self, agent_name: str, operation: str):
        """Context manager for tracing agent operations"""
        with self.tracer.start_as_current_span(
            f"{agent_name}.{operation}"
        ) as span:
            # Add custom attributes
            span.set_attribute("agent.name", agent_name)
            span.set_attribute("agent.operation", operation)
            
            try:
                yield span
            except Exception as e:
                span.record_exception(e)
                span.set_status(trace.Status(trace.StatusCode.ERROR))
                raise

# Usage example
class TracedMultiAgentSystem:
    def __init__(self):
        self.tracing = DistributedTracing()
        self.research_agent = ResearchAgent()
        self.analysis_agent = AnalysisAgent()
    
    async def process_complex_query(self, query: str) -> str:
        with await self.tracing.trace_agent_interaction("orchestrator", "process_query"):
            
            # Research phase
            async with self.tracing.trace_agent_interaction("research", "gather_info") as research_span:
                research_span.set_attribute("query", query)
                research_data = await self.research_agent.gather_info(query)
                research_span.set_attribute("data_points", len(research_data))
            
            # Analysis phase
            async with self.tracing.trace_agent_interaction("analysis", "analyze_data") as analysis_span:
                analysis_span.set_attribute("input_data_size", len(research_data))
                insights = await self.analysis_agent.analyze(research_data)
                analysis_span.set_attribute("insights_generated", len(insights))
            
            return insights

from fastapi import FastAPI, WebSocket
import json

class AgentMonitoringAPI:
    def __init__(self):
        self.app = FastAPI()
        self.metrics_collector = AgentMetricsCollector()
        self.active_connections: List[WebSocket] = []
        self.setup_routes()
    
    def setup_routes(self):
        @self.app.websocket("/ws/metrics")
        async def websocket_endpoint(websocket: WebSocket):
            await websocket.accept()
            self.active_connections.append(websocket)
            
            try:
                while True:
                    # Send real-time metrics every 5 seconds
                    metrics = await self.get_real_time_metrics()
                    await websocket.send_text(json.dumps(metrics))
                    await asyncio.sleep(5)
            except:
                self.active_connections.remove(websocket)
        
        @self.app.get("/api/metrics/summary")
        async def get_metrics_summary():
            """Get aggregated metrics for the last 24 hours"""
            return await self.metrics_collector.get_24h_summary()
        
        @self.app.get("/api/metrics/costs")
        async def get_cost_breakdown():
            """Get cost breakdown by various dimensions"""
            cost_tracker = CostTracker()
            return await cost_tracker.get_cost_breakdown()
    
    async def get_real_time_metrics(self) -> dict:
        """Collect real-time system metrics"""
        return {
            'timestamp': time.time(),
            'active_sessions': await self._count_active_sessions(),
            'requests_per_minute': await self._get_rpm(),
            'average_response_time': await self._get_avg_response_time(),
            'error_rate': await self._get_error_rate(),
            'queue_lengths': await self._get_queue_lengths(),
            'model_health': await self._check_model_health()
        }
    
    async def broadcast_alert(self, alert_data: dict):
        """Broadcast alerts to all connected dashboards"""
        message = json.dumps({
            'type': 'alert',
            'data': alert_data
        })
        
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except:
                # Remove disconnected clients
                self.active_connections.remove(connection)

class ConversationFlowTracer:
    def __init__(self):
        self.conversation_steps = []
    
    def trace_reasoning_step(self, step_type: str, input_data: dict, 
                           output_data: dict, metadata: dict = None):
        """Record individual reasoning steps"""
        step = {
            'timestamp': time.time(),
            'step_type': step_type,
            'input': input_data,
            'output': output_data,
            'metadata': metadata or {},
            'step_id': len(self.conversation_steps)
        }
        self.conversation_steps.append(step)
    
    def generate_flow_diagram(self) -> str:
        """Generate mermaid diagram of conversation flow"""
        diagram = "graph TD\n"
        
        for i, step in enumerate(self.conversation_steps):
            node_label = f"{step['step_type']}\\n{step['metadata'].get('duration', 'N/A')}ms"
            diagram += f"    {i}[{node_label}]\n"
            
            if i > 0:
                diagram += f"    {i-1} --> {i}\n"
        
        return diagram
    
    def export_debugging_info(self) -> dict:
        """Export comprehensive debugging information"""
        return {
            'conversation_steps': self.conversation_steps,
            'total_duration': sum(
                step['metadata'].get('duration', 0) 
                for step in self.conversation_steps
            ),
            'step_count': len(self.conversation_steps),
            'flow_diagram': self.generate_flow_diagram()
        }

class ContextWindowAnalyzer:
    def __init__(self, model_context_limit: int = 4096):
        self.context_limit = model_context_limit
        self.tokenizer = get_tokenizer()
    
    def analyze_context_usage(self, full_context: str) -> dict:
        """Analyze how context window is being utilized"""
        sections = self._split_context_sections(full_context)
        
        analysis = {
            'total_tokens': len(self.tokenizer.encode(full_context)),
            'utilization_percentage': (len(self.tokenizer.encode(full_context)) / self.context_limit) * 100,
            'sections': {}
        }
        
        for section_name, section_content in sections.items():
            section_tokens = len(self.tokenizer.encode(section_content))
            analysis['sections'][section_name] = {
                'token_count': section_tokens,
                'percentage_of_total': (section_tokens / analysis['total_tokens']) * 100
            }
        
        # Identify optimization opportunities
        analysis['optimization_suggestions'] = self._suggest_optimizations(analysis)
        
        return analysis
    
    def _suggest_optimizations(self, analysis: dict) -> List[str]:
        suggestions = []
        
        if analysis['utilization_percentage'] > 90:
            suggestions.append("Context window nearly full - consider summarization")
        
        # Find sections taking up excessive space
        for section, data in analysis['sections'].items():
            if data['percentage_of_total'] > 40:
                suggestions.append(f"Section '{section}' uses {data['percentage_of_total']:.1f}% of context")
        
        return suggestions

class IntelligentAlertManager:
    def __init__(self):
        self.alert_thresholds = {
            'error_rate': 0.05,  # 5% error rate
            'response_time_p95': 10.0,  # 10 seconds
            'daily_cost': 1000.0,  # $1000/day
            'queue_depth': 100  # 100 pending requests
        }
        
        self.alert_cooldowns = {}  # Prevent alert spam
        self.escalation_rules = {}
    
    async def check_and_alert(self, metrics: dict):
        """Check metrics and send alerts if thresholds exceeded"""
        for metric_name, threshold in self.alert_thresholds.items():
            current_value = metrics.get(metric_name, 0)
            
            if current_value > threshold:
                # Check if we're in cooldown period
                if not self._in_cooldown(metric_name):
                    await self._send_alert(metric_name, current_value, threshold)
                    self._set_cooldown(metric_name, duration=300)  # 5 min cooldown
    
    async def _send_alert(self, metric_name: str, current_value: float, 
                         threshold: float):
        """Send alert through configured channels"""
        alert_data = {
            'metric': metric_name,
            'current_value': current_value,
            'threshold': threshold,
            'severity': self._calculate_severity(current_value, threshold),
            'timestamp': time.time()
        }
        
        # Send to multiple channels
        await self._send_slack_alert(alert_data)
        await self._send_email_alert(alert_data)
        
        # Log to metrics system
        await self.metrics_logger.log_alert(alert_data)
    
    def _calculate_severity(self, current: float, threshold: float) -> str:
        ratio = current / threshold
        
        if ratio >= 2.0:
            return "critical"
        elif ratio >= 1.5:
            return "high"
        elif ratio >= 1.2:
            return "medium"
        else:
            return "low"