agent/api_monitor.py

"""
API Usage Monitor - Track all API calls with tokens, costs, and timing
Provides real-time visibility into LLM and GNS3 API usage for judges/users
"""

import time
import logging
import os
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field, asdict
from datetime import datetime
from threading import Lock
import json

logger = logging.getLogger(__name__)


# Pricing per 1M tokens (as of Nov 2024)
PRICING = {
    "gpt-4o": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-3-opus-20240229": {"input": 15.00, "output": 75.00},
    "claude-3-sonnet-20240229": {"input": 3.00, "output": 15.00},
    "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25},
    "claude-3-5-sonnet-20241022": {"input": 3.00, "output": 15.00},
    "anthropic/claude-3.5-sonnet": {"input": 3.00, "output": 15.00},  # OpenRouter
}


@dataclass
class APICall:
    """Record of a single API call"""
    call_id: str
    timestamp: str
    api_type: str  # "llm", "gns3", "netbox", etc.
    provider: str  # "openai", "anthropic", "openrouter", "gns3"
    endpoint: str  # Model name or API endpoint
    status: str  # "success", "error", "in-progress"
    duration_ms: Optional[float] = None
    input_tokens: Optional[int] = None
    output_tokens: Optional[int] = None
    total_tokens: Optional[int] = None
    estimated_cost: Optional[float] = None
    error_message: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        """Convert to dictionary for JSON serialization"""
        return asdict(self)
    
    def format_log_entry(self) -> str:
        """Format as a readable log entry for UI display"""
        icon_map = {
            "llm": "🤖",
            "gns3": "🌐",
            "netbox": "📊",
        }
        
        status_icon = "✅" if self.status == "success" else "❌" if self.status == "error" else "⏳"
        api_icon = icon_map.get(self.api_type, "🔧")
        
        parts = [f"{status_icon} {api_icon} **{self.provider.upper()}**"]
        
        if self.api_type == "llm":
            parts.append(f"`{self.endpoint}`")
            if self.input_tokens and self.output_tokens:
                parts.append(f"📝 {self.input_tokens:,}→{self.output_tokens:,} tokens")
            if self.estimated_cost:
                parts.append(f"💰 ${self.estimated_cost:.5f}")
        elif self.api_type == "gns3":
            parts.append(f"`{self.endpoint}`")
        
        if self.duration_ms:
            parts.append(f"⏱️ {self.duration_ms:.0f}ms")
        
        if self.error_message:
            parts.append(f"⚠️ {self.error_message}")
        
        return " | ".join(parts)


