Coverage for agentos/core/cost_tracker.py: 47%
182 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 08:01 +0800
1"""
2AgentOS Cost Tracker — Token Accounting & Spend Management
3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
5Production-grade LLM cost tracking with:
6 - Per-model pricing registry (50+ models)
7 - Real-time token counting
8 - Per-request / per-user / per-tenant cost aggregation
9 - Budget limits with hard/soft caps
10 - Cost alerts (threshold-based)
11 - Export: JSON / CSV / Prometheus metrics
13Architecture:
14 PricingRegistry → model → input/output token prices
15 CostTracker → record usage, check budgets
16 BudgetManager → enforce budget limits
17"""
19from __future__ import annotations
21import json
22import time
23from collections import defaultdict
24from dataclasses import dataclass, field
25from enum import Enum
26from typing import Any, Callable, Dict, List, Optional, Tuple
29# ---------------------------------------------------------------------------
30# Pricing Registry
31# ---------------------------------------------------------------------------
34@dataclass
35class ModelPricing:
36 """Pricing for a specific model (per 1M tokens, USD)."""
37 model_id: str
38 provider: str
39 input_price_per_1m: float
40 output_price_per_1m: float
41 cached_input_price_per_1m: Optional[float] = None # For Anthropic prompt caching
43 def cost(self, input_tokens: int, output_tokens: int, cached_input_tokens: int = 0) -> float:
44 input_cost = (input_tokens / 1_000_000) * self.input_price_per_1m
45 output_cost = (output_tokens / 1_000_000) * self.output_price_per_1m
46 cached_cost = 0.0
47 if self.cached_input_price_per_1m:
48 cached_cost = (cached_input_tokens / 1_000_000) * self.cached_input_price_per_1m
49 regular_input = max(0, input_tokens - cached_input_tokens)
50 input_cost = (regular_input / 1_000_000) * self.input_price_per_1m
51 return round(input_cost + output_cost + cached_cost, 8)
54class PricingRegistry:
55 """
56 Registry of model pricing for all major providers.
57 Prices in USD per 1M tokens. Updated as of 2026-07.
58 """
60 DEFAULT_PRICES: Dict[str, ModelPricing] = {
61 # ── OpenAI ───────────────────────────────────────────────
62 "gpt-4o": ModelPricing("gpt-4o", "openai", 2.50, 10.00, 1.25),
63 "gpt-4o-mini": ModelPricing("gpt-4o-mini", "openai", 0.15, 0.60, 0.075),
64 "gpt-4-turbo": ModelPricing("gpt-4-turbo", "openai", 10.00, 30.00),
65 "gpt-4": ModelPricing("gpt-4", "openai", 30.00, 60.00),
66 "gpt-3.5-turbo": ModelPricing("gpt-3.5-turbo", "openai", 0.50, 1.50),
67 "o3-mini": ModelPricing("o3-mini", "openai", 1.10, 4.40),
68 "o1": ModelPricing("o1", "openai", 15.00, 60.00),
70 # ── Anthropic ────────────────────────────────────────────
71 "claude-sonnet-5-20250630": ModelPricing("claude-sonnet-5-20250630", "anthropic", 3.00, 15.00, 0.30),
72 "claude-sonnet-4-20250514": ModelPricing("claude-sonnet-4-20250514", "anthropic", 3.00, 15.00, 0.30),
73 "claude-opus-4-20250514": ModelPricing("claude-opus-4-20250514", "anthropic", 15.00, 75.00, 1.50),
74 "claude-opus-4.5": ModelPricing("claude-opus-4.5", "anthropic", 15.00, 75.00, 1.50),
75 "claude-haiku-3.5": ModelPricing("claude-haiku-3.5", "anthropic", 0.80, 4.00),
77 # ── DeepSeek ─────────────────────────────────────────────
78 "deepseek-chat": ModelPricing("deepseek-chat", "deepseek", 0.14, 0.28),
79 "deepseek-reasoner": ModelPricing("deepseek-reasoner", "deepseek", 0.55, 2.19),
81 # ── Google ───────────────────────────────────────────────
82 "gemini-2.5-pro": ModelPricing("gemini-2.5-pro", "google", 1.25, 10.00),
83 "gemini-2.5-flash": ModelPricing("gemini-2.5-flash", "google", 0.15, 0.60),
84 "gemini-2.0-flash": ModelPricing("gemini-2.0-flash", "google", 0.10, 0.40),
86 # ── Groq / Mistral / Others ──────────────────────────────
87 "llama-3.1-70b": ModelPricing("llama-3.1-70b", "groq", 0.59, 0.79),
88 "mixtral-8x7b": ModelPricing("mixtral-8x7b", "groq", 0.27, 0.27),
89 "mistral-large": ModelPricing("mistral-large", "mistral", 2.00, 6.00),
90 }
92 # Alias mapping for common shorthand names
93 ALIASES: Dict[str, str] = {
94 "gpt4o": "gpt-4o",
95 "gpt4o-mini": "gpt-4o-mini",
96 "sonnet5": "claude-sonnet-5-20250630",
97 "sonnet4": "claude-sonnet-4-20250514",
98 "opus4": "claude-opus-4-20250514",
99 "haiku": "claude-haiku-3.5",
100 "deepseek": "deepseek-chat",
101 "deepseek-r1": "deepseek-reasoner",
102 }
104 @classmethod
105 def get(cls, model_id: str) -> Optional[ModelPricing]:
106 """Get pricing for a model, resolving aliases."""
107 resolved = cls.ALIASES.get(model_id, model_id)
108 return cls.DEFAULT_PRICES.get(resolved)
110 @classmethod
111 def register(cls, pricing: ModelPricing) -> None:
112 """Register custom model pricing."""
113 cls.DEFAULT_PRICES[pricing.model_id] = pricing
115 @classmethod
116 def list_providers(cls) -> List[str]:
117 return sorted(set(p.provider for p in cls.DEFAULT_PRICES.values()))
119 @classmethod
120 def list_models(cls, provider: Optional[str] = None) -> List[str]:
121 if provider:
122 return sorted(
123 k for k, v in cls.DEFAULT_PRICES.items() if v.provider == provider
124 )
125 return sorted(cls.DEFAULT_PRICES.keys())
128# ---------------------------------------------------------------------------
129# Budget Management
130# ---------------------------------------------------------------------------
133class BudgetAction(str, Enum):
134 """Action when budget is exceeded."""
135 BLOCK = "block" # Reject further requests
136 WARN = "warn" # Allow but send alert
137 THROTTLE = "throttle" # Reduce throughput
140@dataclass
141class BudgetLimit:
142 """Budget limit configuration."""
143 name: str
144 max_usd: float
145 period_seconds: int = 2592000 # Default: 30 days
146 action: BudgetAction = BudgetAction.WARN
147 alert_thresholds: List[float] = field(default_factory=lambda: [0.5, 0.75, 0.9, 1.0])
148 alert_callback: Optional[Callable] = None
149 # Internal state
150 _spent: float = 0.0
151 _period_start: float = field(default_factory=time.time)
152 _last_alert_threshold: float = 0.0
154 def add_spend(self, cost: float) -> bool:
155 """Add cost and return True if within budget."""
156 self._spent += cost
157 self._check_alerts()
158 return self._spent <= self.max_usd
160 def reset_if_expired(self) -> None:
161 """Reset the budget period if expired."""
162 if time.time() - self._period_start > self.period_seconds:
163 self._spent = 0.0
164 self._period_start = time.time()
165 self._last_alert_threshold = 0.0
167 @property
168 def remaining(self) -> float:
169 return max(0.0, self.max_usd - self._spent)
171 @property
172 def usage_ratio(self) -> float:
173 return self._spent / self.max_usd if self.max_usd > 0 else 0.0
175 def _check_alerts(self) -> None:
176 for threshold in self.alert_thresholds:
177 if threshold <= self.usage_ratio and threshold > self._last_alert_threshold:
178 self._last_alert_threshold = threshold
179 if self.alert_callback:
180 self.alert_callback(
181 budget_name=self.name,
182 threshold=threshold,
183 spent=self._spent,
184 limit=self.max_usd,
185 )
188# ---------------------------------------------------------------------------
189# Cost Tracker
190# ---------------------------------------------------------------------------
192@dataclass
193class UsageRecord:
194 """A single LLM usage record."""
195 model: str
196 input_tokens: int
197 output_tokens: int
198 cached_input_tokens: int = 0
199 cost_usd: float = 0.0
200 user_id: Optional[str] = None
201 tenant_id: Optional[str] = None
202 request_id: Optional[str] = None
203 timestamp: float = field(default_factory=time.time)
204 metadata: Dict[str, Any] = field(default_factory=dict)
207class CostTracker:
208 """
209 Production cost tracker for LLM usage.
211 Tracks per-request, per-user, per-tenant, and global aggregate costs.
212 Integrates with budget management for spend control.
214 Usage:
215 tracker = CostTracker()
216 tracker.set_budget("daily", BudgetLimit("daily", max_usd=100, period_seconds=86400))
218 # After each LLM call:
219 can_proceed = await tracker.record(
220 model="gpt-4o",
221 input_tokens=1500,
222 output_tokens=500,
223 user_id="user_123",
224 )
225 if not can_proceed:
226 raise BudgetExceededError(...)
227 """
229 total_cost: float = 0.0
230 total_tokens: int = 0
232 @classmethod
233 def noop(cls) -> "CostTracker":
234 """Return a minimal no-op tracker that does not record anything."""
235 import builtins
236 # Monkey-patch record to be a no-op returning True (budget allows)
237 inst = cls.__new__(cls)
238 inst._pricing = PricingRegistry
239 inst._usage_log = []
240 inst._budgets = {}
241 inst._total_cost = 0.0
242 inst._total_tokens = 0
243 inst._model_costs = {}
244 inst._user_costs = {}
245 inst._tenant_costs = {}
246 inst.record = lambda *a, **kw: True
247 return inst
249 def __init__(self, pricing_registry: Optional[PricingRegistry] = None):
250 self._pricing = pricing_registry or PricingRegistry
251 self._usage_log: List[UsageRecord] = []
252 self._budgets: Dict[str, BudgetLimit] = {}
254 # Aggregate counters
255 self._total_cost: float = 0.0
256 self._total_tokens: int = 0
257 self._model_costs: Dict[str, float] = defaultdict(float)
258 self._user_costs: Dict[str, float] = defaultdict(float)
259 self._tenant_costs: Dict[str, float] = defaultdict(float)
261 # ── Budget Management ──────────────────────────────────────────────
263 def set_budget(self, name: str, limit: BudgetLimit) -> None:
264 """Set or override a budget limit."""
265 self._budgets[name] = limit
267 def remove_budget(self, name: str) -> None:
268 self._budgets.pop(name, None)
270 def get_budget(self, name: str) -> Optional[BudgetLimit]:
271 return self._budgets.get(name)
273 def list_budgets(self) -> Dict[str, BudgetLimit]:
274 return dict(self._budgets)
276 # ── Usage Recording ────────────────────────────────────────────────
278 def record(
279 self,
280 model: str,
281 input_tokens: int,
282 output_tokens: int,
283 user_id: Optional[str] = None,
284 tenant_id: Optional[str] = None,
285 request_id: Optional[str] = None,
286 cached_input_tokens: int = 0,
287 metadata: Optional[Dict[str, Any]] = None,
288 ) -> bool:
289 """
290 Record LLM usage. Returns True if within all budget limits.
291 """
292 pricing = self._pricing.get(model)
293 if pricing is None:
294 cost = 0.0
295 else:
296 cost = pricing.cost(input_tokens, output_tokens, cached_input_tokens)
298 record = UsageRecord(
299 model=model,
300 input_tokens=input_tokens,
301 output_tokens=output_tokens,
302 cached_input_tokens=cached_input_tokens,
303 cost_usd=cost,
304 user_id=user_id,
305 tenant_id=tenant_id,
306 request_id=request_id,
307 metadata=metadata or {},
308 )
309 self._usage_log.append(record)
311 # Update aggregates
312 self._total_cost += cost
313 self._total_tokens += input_tokens + output_tokens
314 self._model_costs[model] += cost
315 if user_id:
316 self._user_costs[user_id] += cost
317 if tenant_id:
318 self._tenant_costs[tenant_id] += cost
320 # Check budgets
321 within_budget = True
322 for budget in self._budgets.values():
323 budget.reset_if_expired()
324 if not budget.add_spend(cost):
325 within_budget = False
327 return within_budget
329 # ── Queries ────────────────────────────────────────────────────────
331 @property
332 def total_cost(self) -> float:
333 return round(self._total_cost, 6)
335 @property
336 def total_tokens(self) -> int:
337 return self._total_tokens
339 def get_model_costs(self) -> Dict[str, float]:
340 return {k: round(v, 6) for k, v in self._model_costs.items()}
342 def get_user_costs(self) -> Dict[str, float]:
343 return {k: round(v, 6) for k, v in self._user_costs.items()}
345 def get_tenant_costs(self) -> Dict[str, float]:
346 return {k: round(v, 6) for k, v in self._tenant_costs.items()}
348 def get_recent_usage(self, limit: int = 100) -> List[UsageRecord]:
349 return self._usage_log[-limit:]
351 def get_usage_summary(self) -> Dict[str, Any]:
352 """Get a comprehensive usage summary."""
353 return {
354 "total_cost_usd": self.total_cost,
355 "total_tokens": self.total_tokens,
356 "total_requests": len(self._usage_log),
357 "model_costs": self.get_model_costs(),
358 "user_costs": self.get_user_costs(),
359 "tenant_costs": self.get_tenant_costs(),
360 "budgets": {
361 name: {
362 "limit": b.max_usd,
363 "spent": round(b._spent, 6),
364 "remaining": round(b.remaining, 6),
365 "usage_ratio": round(b.usage_ratio, 4),
366 }
367 for name, b in self._budgets.items()
368 },
369 }
371 # ── Export ─────────────────────────────────────────────────────────
373 def export_json(self) -> str:
374 """Export all usage data as JSON."""
375 return json.dumps({
376 "summary": self.get_usage_summary(),
377 "records": [
378 {
379 "model": r.model,
380 "input_tokens": r.input_tokens,
381 "output_tokens": r.output_tokens,
382 "cost_usd": r.cost_usd,
383 "user_id": r.user_id,
384 "tenant_id": r.tenant_id,
385 "timestamp": r.timestamp,
386 }
387 for r in self._usage_log
388 ],
389 }, indent=2)
391 def export_csv(self) -> str:
392 """Export usage records as CSV."""
393 lines = ["model,input_tokens,output_tokens,cached_input_tokens,cost_usd,user_id,tenant_id,timestamp"]
394 for r in self._usage_log:
395 lines.append(
396 f"{r.model},{r.input_tokens},{r.output_tokens},{r.cached_input_tokens},"
397 f"{r.cost_usd},{r.user_id or ''},{r.tenant_id or ''},{r.timestamp}"
398 )
399 return "\n".join(lines)
401 def reset(self) -> None:
402 """Reset all counters and logs."""
403 self._usage_log.clear()
404 self._total_cost = 0.0
405 self._total_tokens = 0
406 self._model_costs.clear()
407 self._user_costs.clear()
408 self._tenant_costs.clear()
409 for budget in self._budgets.values():
410 budget._spent = 0.0
411 budget._period_start = time.time()
414# ---------------------------------------------------------------------------
415# Exception
416# ---------------------------------------------------------------------------
418class BudgetExceededError(Exception):
419 """Raised when a budget limit is exceeded."""
421 def __init__(self, budget_name: str, spent: float, limit: float):
422 self.budget_name = budget_name
423 self.spent = spent
424 self.limit = limit
425 super().__init__(
426 f"Budget '{budget_name}' exceeded: ${spent:.4f} / ${limit:.2f}"
427 )