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