@dataclass
class SessionStats:
    """Cumulative statistics for a session"""
    total_calls: int = 0
    llm_calls: int = 0
    gns3_calls: int = 0
    total_tokens: int = 0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_cost: float = 0.0
    errors: int = 0
    start_time: str = field(default_factory=lambda: datetime.now().isoformat())
    budget_limit: Optional[float] = None
    budget_alert_fraction: float = 0.8
    budget_alert_triggered: bool = False
    
    def to_dict(self) -> Dict:
        """Convert to dictionary"""
        return asdict(self)
    
    def budget_remaining(self) -> Optional[float]:
        """Return remaining budget if configured"""
        if self.budget_limit is None:
            return None
        return max(self.budget_limit - self.total_cost, 0.0)
    
    def budget_usage_ratio(self) -> Optional[float]:
        """Return fraction of budget consumed"""
        if self.budget_limit is None or self.budget_limit == 0:
            return None
        return self.total_cost / self.budget_limit
    
    def budget_status(self) -> str:
        """Human-readable budget status message"""
        if self.budget_limit is None or self.budget_limit == 0:
            return "Budget not configured"
        
        ratio = self.budget_usage_ratio() or 0.0
        remaining = self.budget_remaining()
        
        if ratio >= 1.0:
            return f"🛑 Budget exceeded by ${abs(remaining):.2f}"
        if ratio >= self.budget_alert_fraction:
            return f"⚠️ {ratio*100:.0f}% of budget used (limit ${self.budget_limit:.2f})"
        return f"✅ {ratio*100:.0f}% of budget used, ${remaining:.2f} remaining"
    
    def format_dashboard(self, include_heading: bool = True) -> str:
        """Format as markdown dashboard for UI"""
        uptime = datetime.now() - datetime.fromisoformat(self.start_time)
        uptime_str = f"{int(uptime.total_seconds() / 60)}m {int(uptime.total_seconds() % 60)}s"
        
        budget_line = "n/a"
        if self.budget_limit not in (None, 0):
            remaining = self.budget_remaining()
            budget_line = f"${self.budget_limit:.2f} (remaining ${remaining:.2f})"
        status_line = self.budget_status()
        
        header = "### 📊 Session Statistics\n\n" if include_heading else ""
        return (
            f"{header}"
            "| Metric | Value |\n"
            "|--------|-------|\n"
            f"| ⏰ Session Duration | {uptime_str} |\n"
            f"| 📞 Total API Calls | {self.total_calls:,} |\n"
            f"| 🤖 LLM Calls | {self.llm_calls:,} |\n"
            f"| 🌐 GNS3 Calls | {self.gns3_calls:,} |\n"
            f"| 📝 Total Tokens | {self.total_tokens:,} |\n"
            f"| 💰 **Total Cost** | **${self.total_cost:.4f}** |\n"
            f"| ❌ Errors | {self.errors} |\n"
            f"| 🧭 Budget | {budget_line} |\n"
            f"| 🚨 Budget Status | {status_line} |\n\n"
            "---\n\n"
            "### 💵 Cost Breakdown\n"
            f"- Input tokens: {self.total_input_tokens:,} ({self.total_input_tokens / 1000000:.2f}M)\n"
            f"- Output tokens: {self.total_output_tokens:,} ({self.total_output_tokens / 1000000:.2f}M)\n"
            f"- Avg cost per call: ${self.total_cost / max(self.total_calls, 1):.4f}\n"
        )


