Per-Step Token and Cost Profiling
Measuring token consumption per tool call and per reasoning step.
Per-Step Token and Cost Profiling is a free AI Agents lesson on CoddyKit — lesson 2 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Agents learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
Why Profile Token Usage?
LLM API costs scale directly with token usage. A single agent run can make dozens of LLM calls. Without per-step profiling, you cannot know which step is expensive, where to cache, or how to reduce costs.
Reading Token Usage from OpenAI
Every OpenAI completion response includes a usage object with prompt_tokens, completion_tokens, and total_tokens. Always capture this.
import openai
client = openai.OpenAI(api_key='sk-...')
def call_llm_with_tracking(prompt: str, model: str = 'gpt-4o-mini') -> dict:
response = client.chat.completions.create(
model=model,
messages=[{'role': 'user', 'content': prompt}]
)
usage = response.usage
return {
'content': response.choices[0].message.content,
'prompt_tokens': usage.prompt_tokens,
'completion_tokens': usage.completion_tokens,
'total_tokens': usage.total_tokens,
'model': model
}
result = call_llm_with_tracking('What is the capital of France?')
print(f'Response: {result["content"]}')
print(f'Tokens - Prompt: {result["prompt_tokens"]}, Completion: {result["completion_tokens"]}, Total: {result["total_tokens"]}')Cost Calculation Per Call
Calculate the dollar cost of each LLM call using the pricing table. Costs are typically per million tokens, so: cost = (prompt_tokens / 1_000_000) * input_price + (completion_tokens / 1_000_000) * output_price.
# Pricing per million tokens (as of early 2025 - verify current prices)
MODEL_PRICING = {
'gpt-4o': {'input': 2.50, 'output': 10.00},
'gpt-4o-mini': {'input': 0.15, 'output': 0.60},
'gpt-4-turbo': {'input': 10.00, 'output': 30.00},
'claude-3-5-sonnet-20241022': {'input': 3.00, 'output': 15.00},
'claude-3-haiku-20240307': {'input': 0.25, 'output': 1.25}
}
def calculate_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
pricing = MODEL_PRICING.get(model)
if not pricing:
return 0.0
input_cost = (prompt_tokens / 1_000_000) * pricing['input']
output_cost = (completion_tokens / 1_000_000) * pricing['output']
return input_cost + output_cost
# Example
prompt_tokens = 500
completion_tokens = 200
model = 'gpt-4o-mini'
cost = calculate_cost(prompt_tokens, completion_tokens, model)
print(f'Cost for {prompt_tokens}+{completion_tokens} tokens on {model}: ${cost:.6f}')Cumulative Cost Tracker
Track cumulative cost across an entire agent run. A cost tracker accumulates token usage and cost per step, making it easy to see which step consumed the most budget.
from dataclasses import dataclass, field
from typing import List
@dataclass
class StepCost:
step_name: str
model: str
prompt_tokens: int
completion_tokens: int
cost_usd: float
@dataclass
class CostTracker:
steps: List[StepCost] = field(default_factory=list)
def record(self, step_name: str, model: str, prompt_tokens: int, completion_tokens: int):
cost = calculate_cost(prompt_tokens, completion_tokens, model)
self.steps.append(StepCost(
step_name=step_name,
model=model,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
cost_usd=cost
))
@property
def total_cost(self) -> float:
return sum(s.cost_usd for s in self.steps)
@property
def total_tokens(self) -> int:
return sum(s.prompt_tokens + s.completion_tokens for s in self.steps)
def summary(self) -> str:
lines = ['=== Cost Summary ===']
for step in self.steps:
lines.append(f'{step.step_name}: {step.prompt_tokens}+{step.completion_tokens} tokens = ${step.cost_usd:.6f}')
lines.append(f'TOTAL: {self.total_tokens} tokens = ${self.total_cost:.6f}')
return '\n'.join(lines)
tracker = CostTracker()
tracker.record('entity_extraction', 'gpt-4o-mini', 200, 50)
tracker.record('vector_search_query', 'gpt-4o-mini', 100, 30)
tracker.record('answer_generation', 'gpt-4o-mini', 1500, 300)
print(tracker.summary())Cost Per Tool Call Type
Break down costs by tool call type across many agent runs. Some tools are called much more often and are the dominant cost drivers.
from collections import defaultdict
class ToolCostAnalyzer:
def __init__(self):
self.tool_stats = defaultdict(lambda: {
'call_count': 0,
'total_prompt_tokens': 0,
'total_completion_tokens': 0,
'total_cost_usd': 0.0
})
def record_tool_call(self, tool_name: str, prompt_tokens: int, completion_tokens: int, model: str):
cost = calculate_cost(prompt_tokens, completion_tokens, model)
stats = self.tool_stats[tool_name]
stats['call_count'] += 1
stats['total_prompt_tokens'] += prompt_tokens
stats['total_completion_tokens'] += completion_tokens
stats['total_cost_usd'] += cost
def report(self):
print('=== Tool Cost Breakdown ===')
sorted_tools = sorted(
self.tool_stats.items(),
key=lambda x: x[1]['total_cost_usd'],
reverse=True
)
for tool_name, stats in sorted_tools:
avg_cost = stats['total_cost_usd'] / stats['call_count']
print(f'{tool_name}: {stats["call_count"]} calls, total ${stats["total_cost_usd"]:.4f}, avg ${avg_cost:.6f}/call')
analyzer = ToolCostAnalyzer()
analyzer.record_tool_call('search_web', 800, 200, 'gpt-4o-mini')
analyzer.record_tool_call('search_web', 750, 180, 'gpt-4o-mini')
analyzer.record_tool_call('read_email', 300, 100, 'gpt-4o-mini')
analyzer.record_tool_call('generate_report', 2000, 500, 'gpt-4o')
analyzer.report()Integrating Cost Tracking Into the Agent Loop
Wrap your LLM call function to automatically track cost as part of the agent loop. Pass the tracker around so every call contributes to the session total.
import openai
client = openai.OpenAI(api_key='sk-...')
def tracked_completion(tracker: CostTracker, step_name: str, messages: list, model: str = 'gpt-4o-mini') -> str:
response = client.chat.completions.create(
model=model,
messages=messages
)
usage = response.usage
tracker.record(
step_name=step_name,
model=model,
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens
)
return response.choices[0].message.content
def run_agent_with_cost_tracking(question: str) -> dict:
tracker = CostTracker()
# Step 1: Entity extraction
entities_str = tracked_completion(
tracker, 'entity_extraction',
[{'role': 'user', 'content': f'Extract entities from: {question}'}]
)
# Step 2: Answer generation
answer = tracked_completion(
tracker, 'answer_generation',
[{'role': 'user', 'content': question}]
)
return {
'answer': answer,
'cost_summary': tracker.summary(),
'total_cost_usd': tracker.total_cost
}Estimating Tokens Before Calling
Use tiktoken to estimate token counts before making API calls. This lets you enforce budget limits and detect unexpectedly large prompts early.
import tiktoken
DEFAULT_ENCODER = tiktoken.encoding_for_model('gpt-4o-mini')
def estimate_tokens(text: str, model: str = 'gpt-4o-mini') -> int:
try:
encoding = tiktoken.encoding_for_model(model)
except KeyError:
encoding = DEFAULT_ENCODER
return len(encoding.encode(text))
def check_prompt_budget(messages: list, max_tokens: int = 8000) -> dict:
total = 0
breakdown = []
for msg in messages:
count = estimate_tokens(msg.get('content', ''))
total += count
breakdown.append({'role': msg['role'], 'tokens': count})
return {
'total_tokens': total,
'within_budget': total <= max_tokens,
'budget': max_tokens,
'breakdown': breakdown
}
messages = [
{'role': 'system', 'content': 'You are a helpful assistant that...'},
{'role': 'user', 'content': 'Explain the concept of quantum entanglement in simple terms.'}
]
result = check_prompt_budget(messages)
print(f'Total tokens: {result["total_tokens"]}, Within budget: {result["within_budget"]}')Cost Budgets and Cutoffs
Protect against runaway agent costs by setting per-run and per-day budget limits. If a run exceeds its budget, abort gracefully with a partial result rather than continuing to spend.
class BudgetGuard:
def __init__(self, max_cost_per_run: float = 0.10, max_cost_per_day: float = 5.00):
self.max_run = max_cost_per_run
self.max_day = max_cost_per_day
self.day_spend = 0.0
def check_and_spend(self, tracker: 'CostTracker', about_to_spend_estimate: float = 0.001):
if tracker.total_cost >= self.max_run:
raise RuntimeError(
f'Run budget exceeded: ${tracker.total_cost:.4f} >= ${self.max_run}'
)
if self.day_spend + tracker.total_cost >= self.max_day:
raise RuntimeError(
f'Daily budget exceeded: ${self.day_spend:.4f} daily spend'
)
def finalize_run(self, tracker: 'CostTracker'):
self.day_spend += tracker.total_cost
print(f'Run cost: ${tracker.total_cost:.6f}, Day total: ${self.day_spend:.4f}')
guard = BudgetGuard(max_cost_per_run=0.05, max_cost_per_day=2.00)
tracker = CostTracker()
tracker.record('test_step', 'gpt-4o-mini', 100, 50)
guard.check_and_spend(tracker)
guard.finalize_run(tracker)Storing Cost Data for Analysis
Persist run cost data to a database for trend analysis, billing attribution, and optimization decisions. A simple SQLite table works well for most agents.
import sqlite3
from datetime import datetime
def init_cost_db(db_path: str = 'agent_costs.db'):
conn = sqlite3.connect(db_path)
conn.execute('''
CREATE TABLE IF NOT EXISTS run_costs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL,
step_name TEXT NOT NULL,
model TEXT NOT NULL,
prompt_tokens INTEGER,
completion_tokens INTEGER,
cost_usd REAL,
timestamp TEXT
)
''')
conn.commit()
return conn
def save_run_costs(conn, run_id: str, tracker: 'CostTracker'):
for step in tracker.steps:
conn.execute(
'INSERT INTO run_costs (run_id, step_name, model, prompt_tokens, completion_tokens, cost_usd, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?)',
(run_id, step.step_name, step.model, step.prompt_tokens, step.completion_tokens, step.cost_usd, datetime.utcnow().isoformat())
)
conn.commit()
print(f'Saved {len(tracker.steps)} cost records for run {run_id}')
conn = init_cost_db()
print('Cost database initialized')Token Usage Alerts
Alert when a single agent run exceeds expected token usage. An unexpected spike often indicates a bug: infinitely growing context, repeated tool calls, or missing truncation logic.
def check_token_spike(tracker: 'CostTracker', expected_max_tokens: int = 10000) -> dict:
total = tracker.total_tokens
if total > expected_max_tokens:
# Find the biggest steps
sorted_steps = sorted(tracker.steps, key=lambda s: s.prompt_tokens + s.completion_tokens, reverse=True)
top_steps = [
{'step': s.step_name, 'tokens': s.prompt_tokens + s.completion_tokens}
for s in sorted_steps[:3]
]
message = (
f'Token spike: {total} tokens (expected <= {expected_max_tokens}). '
f'Top consumers: {top_steps}'
)
print(f'ALERT: {message}')
return {'alert': True, 'total_tokens': total, 'message': message, 'top_steps': top_steps}
return {'alert': False, 'total_tokens': total}
tracker = CostTracker()
tracker.record('context_builder', 'gpt-4o-mini', 8000, 200) # Unusually large prompt
result = check_token_spike(tracker, expected_max_tokens=5000)
print('Spike check:', result['alert'], '-', result.get('message', 'OK'))Cost Reporting Query
Query the cost database to generate reports: daily spend by model, most expensive steps, and cost trends over time. This guides optimization decisions.
import sqlite3
from datetime import datetime, timedelta
def cost_report(db_path: str = 'agent_costs.db', days: int = 7) -> dict:
conn = sqlite3.connect(db_path)
since = (datetime.utcnow() - timedelta(days=days)).isoformat()
# Total cost by model
model_costs = conn.execute('''
SELECT model, SUM(cost_usd) as total_cost, COUNT(*) as call_count
FROM run_costs WHERE timestamp >= ?
GROUP BY model ORDER BY total_cost DESC
''', (since,)).fetchall()
# Top expensive steps
step_costs = conn.execute('''
SELECT step_name, SUM(cost_usd) as total_cost, AVG(cost_usd) as avg_cost
FROM run_costs WHERE timestamp >= ?
GROUP BY step_name ORDER BY total_cost DESC LIMIT 10
''', (since,)).fetchall()
conn.close()
return {
'period_days': days,
'by_model': [{'model': r[0], 'total_usd': r[1], 'calls': r[2]} for r in model_costs],
'by_step': [{'step': r[0], 'total_usd': r[1], 'avg_usd': r[2]} for r in step_costs]
}
print('Cost report function defined')Knowledge Check: Token and Cost Profiling
Test your understanding of per-step token and cost profiling.
Cost Profiling Summary
Effective cost profiling requires: capturing usage from every API response, calculating cost per step using model pricing tables, tracking cumulative cost per agent run, breaking down cost by tool call type, enforcing budget limits with guard checks, and persisting cost data for trend analysis and optimization.
Frequently asked questions
Is the “Per-Step Token and Cost Profiling” lesson free?
Yes — the full text of “Per-Step Token and Cost Profiling” is free to read here on the web, and the AI Agents course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Agents course, upgrade to CoddyKit PRO.
What will I learn in “Per-Step Token and Cost Profiling”?
Measuring token consumption per tool call and per reasoning step. You practise AI Agents with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start AI Agents?
No prior experience is required. AI Agents on CoddyKit is structured for beginners through advanced learners; this is — lesson 2 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “Per-Step Token and Cost Profiling” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this AI Agents lesson?
Yes. Every AI Agents lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.
All lessons in this course
- Trace Analysis with LangSmith and Langfuse
- Per-Step Token and Cost Profiling
- Identifying Slow and Expensive Steps
- Root Cause Analysis for Agent Failures