Cost Optimization & Token Management

8 min read Module 7 of 10 Topic 20 of 30

What you'll learn

  • Calculate and track the per-run token cost of an agent workflow
  • Implement model cascading to use expensive models only where necessary
  • Add prompt caching to avoid re-paying for repeated context
  • Compress messages and tool results to reduce input token count without information loss
Building this at your company? For enterprise and company teams taking this to production: book a 30-minute session with our AI engineers for architecture guidance, code review, and a rollout plan for your use case.
Book a Team Session

The Real Cost of Production Agents

Let’s be concrete about what agents cost at scale:

Typical agent run (research task):
- System prompt:   1,500 tokens × 20 calls = 30,000 input tokens
- Tool results:    500 tokens avg × 12 calls = 6,000 input tokens  
- Model outputs:   200 tokens × 20 calls = 4,000 output tokens
- Total per run:   ~36,000 input + ~4,000 output tokens

At GPT-5.6 Sol pricing ($5/M input, $15/M output):
- Input:  36,000 × $5/1M = $0.18
- Output:  4,000 × $15/1M = $0.06
- Per run: $0.24

At 10,000 runs/day: $2,400/day = $72,000/month

That’s real money. The optimizations in this lesson can reduce this by 60-80%.


Token Counting Before You Spend

import tiktoken
from anthropic import Anthropic

# tiktoken.encoding_for_model gives the exact tokenizer the model uses: count is precise not estimated
encoder = tiktoken.encoding_for_model("gpt-5.6-sol")

def count_openai_tokens(messages: list[dict]) -> dict:
    token_count = 0
    for msg in messages:
        token_count += 4  # OpenAI charges 4 tokens of overhead per message for role/delimiters
        content = msg.get("content", "")
        if isinstance(content, str):
            # encoder.encode() splits the string into token IDs: len() gives the token count
            token_count += len(encoder.encode(content))
        elif isinstance(content, list):
            # content can be a list of blocks (text + images) in vision messages
            for block in content:
                if block.get("type") == "text":
                    token_count += len(encoder.encode(block["text"]))
    
    return {"total_tokens": token_count, "estimated_cost_usd": token_count * 5 / 1_000_000}

# Anthropic token counting (exact, uses the API): more accurate than tiktoken for Claude models
anthropic_client = Anthropic()

def count_anthropic_tokens(messages: list, system: str = "") -> int:
    # count_tokens() makes a lightweight API call: no generation happens, just tokenization
    response = anthropic_client.messages.count_tokens(
        model="claude-sonnet-5",
        system=system,
        messages=messages,
    )
    return response.input_tokens

Model Cascading

from enum import Enum
from dataclasses import dataclass

class TaskComplexity(Enum):
    TRIVIAL = "trivial"   # routing, yes/no decisions
    SIMPLE = "simple"     # formatting, extraction
    MEDIUM = "medium"     # summarization, classification
    COMPLEX = "complex"   # multi-step reasoning, code generation

@dataclass
class ModelConfig:
    name: str
    cost_per_m_input: float
    cost_per_m_output: float
    max_context: int

# Map each complexity level to the cheapest model that can reliably handle it
MODELS = {
    TaskComplexity.TRIVIAL: ModelConfig("gpt-5.6-luna", 0.15, 0.60, 128000),  # ~33x cheaper than gpt-5.6-sol
    TaskComplexity.SIMPLE:  ModelConfig("gpt-5.6-luna", 0.15, 0.60, 128000),
    TaskComplexity.MEDIUM:  ModelConfig("gpt-5.6-sol", 5.0, 15.0, 128000),
    TaskComplexity.COMPLEX: ModelConfig("gpt-5.6-sol", 5.0, 15.0, 128000),
}

