AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Cost-tracking middleware for agent tool calls

Cost-Tracking Middleware for Agent Tool Calls

Agent systems can burn through API budgets fast. Every LLM call, every tool invocation, every retry costs tokens and time. Without cost tracking, you discover overspending after the invoice arrives. Cost-tracking middleware wraps every tool call to record token usage, latency, and estimated cost in real time, with optional budget enforcement that stops the agent before it exceeds a limit.

What you are building

Create a CostTracker middleware class that:

Wraps tool calls — intercepts every tool invocation to measure and record metrics before and after execution.
Tracks per-call metrics — token usage (input + output), latency in milliseconds, estimated cost in USD.
Maintains running totals — cumulative tokens, cost, and call count, grouped by tool name.
Enforces budgets — if a configurable budget (in USD or tokens) would be exceeded by the next call, raise a BudgetExceededError instead of executing.
Provides a summary — return a structured cost report with per-tool and aggregate breakdowns.

Why this matters

In production, cost visibility is not optional. A runaway agent loop that calls GPT-4 in a retry spiral can burn hundreds of dollars in minutes. Cost tracking middleware is the standard pattern for: (1) real-time budget enforcement, (2) per-request cost attribution for billing, (3) identifying expensive tool calls for optimization, and (4) comparing agent configurations by cost efficiency.

This is the kind of infrastructure every AI team builds early and never removes.

Constraints

The middleware must work with both sync and async tool functions.
Token counts come from the tool's return value (assume tools return a dict with an optional usage field).
The pricing table maps model/tool names to per-1K-token rates.
Budget enforcement must check BEFORE executing, not after (to prevent overspend).
All tracking data must be thread-safe for concurrent agent execution.

Agents / intermediate / Step 8 of 8

Practice stage

Agent architecture patterns

Hint

Focus on explicit control surfaces — function schemas, state machines, and structured outputs. Agents are most useful when the task requires dynamic tool selection or multi-step reasoning. Start with the simplest pattern (single tool call) before reaching for ReAct loops.

Next drill

This is the end of the current mini-sequence.

Success criteria

- Tools return structured, typed responses
- Agent completes the task within a bounded number of steps
- All tool calls include error handling and retries
- Memory/state management prevents unbounded context growth

Review checklist

- Tool schemas validate inputs and handle errors gracefully
- Agent loop has explicit termination conditions
- State is serializable and inspectable between steps
- Cost and token usage are tracked per invocation
- Fallback behavior exists for tool call failures

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable
from collections import defaultdict

class BudgetExceededError(Exception):
    """Raised when a tool call would exceed the configured budget."""
    def __init__(self, budget: float, current_cost: float, estimated_cost: float):
        self.budget = budget
        self.current_cost = current_cost
        self.estimated_cost = estimated_cost
        super().__init__(
            f"Budget ${budget:.4f} would be exceeded: "
            f"current=${current_cost:.4f}, estimated next=${estimated_cost:.4f}"
        )

@dataclass
class CallRecord:
    """Record of a single tool call."""
    tool_name: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    cost_usd: float
    timestamp: float

@dataclass
class ToolStats:
    """Aggregate stats for a single tool."""
    call_count: int = 0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_cost_usd: float = 0.0
    total_latency_ms: float = 0.0
    avg_latency_ms: float = 0.0

@dataclass
class CostReport:
    """Complete cost report."""
    total_calls: int = 0
    total_tokens: int = 0
    total_cost_usd: float = 0.0
    by_tool: dict[str, ToolStats] = field(default_factory=dict)
    budget_remaining: float | None = None

class CostTracker:
    """Middleware that wraps tool calls to track cost and enforce budgets."""

def __init__(
        self,
        pricing: dict[str, dict[str, float]] | None = None,
        budget_usd: float | None = None,
        budget_tokens: int | None = None,
    ) -> None:
        # pricing: {"tool_name": {"input": rate_per_1k, "output": rate_per_1k}}
        self.pricing = pricing or {}
        self.budget_usd = budget_usd
        self.budget_tokens = budget_tokens
        self._records: list[CallRecord] = []
        self._lock = threading.Lock()

def _estimate_cost(self, tool_name: str, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost for a tool call based on pricing table."""
        # TODO: look up pricing for tool_name, compute cost
        raise NotImplementedError

def _check_budget(self, estimated_cost: float, estimated_tokens: int) -> None:
        """Raise BudgetExceededError if budget would be exceeded."""
        # TODO: check USD and token budgets
        raise NotImplementedError

def wrap(self, tool_name: str, fn: Callable) -> Callable:
        """Return a wrapped version of fn that tracks costs."""
        # TODO: create wrapper that:
        #   1. Estimates cost and checks budget BEFORE calling
        #   2. Calls the function and measures latency
        #   3. Extracts usage from result
        #   4. Records the call
        raise NotImplementedError

def get_report(self) -> CostReport:
        """Generate a cost report."""
        # TODO: aggregate records into a CostReport
        raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.