class APIMonitor:
    """
    Singleton monitor for tracking all API usage in the application
    Thread-safe for concurrent access
    """
    
    _instance = None
    _lock = Lock()
    
    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        
        self._initialized = True
        self.calls: List[APICall] = []
        self.active_calls: Dict[str, float] = {}  # call_id -> start_time
        self._call_index: Dict[str, APICall] = {}
        
        # Budget configuration (env vars for quick tuning)
        budget_env = os.getenv("API_BUDGET_USD")
        budget_limit = float(budget_env) if budget_env else None
        budget_alert_fraction = float(os.getenv("API_BUDGET_ALERT_FRACTION", "0.8"))
        
        self.stats = SessionStats(
            budget_limit=budget_limit,
            budget_alert_fraction=budget_alert_fraction
        )
        logger.info("API Monitor initialized")
    
    def start_call(self, call_id: str, api_type: str, provider: str, endpoint: str, **metadata) -> APICall:
        """
        Start tracking an API call
        Returns an APICall object that should be updated when complete
        """
        with self._lock:
            call = APICall(
                call_id=call_id,
                timestamp=datetime.now().isoformat(),
                api_type=api_type,
                provider=provider,
                endpoint=endpoint,
                status="in-progress",
                metadata=metadata
            )
            self.calls.append(call)
            self.active_calls[call_id] = time.time()
            self._call_index[call_id] = call
            return call
    
    def complete_call(
        self,
        call_id: str,
        success: bool = True,
        input_tokens: Optional[int] = None,
        output_tokens: Optional[int] = None,
        error_message: Optional[str] = None,
        **metadata
    ):
        """Mark a call as completed and update statistics"""
        with self._lock:
            start_time = self.active_calls.pop(call_id, None)
            duration_ms = (time.time() - start_time) * 1000 if start_time else None
            
            # Find the matching call by call_id
            call = self._call_index.pop(call_id, None)
            if not call:
                # Fallback to the most recent in-progress call
                for c in reversed(self.calls):
                    if c.status == "in-progress":
                        call = c
                        break
            
            if not call:
                logger.warning(f"Could not find call {call_id} to complete")
                return
            
            # Update call details
            call.status = "success" if success else "error"
            call.duration_ms = duration_ms
            call.input_tokens = input_tokens
            call.output_tokens = output_tokens
            call.error_message = error_message
            call.metadata.update(metadata)
            
            # Handle token accounting even if only one side is present
            if input_tokens is not None or output_tokens is not None:
                in_tokens = input_tokens or 0
                out_tokens = output_tokens or 0
                call.total_tokens = in_tokens + out_tokens
                
                # Calculate cost if it's an LLM call
                if call.api_type == "llm" and call.endpoint in PRICING:
                    pricing = PRICING[call.endpoint]
                    input_cost = (in_tokens / 1_000_000) * pricing["input"]
                    output_cost = (out_tokens / 1_000_000) * pricing["output"]
                    call.estimated_cost = input_cost + output_cost
            
            # Update session stats
            self.stats.total_calls += 1
            
            if call.api_type == "llm":
                self.stats.llm_calls += 1
            elif call.api_type == "gns3":
                self.stats.gns3_calls += 1
            
            if call.total_tokens is not None:
                self.stats.total_tokens += call.total_tokens
            if call.input_tokens is not None:
                self.stats.total_input_tokens += call.input_tokens
            if call.output_tokens is not None:
                self.stats.total_output_tokens += call.output_tokens
            if call.estimated_cost:
                self.stats.total_cost += call.estimated_cost
                
                # Budget monitoring
                if self.stats.budget_limit not in (None, 0):
                    usage_ratio = self.stats.budget_usage_ratio() or 0.0
                    if usage_ratio >= self.stats.budget_alert_fraction:
                        self.stats.budget_alert_triggered = True
            
            if not success:
                self.stats.errors += 1
            
            logger.info(f"Completed API call: {call.format_log_entry()}")
    
    def get_recent_calls(self, limit: int = 20) -> List[APICall]:
        """Get the most recent API calls"""
        with self._lock:
            return self.calls[-limit:]
    
    def get_recent_errors(self, limit: int = 5) -> List[APICall]:
        """Get recent errored API calls"""
        with self._lock:
            errors = [c for c in self.calls if c.status == "error"]
            return errors[-limit:]
    
    def get_all_calls(self) -> List[APICall]:
        """Get all API calls"""
        with self._lock:
            return self.calls.copy()
    
    def get_stats(self) -> SessionStats:
        """Get current session statistics"""
        with self._lock:
            return self.stats
    
    def format_activity_feed(self, limit: int = 20) -> str:
        """Format recent API activity as markdown for UI display"""
        recent = self.get_recent_calls(limit)
        
        if not recent:
            return "## 🔍 API Activity Feed\n\n*No API calls yet. Start using the system to see activity here.*"
        
        lines = ["## 🔍 API Activity Feed\n"]
        lines.append("*Most recent calls first*\n")
        
        for call in reversed(recent):
            lines.append(f"- {call.format_log_entry()}")
        
        if self.stats.budget_alert_triggered:
            lines.append("\n> 🚨 Budget alert: " + self.stats.budget_status())
        
        return "\n".join(lines)
    
    def reset(self):
        """Reset all tracking (for testing or new sessions)"""
        with self._lock:
            self.calls.clear()
            budget_limit = self.stats.budget_limit
            alert_fraction = self.stats.budget_alert_fraction
            self.stats = SessionStats(
                budget_limit=budget_limit,
                budget_alert_fraction=alert_fraction
            )
            self.active_calls.clear()
            self._call_index.clear()
            logger.info("API Monitor reset")
    
    def export_json(self) -> str:
        """Export all data as JSON"""
        with self._lock:
            return json.dumps({
                "calls": [c.to_dict() for c in self.calls],
                "stats": self.stats.to_dict()
            }, indent=2)


# Global singleton instance
monitor = APIMonitor()