def classify_task_complexity(task_description: str) -> TaskComplexity:
    # Use gpt-5.6-luna to classify: the classification itself is a simple task,
    # so we avoid paying gpt-5.6-sol prices just to decide which model to use
    response = ChatOpenAI(model="gpt-5.6-luna", temperature=0).invoke([
        HumanMessage(content=f"""Classify this task:
{task_description}

Return ONE word: trivial, simple, medium, or complex

trivial: yes/no decisions, routing, simple lookups
simple: data extraction, format conversion, summarizing one paragraph
medium: multi-paragraph synthesis, code debugging, analysis
complex: multi-step research, architecture design, complex coding""")
    ])
    
    complexity_str = response.content.strip().lower()
    # Fall back to MEDIUM if the model returns an unexpected value rather than crashing
    return TaskComplexity(complexity_str) if complexity_str in [c.value for c in TaskComplexity] else TaskComplexity.MEDIUM

def get_model_for_task(task: str) -> str:
    complexity = classify_task_complexity(task)
    return MODELS[complexity].name

# In your agent nodes:
def smart_model_node(state: AgentState) -> dict:
    # routing_model handles the cheap decisions: classifying, selecting next step
    routing_model = ChatOpenAI(model="gpt-5.6-luna", temperature=0)
    
    # reasoning_model is reserved for tasks that actually need deep capability
    reasoning_model = ChatOpenAI(model="gpt-5.6-sol", temperature=0)
    
    last_msg = state["messages"][-1].content if state["messages"] else ""
    
    # Short messages or explicit routing instructions are cheap decisions: use the mini model
    if "route to" in last_msg.lower() or len(last_msg) < 100:
        response = routing_model.invoke(state["messages"])
    else:
        response = reasoning_model.invoke(state["messages"])
    
    return {"messages": [response]}

This two-tier cascade works with any provider’s pricing tiers, not just OpenAI’s. If your reasoning tier is cost-sensitive, GLM-5.2 is worth benchmarking against gpt-5.6-sol for the MEDIUM and COMPLEX rows: at roughly a third of the input cost and a fifth of the output cost, swapping just the reasoning_model line to ChatOpenAI(model="glm-5.2", base_url="https://api.z.ai/api/paas/v4/") (Z.ai exposes an OpenAI-compatible endpoint) can meaningfully cut the expensive half of a cascade, provided your eval suite confirms it holds up on your actual task distribution.


Prompt Caching

# Anthropic Prompt Caching: reduces cost for repeated system prompts by ~90%
from anthropic import AsyncAnthropic

cached_client = AsyncAnthropic()

STATIC_SYSTEM = """You are a research assistant with access to tools..."""  # long system prompt

async def call_with_caching(user_message: str) -> str:
    response = await cached_client.messages.create(
        model="claude-sonnet-5",
        max_tokens=2000,
        system=[
            {
                "type": "text",
                "text": STATIC_SYSTEM,
                # cache_control: ephemeral tells Anthropic to cache this prefix
                # saves tokens on repeated calls: "ephemeral" means cached for ~5 minutes
                "cache_control": {"type": "ephemeral"},
            }
        ],
        messages=[{"role": "user", "content": user_message}],
    )
    
    # First call: full price ($3/M input tokens for claude-sonnet-5)
    # Subsequent calls with the same system prefix: $0.30/M for the cached portion (10% of original)
    # Check response.usage.cache_read_input_tokens to confirm cache hits
    return response.content[0].text

# OpenAI Prompt Caching (automatic for prompts > 1024 tokens)
# No special code needed: the cache is keyed on the exact prompt prefix automatically
# Verify cache hits via usage.prompt_tokens_details.cached_tokens in the response

Tool Result Compression

from bs4 import BeautifulSoup
import re

def compress_web_result(raw_html: str, max_chars: int = 2000) -> str:
    """Extract essential text from HTML, drop boilerplate."""
    soup = BeautifulSoup(raw_html, 'html.parser')
    
    # decompose() removes these tags from the parse tree: script/style/nav are pure noise
    for tag in soup(['script', 'style', 'nav', 'footer', 'header', 'aside', 'ads']):
        tag.decompose()
    
    # Try semantic tags first (main, article) before falling back to the full body
    main = soup.find('main') or soup.find('article') or soup.body
    # ' '.join(...split()) collapses all whitespace/newlines into single spaces
    text = ' '.join(main.get_text().split()) if main else ""
    
    # Hard cap at max_chars: raw HTML can be 100k+ chars; we only need the key passage
    return text[:max_chars]

