AI Agent Cost Tracking and Production Monitoring Complete Guide 2026

from dataclasses import dataclass
from typing import List, Optional, Dict
from datetime import datetime
import uuid

@dataclass
class TokenUsage:
    """Token consumption for a single LLM call"""
    input_tokens: int
    output_tokens: int
    cached_tokens: int = 0  # Prompt caching savings

@dataclass
class AgentCall:
    """Single agent invocation"""
    call_id: str
    agent_id: str
    parent_call_id: Optional[str]  # For hierarchy
    model: str
    tokens: TokenUsage
    latency_ms: float
    timestamp: datetime
    metadata: Dict  # user_id, workflow_id, etc.

class HierarchicalAgent:
    """
    Example hierarchical agent structure
    Shows why cost attribution is complex
    """

    def process_request(self, user_query: str) -> str:
        """
        Coordinator agent that delegates to specialized agents

        Cost attribution challenge:
        - Coordinator makes 1 LLM call
        - Delegates to 3 specialized agents (3 calls)
        - Each specialized agent may delegate further
        - How do we attribute total cost?
        """

        # Coordinator analyzes query (uses tokens)
        analysis = self.llm_call(
            agent_id="coordinator",
            prompt=f"Analyze this query and determine required agents: {user_query}"
        )

        # Delegate to specialized agents based on analysis
        if "code" in analysis:
            code_result = self.code_agent.generate(user_query)  # More LLM calls

        if "data" in analysis:
            data_result = self.data_agent.query(user_query)  # More LLM calls

        if "explanation" in analysis:
            explanation = self.explainer_agent.explain(user_query)  # More LLM calls

        # Coordinator synthesizes results (more tokens)
        final_response = self.llm_call(
            agent_id="coordinator",
            prompt=f"Synthesize: {code_result} {data_result} {explanation}"
        )

        # Question: What's the total cost of this request?
        # Answer: Sum of ALL agent calls in the hierarchy

from collections import defaultdict
from typing import Dict, List
import time

