Coverage for agentos/core/cost_tracker.py: 100%
182 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +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 collections.abc import Callable
25from dataclasses import dataclass, field
26from enum import StrEnum
27from typing import Any
29# ---------------------------------------------------------------------------
30# Pricing Registry
31# ---------------------------------------------------------------------------
34@dataclass
35class ModelPricing:
36 """Pricing for a specific model (per 1M tokens, USD)."""
38 model_id: str
39 provider: str
40 input_price_per_1m: float
41 output_price_per_1m: float
42 cached_input_price_per_1m: float | None = None # For Anthropic prompt caching
44 def cost(self, input_tokens: int, output_tokens: int, cached_input_tokens: int = 0) -> float:
45 input_cost = (input_tokens / 1_000_000) * self.input_price_per_1m
46 output_cost = (output_tokens / 1_000_000) * self.output_price_per_1m
47 cached_cost = 0.0
48 if self.cached_input_price_per_1m:
49 cached_cost = (cached_input_tokens / 1_000_000) * self.cached_input_price_per_1m
50 regular_input = max(0, input_tokens - cached_input_tokens)
51 input_cost = (regular_input / 1_000_000) * self.input_price_per_1m
52 return round(input_cost + output_cost + cached_cost, 8)
55class PricingRegistry:
56 """
57 Registry of model pricing for all major providers.
58 Prices in USD per 1M tokens. Updated as of 2026-07.
59 """
61 DEFAULT_PRICES: dict[str, ModelPricing] = {
62 # ── OpenAI ───────────────────────────────────────────────
63 "gpt-4o": ModelPricing("gpt-4o", "openai", 2.50, 10.00, 1.25),
64 "gpt-4o-mini": ModelPricing("gpt-4o-mini", "openai", 0.15, 0.60, 0.075),
65 "gpt-4-turbo": ModelPricing("gpt-4-turbo", "openai", 10.00, 30.00),
66 "gpt-4": ModelPricing("gpt-4", "openai", 30.00, 60.00),
67 "gpt-3.5-turbo": ModelPricing("gpt-3.5-turbo", "openai", 0.50, 1.50),
68 "o3-mini": ModelPricing("o3-mini", "openai", 1.10, 4.40),
69 "o1": ModelPricing("o1", "openai", 15.00, 60.00),
70 # ── Anthropic ────────────────────────────────────────────
71 "claude-sonnet-5-20250630": ModelPricing(
72 "claude-sonnet-5-20250630", "anthropic", 3.00, 15.00, 0.30
73 ),
74 "claude-sonnet-4-20250514": ModelPricing(
75 "claude-sonnet-4-20250514", "anthropic", 3.00, 15.00, 0.30
76 ),
77 "claude-opus-4-20250514": ModelPricing(
78 "claude-opus-4-20250514", "anthropic", 15.00, 75.00, 1.50
79 ),
80 "claude-opus-4.5": ModelPricing("claude-opus-4.5", "anthropic", 15.00, 75.00, 1.50),
81 "claude-haiku-3.5": ModelPricing("claude-haiku-3.5", "anthropic", 0.80, 4.00),
82 # ── DeepSeek ─────────────────────────────────────────────
83 "deepseek-chat": ModelPricing("deepseek-chat", "deepseek", 0.14, 0.28),
84 "deepseek-reasoner": ModelPricing("deepseek-reasoner", "deepseek", 0.55, 2.19),
85 # ── Google ───────────────────────────────────────────────
86 "gemini-2.5-pro": ModelPricing("gemini-2.5-pro", "google", 1.25, 10.00),
87 "gemini-2.5-flash": ModelPricing("gemini-2.5-flash", "google", 0.15, 0.60),
88 "gemini-2.0-flash": ModelPricing("gemini-2.0-flash", "google", 0.10, 0.40),
89 # ── Groq / Mistral / Others ──────────────────────────────
90 "llama-3.1-70b": ModelPricing("llama-3.1-70b", "groq", 0.59, 0.79),
91 "mixtral-8x7b": ModelPricing("mixtral-8x7b", "groq", 0.27, 0.27),
92 "mistral-large": ModelPricing("mistral-large", "mistral", 2.00, 6.00),
93 }
95 # Alias mapping for common shorthand names
96 ALIASES: dict[str, str] = {
97 "gpt4o": "gpt-4o",
98 "gpt4o-mini": "gpt-4o-mini",
99 "sonnet5": "claude-sonnet-5-20250630",
100 "sonnet4": "claude-sonnet-4-20250514",
101 "opus4": "claude-opus-4-20250514",
102 "haiku": "claude-haiku-3.5",
103 "deepseek": "deepseek-chat",
104 "deepseek-r1": "deepseek-reasoner",
105 }
107 @classmethod
108 def get(cls, model_id: str) -> ModelPricing | None:
109 """Get pricing for a model, resolving aliases."""
110 resolved = cls.ALIASES.get(model_id, model_id)
111 return cls.DEFAULT_PRICES.get(resolved)
113 @classmethod
114 def register(cls, pricing: ModelPricing) -> None:
115 """Register custom model pricing."""
116 cls.DEFAULT_PRICES[pricing.model_id] = pricing
118 @classmethod
119 def list_providers(cls) -> list[str]:
120 return sorted(set(p.provider for p in cls.DEFAULT_PRICES.values()))
122 @classmethod
123 def list_models(cls, provider: str | None = None) -> list[str]:
124 if provider:
125 return sorted(k for k, v in cls.DEFAULT_PRICES.items() if v.provider == provider)
126 return sorted(cls.DEFAULT_PRICES.keys())
129# ---------------------------------------------------------------------------
130# Budget Management
131# ---------------------------------------------------------------------------
134class BudgetAction(StrEnum):
135 """Action when budget is exceeded."""
137 BLOCK = "block" # Reject further requests
138 WARN = "warn" # Allow but send alert
139 THROTTLE = "throttle" # Reduce throughput
142@dataclass
143class BudgetLimit:
144 """Budget limit configuration."""
146 name: str
147 max_usd: float
148 period_seconds: int = 2592000 # Default: 30 days
149 action: BudgetAction = BudgetAction.WARN
150 alert_thresholds: list[float] = field(default_factory=lambda: [0.5, 0.75, 0.9, 1.0])
151 alert_callback: Callable | None = None
152 # Internal state
153 _spent: float = 0.0
154 _period_start: float = field(default_factory=time.time)
155 _last_alert_threshold: float = 0.0
157 def add_spend(self, cost: float) -> bool:
158 """Add cost and return True if within budget."""
159 self._spent += cost
160 self._check_alerts()
161 return self._spent <= self.max_usd
163 def reset_if_expired(self) -> None:
164 """Reset the budget period if expired."""
165 if time.time() - self._period_start > self.period_seconds:
166 self._spent = 0.0
167 self._period_start = time.time()
168 self._last_alert_threshold = 0.0
170 @property
171 def remaining(self) -> float:
172 return max(0.0, self.max_usd - self._spent)
174 @property
175 def usage_ratio(self) -> float:
176 return self._spent / self.max_usd if self.max_usd > 0 else 0.0
178 def _check_alerts(self) -> None:
179 for threshold in self.alert_thresholds:
180 if threshold <= self.usage_ratio and threshold > self._last_alert_threshold:
181 self._last_alert_threshold = threshold
182 if self.alert_callback:
183 self.alert_callback(
184 budget_name=self.name,
185 threshold=threshold,
186 spent=self._spent,
187 limit=self.max_usd,
188 )
191# ---------------------------------------------------------------------------
192# Cost Tracker
193# ---------------------------------------------------------------------------
196@dataclass
197class UsageRecord:
198 """A single LLM usage record."""
200 model: str
201 input_tokens: int
202 output_tokens: int
203 cached_input_tokens: int = 0
204 cost_usd: float = 0.0
205 user_id: str | None = None
206 tenant_id: str | None = None
207 request_id: str | None = None
208 timestamp: float = field(default_factory=time.time)
209 metadata: dict[str, Any] = field(default_factory=dict)
212class CostTracker:
213 """
214 Production cost tracker for LLM usage.
216 Tracks per-request, per-user, per-tenant, and global aggregate costs.
217 Integrates with budget management for spend control.
219 Usage:
220 tracker = CostTracker()
221 tracker.set_budget("daily", BudgetLimit("daily", max_usd=100, period_seconds=86400))
223 # After each LLM call:
224 can_proceed = await tracker.record(
225 model="gpt-4o",
226 input_tokens=1500,
227 output_tokens=500,
228 user_id="user_123",
229 )
230 if not can_proceed:
231 raise BudgetExceededError(...)
232 """
234 total_cost: float = 0.0
235 total_tokens: int = 0
237 @classmethod
238 def noop(cls) -> CostTracker:
239 """Return a minimal no-op tracker that does not record anything."""
240 # Monkey-patch record to be a no-op returning True (budget allows)
241 inst = cls.__new__(cls)
242 inst._pricing = PricingRegistry
243 inst._usage_log = []
244 inst._budgets = {}
245 inst._total_cost = 0.0
246 inst._total_tokens = 0
247 inst._model_costs = {}
248 inst._user_costs = {}
249 inst._tenant_costs = {}
250 inst.record = lambda *a, **kw: True
251 return inst
253 def __init__(self, pricing_registry: PricingRegistry | None = None):
254 self._pricing = pricing_registry or PricingRegistry
255 self._usage_log: list[UsageRecord] = []
256 self._budgets: dict[str, BudgetLimit] = {}
258 # Aggregate counters
259 self._total_cost: float = 0.0
260 self._total_tokens: int = 0
261 self._model_costs: dict[str, float] = defaultdict(float)
262 self._user_costs: dict[str, float] = defaultdict(float)
263 self._tenant_costs: dict[str, float] = defaultdict(float)
265 # ── Budget Management ──────────────────────────────────────────────
267 def set_budget(self, name: str, limit: BudgetLimit) -> None:
268 """Set or override a budget limit."""
269 self._budgets[name] = limit
271 def remove_budget(self, name: str) -> None:
272 self._budgets.pop(name, None)
274 def get_budget(self, name: str) -> BudgetLimit | None:
275 return self._budgets.get(name)
277 def list_budgets(self) -> dict[str, BudgetLimit]:
278 return dict(self._budgets)
280 # ── Usage Recording ────────────────────────────────────────────────
282 def record(
283 self,
284 model: str,
285 input_tokens: int,
286 output_tokens: int,
287 user_id: str | None = None,
288 tenant_id: str | None = None,
289 request_id: str | None = None,
290 cached_input_tokens: int = 0,
291 metadata: dict[str, Any] | None = None,
292 ) -> bool:
293 """
294 Record LLM usage. Returns True if within all budget limits.
295 """
296 pricing = self._pricing.get(model)
297 if pricing is None:
298 cost = 0.0
299 else:
300 cost = pricing.cost(input_tokens, output_tokens, cached_input_tokens)
302 record = UsageRecord(
303 model=model,
304 input_tokens=input_tokens,
305 output_tokens=output_tokens,
306 cached_input_tokens=cached_input_tokens,
307 cost_usd=cost,
308 user_id=user_id,
309 tenant_id=tenant_id,
310 request_id=request_id,
311 metadata=metadata or {},
312 )
313 self._usage_log.append(record)
315 # Update aggregates
316 self._total_cost += cost
317 self._total_tokens += input_tokens + output_tokens
318 self._model_costs[model] += cost
319 if user_id:
320 self._user_costs[user_id] += cost
321 if tenant_id:
322 self._tenant_costs[tenant_id] += cost
324 # Check budgets
325 within_budget = True
326 for budget in self._budgets.values():
327 budget.reset_if_expired()
328 if not budget.add_spend(cost):
329 within_budget = False
331 return within_budget
333 # ── Queries ────────────────────────────────────────────────────────
335 @property
336 def total_cost(self) -> float:
337 return round(self._total_cost, 6)
339 @property
340 def total_tokens(self) -> int:
341 return self._total_tokens
343 def get_model_costs(self) -> dict[str, float]:
344 return {k: round(v, 6) for k, v in self._model_costs.items()}
346 def get_user_costs(self) -> dict[str, float]:
347 return {k: round(v, 6) for k, v in self._user_costs.items()}
349 def get_tenant_costs(self) -> dict[str, float]:
350 return {k: round(v, 6) for k, v in self._tenant_costs.items()}
352 def get_recent_usage(self, limit: int = 100) -> list[UsageRecord]:
353 return self._usage_log[-limit:]
355 def get_usage_summary(self) -> dict[str, Any]:
356 """Get a comprehensive usage summary."""
357 return {
358 "total_cost_usd": self.total_cost,
359 "total_tokens": self.total_tokens,
360 "total_requests": len(self._usage_log),
361 "model_costs": self.get_model_costs(),
362 "user_costs": self.get_user_costs(),
363 "tenant_costs": self.get_tenant_costs(),
364 "budgets": {
365 name: {
366 "limit": b.max_usd,
367 "spent": round(b._spent, 6),
368 "remaining": round(b.remaining, 6),
369 "usage_ratio": round(b.usage_ratio, 4),
370 }
371 for name, b in self._budgets.items()
372 },
373 }
375 # ── Export ─────────────────────────────────────────────────────────
377 def export_json(self) -> str:
378 """Export all usage data as JSON."""
379 return json.dumps(
380 {
381 "summary": self.get_usage_summary(),
382 "records": [
383 {
384 "model": r.model,
385 "input_tokens": r.input_tokens,
386 "output_tokens": r.output_tokens,
387 "cost_usd": r.cost_usd,
388 "user_id": r.user_id,
389 "tenant_id": r.tenant_id,
390 "timestamp": r.timestamp,
391 }
392 for r in self._usage_log
393 ],
394 },
395 indent=2,
396 )
398 def export_csv(self) -> str:
399 """Export usage records as CSV."""
400 lines = [
401 "model,input_tokens,output_tokens,cached_input_tokens,cost_usd,user_id,tenant_id,timestamp"
402 ]
403 for r in self._usage_log:
404 lines.append(
405 f"{r.model},{r.input_tokens},{r.output_tokens},{r.cached_input_tokens},"
406 f"{r.cost_usd},{r.user_id or ''},{r.tenant_id or ''},{r.timestamp}"
407 )
408 return "\n".join(lines)
410 def reset(self) -> None:
411 """Reset all counters and logs."""
412 self._usage_log.clear()
413 self._total_cost = 0.0
414 self._total_tokens = 0
415 self._model_costs.clear()
416 self._user_costs.clear()
417 self._tenant_costs.clear()
418 for budget in self._budgets.values():
419 budget._spent = 0.0
420 budget._period_start = time.time()
423# ---------------------------------------------------------------------------
424# Exception
425# ---------------------------------------------------------------------------
428class BudgetExceededError(Exception):
429 """Raised when a budget limit is exceeded."""
431 def __init__(self, budget_name: str, spent: float, limit: float):
432 self.budget_name = budget_name
433 self.spent = spent
434 self.limit = limit
435 super().__init__(f"Budget '{budget_name}' exceeded: ${spent:.4f} / ${limit:.2f}")