def compress_api_response(response: dict, keep_fields: list[str]) -> dict:
    # Drop all fields the agent doesn't need: metadata, headers, pagination cursors, etc.
    return {k: v for k, v in response.items() if k in keep_fields}

def compress_search_results(results: list[dict]) -> str:
    compressed = []
    for r in results[:5]:  # cap at 5 results, more than 5 rarely improves agent quality
        compressed.append(
            f"**{r.get('title', '')}** ({r.get('url', '')})\n"
            f"{r.get('snippet', '')[:300]}"  # truncate snippet to 300 chars, enough context without bloat
        )
    return "\n\n".join(compressed)

Cost Dashboard Pattern

from dataclasses import dataclass, field
import json

@dataclass
class CostTracker:
    session_id: str
    _calls: list[dict] = field(default_factory=list)  # stores per-call records for the session
    
    # Prices are per million tokens (input, output): update when providers change pricing
    PRICING = {
        "gpt-5.6-sol": (5.0, 15.0),
        "gpt-5.6-luna": (0.15, 0.60),
        "claude-sonnet-5": (3.0, 15.0),
        "claude-haiku-4-5-20251001": (0.25, 1.25),
        "text-embedding-3-small": (0.02, 0.0),  # embeddings only have input cost
    }
    
    def record(self, model: str, input_tokens: int, output_tokens: int, node: str = ""):
        input_rate, output_rate = self.PRICING.get(model, (5.0, 15.0))  # default to gpt-5.6-sol prices if model unknown
        # Divide by 1_000_000 to convert from per-million-token rate to per-token cost
        cost = (input_tokens * input_rate + output_tokens * output_rate) / 1_000_000
        
        self._calls.append({
            "model": model, "node": node,  # node lets you see which graph node is most expensive
            "input_tokens": input_tokens, "output_tokens": output_tokens,
            "cost_usd": cost,
        })
    
    def summary(self) -> dict:
        total_cost = sum(c["cost_usd"] for c in self._calls)
        total_input = sum(c["input_tokens"] for c in self._calls)
        total_output = sum(c["output_tokens"] for c in self._calls)
        
        # Group by model: shows exactly which models are driving your costs
        by_model = {}
        for call in self._calls:
            m = call["model"]
            by_model[m] = by_model.get(m, {"calls": 0, "cost": 0.0})
            by_model[m]["calls"] += 1
            by_model[m]["cost"] += call["cost_usd"]
        
        return {
            "total_cost_usd": round(total_cost, 6),
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "by_model": by_model,  # e.g. {"gpt-5.6-sol": {"calls": 5, "cost": 0.12}, "gpt-5.6-luna": {...}}
        }

A 60-80% cost reduction on a $72,000/month agent system is $43,000-$58,000 in savings, and the techniques above (model cascading, prompt caching, tool compression) can realistically achieve that with minimal quality impact.

Knowledge Check

3 questions to test your understanding

1 An agent makes 20 tool calls, each including the full 2000-token system prompt. How many system prompt tokens does it consume, and how can you reduce this?

2 What is model cascading and when does it save money?

3 Why is tool result compression important for long-running agents?

Discussion

Questions and notes from learners on this topic

Loading discussion…

Go further with expert guidance

Ready to build production AI?
Talk to our R&D team.

These courses give you the foundation. Our embedded AI teams take you from prototype to production in 30–90 days, with your team, your codebase, your goals. Book a free strategy call to see how we can accelerate your AI initiative.

30 minutes · No obligation · Expert AI engineers, not sales reps

AI Architecture Review

Audit your current stack and identify high-impact improvements

Project Review

Get expert feedback on your AI implementation and codebase

Team Mentoring

Upskill your engineers with hands-on AI coaching sessions

AI Strategy

Define your AI roadmap, prioritization, and implementation plan