class AgentCostTracker:
    """
    Production-ready hierarchical cost attribution system
    Tracks costs across complex multi-agent architectures
    """

    def __init__(self, pricing: Dict[str, Dict[str, float]]):
        """
        Initialize with model-specific pricing

        Args:
            pricing: {
                'gpt-5.2': {'input': 3.0, 'output': 15.0},  # per 1M tokens
                'claude-sonnet-4-5': {'input': 3.0, 'output': 15.0},
                'gemini-3-pro': {'input': 2.5, 'output': 10.0}
            }
        """
        self.pricing = pricing
        self.calls: List[AgentCall] = []
        self.hierarchy: Dict[str, List[str]] = defaultdict(list)  # parent -> children

    def track_call(
        self,
        agent_id: str,
        parent_call_id: Optional[str],
        model: str,
        tokens: TokenUsage,
        latency_ms: float,
        metadata: Dict = None
    ) -> str:
        """
        Track a single agent call

        Args:
            agent_id: Identifier for this agent
            parent_call_id: ID of parent call (None for root)
            model: Model used (gpt-5.2, claude-sonnet-4-5, etc.)
            tokens: Token usage breakdown
            latency_ms: Call latency in milliseconds
            metadata: Additional context (user_id, workflow_id, etc.)

        Returns:
            call_id for this invocation
        """
        call_id = str(uuid.uuid4())

        call = AgentCall(
            call_id=call_id,
            agent_id=agent_id,
            parent_call_id=parent_call_id,
            model=model,
            tokens=tokens,
            latency_ms=latency_ms,
            timestamp=datetime.now(),
            metadata=metadata or {}
        )

        self.calls.append(call)

        # Track hierarchy
        if parent_call_id:
            self.hierarchy[parent_call_id].append(call_id)

        return call_id

    def calculate_call_cost(self, call: AgentCall) -> float:
        """
        Calculate cost for a single call

        Returns:
            Cost in dollars
        """
        if call.model not in self.pricing:
            print(f"Warning: No pricing for model {call.model}")
            return 0.0

        pricing = self.pricing[call.model]

        # Calculate input cost (subtract cached tokens which are free/cheaper)
        billable_input = call.tokens.input_tokens - call.tokens.cached_tokens
        input_cost = (billable_input / 1_000_000) * pricing['input']

        # Calculate output cost
        output_cost = (call.tokens.output_tokens / 1_000_000) * pricing['output']

        total_cost = input_cost + output_cost
        return round(total_cost, 6)

    def get_hierarchical_cost(self, call_id: str) -> Dict:
        """
        Get total cost including all child calls

        Args:
            call_id: Root call ID to calculate from

        Returns:
            {
                'direct_cost': 0.045,  # This call only
                'child_cost': 0.070,   # All children
                'total_cost': 0.115,   # Direct + children
                'children': [...]      # Child call details
            }
        """
        # Find the call
        call = next((c for c in self.calls if c.call_id == call_id), None)
        if not call:
            return {'error': 'Call not found'}

        # Calculate direct cost
        direct_cost = self.calculate_call_cost(call)

        # Recursively calculate child costs
        child_calls = self.hierarchy.get(call_id, [])
        child_costs = []

        for child_id in child_calls:
            child_result = self.get_hierarchical_cost(child_id)
            child_costs.append(child_result)

        total_child_cost = sum(c['total_cost'] for c in child_costs)

        return {
            'call_id': call_id,
            'agent_id': call.agent_id,
            'model': call.model,
            'direct_cost': direct_cost,
            'child_cost': total_child_cost,
            'total_cost': direct_cost + total_child_cost,
            'tokens': {
                'input': call.tokens.input_tokens,
                'output': call.tokens.output_tokens,
                'cached': call.tokens.cached_tokens
            },
            'latency_ms': call.latency_ms,
            'children': child_costs
        }

    def get_cost_by_attribute(
        self,
        attribute: str,  # 'user_id', 'workflow_id', 'team_id'
        timeframe_hours: int = 24
    ) -> Dict[str, float]:
        """
        Aggregate costs by metadata attribute

        Args:
            attribute: Metadata key to group by
            timeframe_hours: Look back window

        Returns:
            Dictionary mapping attribute value to total cost
        """
        cutoff = datetime.now() - timedelta(hours=timeframe_hours)
        recent = [c for c in self.calls if c.timestamp > cutoff]

        # Group by attribute
        costs = defaultdict(float)

        for call in recent:
            attr_value = call.metadata.get(attribute, 'unknown')

            # Only count root-level calls to avoid double-counting
            if not call.parent_call_id:
                hierarchical_cost = self.get_hierarchical_cost(call.call_id)
                costs[attr_value] += hierarchical_cost['total_cost']

        return dict(costs)

    def get_cost_by_agent_type(
        self,
        timeframe_hours: int = 24
    ) -> Dict[str, Dict]:
        """
        Analyze costs by agent type

        Returns:
            {
                'coordinator': {
                    'total_cost': 5.23,
                    'call_count': 1250,
                    'avg_cost_per_call': 0.00418
                },
                ...
            }
        """
        cutoff = datetime.now() - timedelta(hours=timeframe_hours)
        recent = [c for c in self.calls if c.timestamp > cutoff]

        # Group by agent_id
        agent_stats = defaultdict(lambda: {'total_cost': 0.0, 'call_count': 0})

        for call in recent:
            cost = self.calculate_call_cost(call)
            agent_stats[call.agent_id]['total_cost'] += cost
            agent_stats[call.agent_id]['call_count'] += 1

        # Calculate averages
        results = {}
        for agent_id, stats in agent_stats.items():
            results[agent_id] = {
                'total_cost': round(stats['total_cost'], 2),
                'call_count': stats['call_count'],
                'avg_cost_per_call': round(
                    stats['total_cost'] / stats['call_count'], 5
                ) if stats['call_count'] > 0 else 0
            }

        # Sort by total cost descending
        results = dict(
            sorted(results.items(), key=lambda x: x[1]['total_cost'], reverse=True)
        )

        return results

# Example usage - simulating multi-agent request
pricing = {
    'gpt-5.2': {'input': 3.0, 'output': 15.0},
    'claude-sonnet-4-5': {'input': 3.0, 'output': 15.0},
    'gemini-3-pro': {'input': 2.5, 'output': 10.0}
}

tracker = AgentCostTracker(pricing)

# Simulate hierarchical agent calls
# Root: Coordinator
coordinator_id = tracker.track_call(
    agent_id='coordinator',
    parent_call_id=None,  # Root call
    model='gpt-5.2',
    tokens=TokenUsage(input_tokens=1500, output_tokens=500),
    latency_ms=1200,
    metadata={'user_id': 'user_123', 'workflow_id': 'code_generation'}
)

# Child 1: Code Agent (delegates to coordinator)
code_agent_id = tracker.track_call(
    agent_id='code_agent',
    parent_call_id=coordinator_id,
    model='gpt-5.2',
    tokens=TokenUsage(input_tokens=2000, output_tokens=800),
    latency_ms=1800,
    metadata={'user_id': 'user_123', 'workflow_id': 'code_generation'}
)

