Coverage for agentos/cost/tracker.py: 71%
169 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 17:45 +0800
1""" # noqa: E501
2v1.10.0: Cost Tracker — token counting & pricing across all providers.
4Tracks token usage and cost for: OpenAI, Anthropic, Google, DeepSeek, Groq.
5Features: per-request tracking, budget management, usage reporting.
6"""
8from __future__ import annotations
10from collections import defaultdict
11from dataclasses import dataclass
12from datetime import UTC, datetime
13from enum import StrEnum
14from typing import Any
16# ── Data Classes ──────────────────────────────────────────────────
19class ProviderPricing(StrEnum):
20 OPENAI = "openai"
21 ANTHROPIC = "anthropic"
22 GOOGLE = "google"
23 DEEPSEEK = "deepseek"
24 GROQ = "groq"
25 CUSTOM = "custom"
28@dataclass
29class TokenPricing:
30 """Pricing per 1M tokens (input/output)."""
32 provider: ProviderPricing
33 model: str
34 input_price_per_1m: float # USD per 1M input tokens
35 output_price_per_1m: float # USD per 1M output tokens
36 cache_write_price_per_1m: float = 0.0
37 cache_read_price_per_1m: float = 0.0
39 def cost(
40 self, input_tokens: int, output_tokens: int, cache_write: int = 0, cache_read: int = 0
41 ) -> float:
42 return (
43 (input_tokens / 1_000_000) * self.input_price_per_1m
44 + (output_tokens / 1_000_000) * self.output_price_per_1m
45 + (cache_write / 1_000_000) * self.cache_write_price_per_1m
46 + (cache_read / 1_000_000) * self.cache_read_price_per_1m
47 )
50@dataclass
51class TokenUsage:
52 """Token usage for a single API call."""
54 model: str
55 input_tokens: int = 0
56 output_tokens: int = 0
57 cache_write_tokens: int = 0
58 cache_read_tokens: int = 0
59 total_tokens: int = 0
60 cost: float = 0.0
61 latency_ms: float = 0.0
62 timestamp: str = ""
64 def __post_init__(self):
65 if not self.total_tokens:
66 self.total_tokens = self.input_tokens + self.output_tokens
67 if not self.timestamp:
68 self.timestamp = datetime.now(UTC).isoformat()
71@dataclass
72class Budget:
73 """Spending budget configuration."""
75 name: str
76 limit: float # USD
77 period: str = "monthly" # daily / weekly / monthly / total
78 current_spend: float = 0.0
79 alert_threshold: float = 0.8 # Alert at 80% of limit
80 hard_stop: bool = False # Block requests when exceeded
82 @property
83 def remaining(self) -> float:
84 return max(0.0, self.limit - self.current_spend)
86 @property
87 def pct_used(self) -> float:
88 return (self.current_spend / self.limit * 100) if self.limit > 0 else 0.0
90 @property
91 def exceeded(self) -> bool:
92 return self.current_spend >= self.limit
94 @property
95 def should_alert(self) -> bool:
96 return self.pct_used >= self.alert_threshold * 100
99# ── Default Pricing (as of 2025-07) ────────────────────────────────
101DEFAULT_PRICING: dict[str, TokenPricing] = {
102 # OpenAI
103 "gpt-4o": TokenPricing(ProviderPricing.OPENAI, "gpt-4o", 2.50, 10.00),
104 "gpt-4o-mini": TokenPricing(ProviderPricing.OPENAI, "gpt-4o-mini", 0.15, 0.60),
105 "gpt-4-turbo": TokenPricing(ProviderPricing.OPENAI, "gpt-4-turbo", 10.00, 30.00),
106 "gpt-3.5-turbo": TokenPricing(ProviderPricing.OPENAI, "gpt-3.5-turbo", 0.50, 1.50),
107 "o3-mini": TokenPricing(ProviderPricing.OPENAI, "o3-mini", 1.10, 4.40),
108 # Anthropic
109 "claude-3-5-sonnet": TokenPricing(
110 ProviderPricing.ANTHROPIC,
111 "claude-3-5-sonnet",
112 3.00,
113 15.00,
114 cache_write_price_per_1m=3.75,
115 cache_read_price_per_1m=0.30,
116 ),
117 "claude-3-haiku": TokenPricing(ProviderPricing.ANTHROPIC, "claude-3-haiku", 0.25, 1.25),
118 "claude-3-opus": TokenPricing(ProviderPricing.ANTHROPIC, "claude-3-opus", 15.00, 75.00),
119 # Google
120 "gemini-2.0-flash": TokenPricing(ProviderPricing.GOOGLE, "gemini-2.0-flash", 0.10, 0.40),
121 "gemini-2.0-pro": TokenPricing(ProviderPricing.GOOGLE, "gemini-2.0-pro", 1.25, 5.00),
122 "gemini-1.5-pro": TokenPricing(ProviderPricing.GOOGLE, "gemini-1.5-pro", 1.25, 5.00),
123 # DeepSeek
124 "deepseek-chat": TokenPricing(ProviderPricing.DEEPSEEK, "deepseek-chat", 0.27, 1.10),
125 "deepseek-reasoner": TokenPricing(ProviderPricing.DEEPSEEK, "deepseek-reasoner", 0.55, 2.19),
126 # Groq
127 "llama-3.3-70b": TokenPricing(ProviderPricing.GROQ, "llama-3.3-70b", 0.59, 0.79),
128 "mixtral-8x7b": TokenPricing(ProviderPricing.GROQ, "mixtral-8x7b", 0.24, 0.24),
129 "gemma2-9b-it": TokenPricing(ProviderPricing.GROQ, "gemma2-9b-it", 0.20, 0.20),
130}
133# ── Token Counter (heuristic-based, provider-agnostic) ────────────
136class TokenCounter:
137 """Approximate token counter based on word count + code heuristics.
139 For exact counts, use provider-specific tokenizers (tiktoken, etc.).
140 This provides fast, offline estimates within ~10% accuracy.
141 """
143 # Rough tokens-per-word ratios (language-dependent)
144 TOKENS_PER_WORD: dict[str, float] = {
145 "en": 1.3, # ~4 chars/token for English
146 "zh": 0.5, # ~2 chars/token for Chinese (character-based)
147 "ja": 0.6,
148 "ko": 0.6,
149 "code": 0.7, # Code tends to be denser in tokens per word
150 "default": 1.0,
151 }
153 @classmethod
154 def count(cls, text: str, source: str = "default") -> int:
155 """Estimate token count."""
156 if not text:
157 return 0
159 ratio = cls.TOKENS_PER_WORD.get(source, cls.TOKENS_PER_WORD["default"])
160 chars = len(text)
162 # For Chinese (high CJK ratio), use character-based estimation
163 cjk_chars = sum(1 for c in text if "\u4e00" <= c <= "\u9fff" or "\u3040" <= c <= "\u30ff")
164 cjk_ratio = cjk_chars / max(chars, 1)
166 if cjk_ratio > 0.3:
167 # Mostly Chinese/Japanese — use CJK character ratio
168 non_cjk = chars - cjk_chars
169 return int(cjk_chars * cls.TOKENS_PER_WORD["zh"] + non_cjk * 0.25)
171 if source == "code" or cls._is_code(text):
172 ratio = cls.TOKENS_PER_WORD["code"]
174 words = len(text.split())
175 return max(1, int(words * ratio))
177 @staticmethod
178 def _is_code(text: str) -> bool:
179 """Heuristic: detect if text is code."""
180 code_indicators = [
181 "def ",
182 "class ",
183 "import ",
184 "from ",
185 "function",
186 "const ",
187 "let ",
188 "var ",
189 "{",
190 "}",
191 "=>",
192 "return ",
193 ]
194 count = sum(1 for ind in code_indicators if ind in text)
195 return count >= 3
198# ── Cost Tracker ───────────────────────────────────────────────────
201class CostTracker:
202 """Track token usage and costs across all provider calls.
204 Usage:
205 tracker = CostTracker()
206 tracker.record("gpt-4o", input_tokens=500, output_tokens=200)
207 tracker.record("claude-3-5-sonnet", input_tokens=1000, output_tokens=500)
208 report = tracker.report()
209 """
211 total_cost: float = 0.0
212 total_tokens: int = 0
214 @classmethod
215 def noop(cls) -> CostTracker:
216 """Return a minimal no-op tracker that does not record anything."""
217 inst = cls.__new__(cls)
218 inst.pricing = {}
219 inst.budgets = {}
220 inst.usage_log = []
221 inst._model_totals = {}
222 inst.record = lambda *a, **kw: True
223 return inst
225 def __init__(
226 self,
227 custom_pricing: dict[str, TokenPricing] | None = None,
228 budgets: list[Budget] | None = None,
229 ):
230 self.pricing: dict[str, TokenPricing] = {**DEFAULT_PRICING}
231 if custom_pricing:
232 self.pricing.update(custom_pricing)
234 self.budgets: dict[str, Budget] = {}
235 if budgets:
236 for b in budgets:
237 self.budgets[b.name] = b
239 self.usage_log: list[TokenUsage] = []
240 self._model_totals: dict[str, dict[str, float]] = defaultdict(
241 lambda: {"input_tokens": 0, "output_tokens": 0, "cost": 0.0, "calls": 0}
242 )
244 def get_price(self, model: str) -> TokenPricing:
245 """Get pricing for a model. Falls back to default if unknown."""
246 if model in self.pricing:
247 return self.pricing[model]
249 # Best-effort fallback based on model name
250 if "gpt-4" in model:
251 return TokenPricing(ProviderPricing.OPENAI, model, 2.50, 10.00)
252 if "gpt-3" in model:
253 return TokenPricing(ProviderPricing.OPENAI, model, 0.50, 1.50)
254 if "claude" in model:
255 return TokenPricing(ProviderPricing.ANTHROPIC, model, 3.00, 15.00)
256 if "gemini" in model:
257 return TokenPricing(ProviderPricing.GOOGLE, model, 0.10, 0.40)
258 if "deepseek" in model:
259 return TokenPricing(ProviderPricing.DEEPSEEK, model, 0.27, 1.10)
260 if any(m in model for m in ["llama", "mixtral", "gemma"]):
261 return TokenPricing(ProviderPricing.GROQ, model, 0.20, 0.20)
263 return TokenPricing(ProviderPricing.CUSTOM, model, 1.00, 1.00)
265 def record(
266 self,
267 model: str,
268 input_tokens: int = 0,
269 output_tokens: int = 0,
270 cache_write_tokens: int = 0,
271 cache_read_tokens: int = 0,
272 latency_ms: float = 0.0,
273 ) -> TokenUsage:
274 """Record a token usage event. Returns the TokenUsage with cost."""
275 pricing = self.get_price(model)
276 cost = pricing.cost(input_tokens, output_tokens, cache_write_tokens, cache_read_tokens)
278 usage = TokenUsage(
279 model=model,
280 input_tokens=input_tokens,
281 output_tokens=output_tokens,
282 cache_write_tokens=cache_write_tokens,
283 cache_read_tokens=cache_read_tokens,
284 cost=cost,
285 latency_ms=latency_ms,
286 )
287 self.usage_log.append(usage)
289 # Update model totals
290 mt = self._model_totals[model]
291 mt["input_tokens"] += input_tokens
292 mt["output_tokens"] += output_tokens
293 mt["cost"] += cost
294 mt["calls"] += 1
296 # Update budgets
297 for budget in self.budgets.values():
298 budget.current_spend += cost
300 return usage
302 def check_budget(self) -> list[str]:
303 """Check all budgets. Returns list of alert messages."""
304 alerts = []
305 for budget in self.budgets.values():
306 if budget.exceeded and budget.hard_stop:
307 alerts.append(
308 f"BUDGET EXCEEDED: {budget.name} (${budget.current_spend:.2f}/${budget.limit:.2f})"
309 )
310 elif budget.should_alert:
311 alerts.append(
312 f"Budget alert: {budget.name} at {budget.pct_used:.0f}% (${budget.current_spend:.2f}/${budget.limit:.2f})" # noqa: E501
313 )
314 return alerts
316 def report(self) -> str:
317 """Generate a human-readable cost report."""
318 total_cost = sum(u.cost for u in self.usage_log)
319 total_tokens = sum(u.total_tokens for u in self.usage_log)
320 total_calls = len(self.usage_log)
322 lines = [
323 "╔══ Cost Report ══╗",
324 f"║ Total calls: {total_calls}",
325 f"║ Total tokens: {total_tokens:,}",
326 f"║ Total cost: ${total_cost:.4f}",
327 "╚════════════════╝",
328 "",
329 "By model:",
330 ]
331 for model, totals in sorted(self._model_totals.items(), key=lambda x: -x[1]["cost"]):
332 lines.append(
333 f" {model:<30} {totals['calls']:>4} calls "
334 f"{totals['input_tokens']+totals['output_tokens']:>12,} tokens "
335 f"${totals['cost']:>8.4f}"
336 )
338 if self.budgets:
339 lines.append("\nBudgets:")
340 for budget in self.budgets.values():
341 status = "EXCEEDED" if budget.exceeded else "OK"
342 lines.append(
343 f" {budget.name:<20} ${budget.current_spend:.2f}/${budget.limit:.2f} "
344 f"({budget.pct_used:.0f}%) [{status}]"
345 )
347 return "\n".join(lines)
349 def report_dict(self) -> dict[str, Any]:
350 """Generate a machine-readable cost report."""
351 return {
352 "total_calls": len(self.usage_log),
353 "total_tokens": sum(u.total_tokens for u in self.usage_log),
354 "total_cost": sum(u.cost for u in self.usage_log),
355 "by_model": {model: dict(totals) for model, totals in self._model_totals.items()},
356 "recent": [
357 {
358 "model": u.model,
359 "input_tokens": u.input_tokens,
360 "output_tokens": u.output_tokens,
361 "cost": u.cost,
362 "timestamp": u.timestamp,
363 }
364 for u in self.usage_log[-20:] # Last 20 calls
365 ],
366 }
368 def reset(self) -> None:
369 """Reset all counters (keeps pricing and budgets)."""
370 self.usage_log.clear()
371 self._model_totals.clear()
372 for budget in self.budgets.values():
373 budget.current_spend = 0.0
375 def set_budget(self, name: str, limit: float, hard_stop: bool = False) -> Budget:
376 """Create or update a budget."""
377 budget = Budget(name=name, limit=limit, hard_stop=hard_stop)
378 self.budgets[name] = budget
379 return budget
382# ── Backward Compatibility Aliases (v1.2.7-) ──────────────────────
383# Old names → new equivalents
384RunCostSession = CostTracker # CostTracker was RunCostSession
385ModelPricing = TokenPricing # ModelPricing → TokenPricing
386UsageRecord = TokenUsage # UsageRecord → TokenUsage
387PRICING = DEFAULT_PRICING # PRICING → DEFAULT_PRICING