Coverage for agentos/core/cost_tracker.py: 100%
181 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +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
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 # Monkey-patch record to be a no-op returning True (budget allows)
236 inst = cls.__new__(cls)
237 inst._pricing = PricingRegistry
238 inst._usage_log = []
239 inst._budgets = {}
240 inst._total_cost = 0.0
241 inst._total_tokens = 0
242 inst._model_costs = {}
243 inst._user_costs = {}
244 inst._tenant_costs = {}
245 inst.record = lambda *a, **kw: True
246 return inst
248 def __init__(self, pricing_registry: Optional[PricingRegistry] = None):
249 self._pricing = pricing_registry or PricingRegistry
250 self._usage_log: List[UsageRecord] = []
251 self._budgets: Dict[str, BudgetLimit] = {}
253 # Aggregate counters
254 self._total_cost: float = 0.0
255 self._total_tokens: int = 0
256 self._model_costs: Dict[str, float] = defaultdict(float)
257 self._user_costs: Dict[str, float] = defaultdict(float)
258 self._tenant_costs: Dict[str, float] = defaultdict(float)
260 # ── Budget Management ──────────────────────────────────────────────
262 def set_budget(self, name: str, limit: BudgetLimit) -> None:
263 """Set or override a budget limit."""
264 self._budgets[name] = limit
266 def remove_budget(self, name: str) -> None:
267 self._budgets.pop(name, None)
269 def get_budget(self, name: str) -> Optional[BudgetLimit]:
270 return self._budgets.get(name)
272 def list_budgets(self) -> Dict[str, BudgetLimit]:
273 return dict(self._budgets)
275 # ── Usage Recording ────────────────────────────────────────────────
277 def record(
278 self,
279 model: str,
280 input_tokens: int,
281 output_tokens: int,
282 user_id: Optional[str] = None,
283 tenant_id: Optional[str] = None,
284 request_id: Optional[str] = None,
285 cached_input_tokens: int = 0,
286 metadata: Optional[Dict[str, Any]] = None,
287 ) -> bool:
288 """
289 Record LLM usage. Returns True if within all budget limits.
290 """
291 pricing = self._pricing.get(model)
292 if pricing is None:
293 cost = 0.0
294 else:
295 cost = pricing.cost(input_tokens, output_tokens, cached_input_tokens)
297 record = UsageRecord(
298 model=model,
299 input_tokens=input_tokens,
300 output_tokens=output_tokens,
301 cached_input_tokens=cached_input_tokens,
302 cost_usd=cost,
303 user_id=user_id,
304 tenant_id=tenant_id,
305 request_id=request_id,
306 metadata=metadata or {},
307 )
308 self._usage_log.append(record)
310 # Update aggregates
311 self._total_cost += cost
312 self._total_tokens += input_tokens + output_tokens
313 self._model_costs[model] += cost
314 if user_id:
315 self._user_costs[user_id] += cost
316 if tenant_id:
317 self._tenant_costs[tenant_id] += cost
319 # Check budgets
320 within_budget = True
321 for budget in self._budgets.values():
322 budget.reset_if_expired()
323 if not budget.add_spend(cost):
324 within_budget = False
326 return within_budget
328 # ── Queries ────────────────────────────────────────────────────────
330 @property
331 def total_cost(self) -> float:
332 return round(self._total_cost, 6)
334 @property
335 def total_tokens(self) -> int:
336 return self._total_tokens
338 def get_model_costs(self) -> Dict[str, float]:
339 return {k: round(v, 6) for k, v in self._model_costs.items()}
341 def get_user_costs(self) -> Dict[str, float]:
342 return {k: round(v, 6) for k, v in self._user_costs.items()}
344 def get_tenant_costs(self) -> Dict[str, float]:
345 return {k: round(v, 6) for k, v in self._tenant_costs.items()}
347 def get_recent_usage(self, limit: int = 100) -> List[UsageRecord]:
348 return self._usage_log[-limit:]
350 def get_usage_summary(self) -> Dict[str, Any]:
351 """Get a comprehensive usage summary."""
352 return {
353 "total_cost_usd": self.total_cost,
354 "total_tokens": self.total_tokens,
355 "total_requests": len(self._usage_log),
356 "model_costs": self.get_model_costs(),
357 "user_costs": self.get_user_costs(),
358 "tenant_costs": self.get_tenant_costs(),
359 "budgets": {
360 name: {
361 "limit": b.max_usd,
362 "spent": round(b._spent, 6),
363 "remaining": round(b.remaining, 6),
364 "usage_ratio": round(b.usage_ratio, 4),
365 }
366 for name, b in self._budgets.items()
367 },
368 }
370 # ── Export ─────────────────────────────────────────────────────────
372 def export_json(self) -> str:
373 """Export all usage data as JSON."""
374 return json.dumps({
375 "summary": self.get_usage_summary(),
376 "records": [
377 {
378 "model": r.model,
379 "input_tokens": r.input_tokens,
380 "output_tokens": r.output_tokens,
381 "cost_usd": r.cost_usd,
382 "user_id": r.user_id,
383 "tenant_id": r.tenant_id,
384 "timestamp": r.timestamp,
385 }
386 for r in self._usage_log
387 ],
388 }, indent=2)
390 def export_csv(self) -> str:
391 """Export usage records as CSV."""
392 lines = ["model,input_tokens,output_tokens,cached_input_tokens,cost_usd,user_id,tenant_id,timestamp"]
393 for r in self._usage_log:
394 lines.append(
395 f"{r.model},{r.input_tokens},{r.output_tokens},{r.cached_input_tokens},"
396 f"{r.cost_usd},{r.user_id or ''},{r.tenant_id or ''},{r.timestamp}"
397 )
398 return "\n".join(lines)
400 def reset(self) -> None:
401 """Reset all counters and logs."""
402 self._usage_log.clear()
403 self._total_cost = 0.0
404 self._total_tokens = 0
405 self._model_costs.clear()
406 self._user_costs.clear()
407 self._tenant_costs.clear()
408 for budget in self._budgets.values():
409 budget._spent = 0.0
410 budget._period_start = time.time()
413# ---------------------------------------------------------------------------
414# Exception
415# ---------------------------------------------------------------------------
417class BudgetExceededError(Exception):
418 """Raised when a budget limit is exceeded."""
420 def __init__(self, budget_name: str, spent: float, limit: float):
421 self.budget_name = budget_name
422 self.spent = spent
423 self.limit = limit
424 super().__init__(
425 f"Budget '{budget_name}' exceeded: ${spent:.4f} / ${limit:.2f}"
426 )