# Grandchild: Syntax Analyzer (delegates to code agent)
tracker.track_call(
    agent_id='syntax_analyzer',
    parent_call_id=code_agent_id,
    model='gemini-3-pro',
    tokens=TokenUsage(input_tokens=500, output_tokens=200),
    latency_ms=600,
    metadata={'user_id': 'user_123', 'workflow_id': 'code_generation'}
)

# Get hierarchical cost breakdown
cost_breakdown = tracker.get_hierarchical_cost(coordinator_id)
print(f"Total request cost: ${cost_breakdown['total_cost']:.4f}")
print(f"  Coordinator direct: ${cost_breakdown['direct_cost']:.4f}")
print(f"  Children total: ${cost_breakdown['child_cost']:.4f}")

# Get costs by user
user_costs = tracker.get_cost_by_attribute('user_id', timeframe_hours=24)
print(f"\nCosts by user: {user_costs}")

# Get costs by agent type
agent_costs = tracker.get_cost_by_agent_type(timeframe_hours=24)
print(f"\nCosts by agent type:")
for agent, stats in agent_costs.items():
    print(f"  {agent}: ${stats['total_cost']:.2f} ({stats['call_count']} calls)")

from datetime import datetime, timedelta
from typing import Dict, Optional
import time

class AgentBudgetManager:
    """
    Production budget management with auto-throttling
    Prevents budget overruns before they happen
    """

    def __init__(
        self,
        cost_tracker: AgentCostTracker,
        budgets: Dict[str, Dict]
    ):
        """
        Initialize with budgets by dimension

        Args:
            cost_tracker: Existing cost tracking system
            budgets: {
                'user_id': {
                    'user_123': {'daily': 10.0, 'monthly': 200.0},
                    'user_456': {'daily': 5.0, 'monthly': 100.0}
                },
                'workflow_id': {
                    'code_generation': {'daily': 50.0, 'monthly': 1000.0}
                }
            }
        """
        self.cost_tracker = cost_tracker
        self.budgets = budgets
        self.alerts_sent = set()  # Track already-sent alerts

    def check_budget(
        self,
        metadata: Dict,
        estimated_cost: float
    ) -> Dict:
        """
        Check if request would exceed budget

        Args:
            metadata: Request metadata (user_id, workflow_id, etc.)
            estimated_cost: Projected cost for this request

        Returns:
            {
                'allowed': True/False,
                'reason': 'Budget OK' or reason for denial,
                'remaining_budget': {
                    'daily': 5.23,
                    'monthly': 87.45
                },
                'current_usage': {
                    'daily': 4.77,
                    'monthly': 112.55
                }
            }
        """
        result = {
            'allowed': True,
            'reason': 'Budget OK',
            'remaining_budget': {},
            'current_usage': {},
            'warnings': []
        }

        # Check each budget dimension
        for dimension, dimension_budgets in self.budgets.items():
            if dimension not in metadata:
                continue

            value = metadata[dimension]
            if value not in dimension_budgets:
                continue

            budget_limits = dimension_budgets[value]

            # Check daily budget
            if 'daily' in budget_limits:
                daily_usage = sum(
                    self.cost_tracker.get_cost_by_attribute(
                        attribute=dimension,
                        timeframe_hours=24
                    ).get(value, 0.0)
                    for _ in [None]  # Single iteration
                )

                daily_limit = budget_limits['daily']
                daily_remaining = daily_limit - daily_usage

                result['current_usage'][f'{dimension}_daily'] = round(daily_usage, 2)
                result['remaining_budget'][f'{dimension}_daily'] = round(daily_remaining, 2)

                # Would this request exceed daily budget?
                if daily_usage + estimated_cost > daily_limit:
                    result['allowed'] = False
                    result['reason'] = f"Daily budget exceeded for {dimension}={value}"
                    return result

                # Warning at 80% usage
                if daily_usage / daily_limit > 0.8:
                    result['warnings'].append(
                        f"{dimension}={value} at {daily_usage/daily_limit*100:.0f}% of daily budget"
                    )

            # Check monthly budget
            if 'monthly' in budget_limits:
                monthly_usage = sum(
                    self.cost_tracker.get_cost_by_attribute(
                        attribute=dimension,
                        timeframe_hours=24*30
                    ).get(value, 0.0)
                    for _ in [None]
                )

                monthly_limit = budget_limits['monthly']
                monthly_remaining = monthly_limit - monthly_usage

                result['current_usage'][f'{dimension}_monthly'] = round(monthly_usage, 2)
                result['remaining_budget'][f'{dimension}_monthly'] = round(monthly_remaining, 2)

                if monthly_usage + estimated_cost > monthly_limit:
                    result['allowed'] = False
                    result['reason'] = f"Monthly budget exceeded for {dimension}={value}"
                    return result

                if monthly_usage / monthly_limit > 0.8:
                    result['warnings'].append(
                        f"{dimension}={value} at {monthly_usage/monthly_limit*100:.0f}% of monthly budget"
                    )

        return result

    def auto_throttle(
        self,
        metadata: Dict,
        estimated_cost: float
    ) -> Dict:
        """
        Auto-throttle request if approaching budget limits

        Returns:
            {
                'action': 'allow'|'throttle'|'deny',
                'wait_seconds': 0 (if throttled),
                'explanation': 'why throttled/denied'
            }
        """
        budget_check = self.check_budget(metadata, estimated_cost)

        if not budget_check['allowed']:
            return {
                'action': 'deny',
                'wait_seconds': 0,
                'explanation': budget_check['reason']
            }

        # Check if close to limits (>90% usage)
        for dimension, usage in budget_check['current_usage'].items():
            limit_key = dimension.replace('current_usage', 'remaining_budget')
            remaining = budget_check['remaining_budget'].get(dimension, float('inf'))
            total = usage + remaining

            if usage / total > 0.9:
                # Throttle: delay request
                return {
                    'action': 'throttle',
                    'wait_seconds': 5,  # Wait 5 seconds
                    'explanation': f"Throttling due to {dimension} at {usage/total*100:.0f}%"
                }

        return {
            'action': 'allow',
            'wait_seconds': 0,
            'explanation': 'Within budget limits'
        }

    def send_alert(
        self,
        alert_type: str,
        message: str,
        metadata: Dict
    ):
        """
        Send budget alert (email, Slack, PagerDuty, etc.)

        Args:
            alert_type: 'warning'|'critical'|'info'
            message: Alert message
            metadata: Context information
        """
        alert_key = f"{alert_type}:{message}:{metadata.get('user_id', 'unknown')}"

        # Deduplicate alerts (don't spam)
        if alert_key in self.alerts_sent:
            return

        print(f"\n{'='*60}")
        print(f"[{alert_type.upper()}] Budget Alert")
        print(f"{message}")
        print(f"Context: {metadata}")
        print(f"{'='*60}\n")

        # In production: send to alerting system
        # slack_webhook.send(message)
        # pagerduty.trigger_alert(message, severity=alert_type)

        self.alerts_sent.add(alert_key)

    def monitor_budgets(self):
        """
        Background task to monitor budget usage
        Runs periodically (every hour) to send proactive alerts
        """
        for dimension, dimension_budgets in self.budgets.items():
            usage = self.cost_tracker.get_cost_by_attribute(
                attribute=dimension,
                timeframe_hours=24
            )

            for value, limits in dimension_budgets.items():
                current_usage = usage.get(value, 0.0)

                if 'daily' in limits:
                    daily_limit = limits['daily']
                    usage_pct = (current_usage / daily_limit) * 100

                    if usage_pct >= 90:
                        self.send_alert(
                            'critical',
                            f"{dimension}={value} at {usage_pct:.0f}% of daily budget (${current_usage:.2f}/${daily_limit:.2f})",
                            {'dimension': dimension, 'value': value}
                        )
                    elif usage_pct >= 75:
                        self.send_alert(
                            'warning',
                            f"{dimension}={value} at {usage_pct:.0f}% of daily budget",
                            {'dimension': dimension, 'value': value}
                        )

# Example usage
budgets = {
    'user_id': {
        'user_123': {'daily': 10.0, 'monthly': 200.0},
        'user_456': {'daily': 5.0, 'monthly': 100.0}
    },
    'workflow_id': {
        'code_generation': {'daily': 50.0, 'monthly': 1000.0},
        'data_analysis': {'daily': 30.0, 'monthly': 600.0}
    }
}

budget_manager = AgentBudgetManager(tracker, budgets)

# Before making agent call, check budget
metadata = {'user_id': 'user_123', 'workflow_id': 'code_generation'}
estimated_cost = 0.085

throttle_decision = budget_manager.auto_throttle(metadata, estimated_cost)

if throttle_decision['action'] == 'deny':
    print(f"Request denied: {throttle_decision['explanation']}")
elif throttle_decision['action'] == 'throttle':
    print(f"Throttling: {throttle_decision['explanation']}")
    time.sleep(throttle_decision['wait_seconds'])
    # Then proceed with request
else:
    print("Request allowed - proceeding")
    # Make agent call

# Run periodic monitoring (in production: background job)
budget_manager.monitor_budgets()