Coverage for src / osiris_cli / cost_tracker.py: 0%
121 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-31 05:01 +0200
1"""
2Cost tracking and analytics for Osiris CLI
4Tracks token usage and costs across providers, provides
5cost transparency, budget alerts, and usage analytics.
6"""
8import json
9from pathlib import Path
10from datetime import datetime, timedelta
11from typing import Dict, List, Optional, Tuple
12from dataclasses import dataclass, asdict
13from collections import defaultdict
15from .logger import get_logger
16from .config import CONFIG_DIR
18logger = get_logger()
20# Cost tracking directory
21COST_DIR = CONFIG_DIR / "costs"
22COST_DIR.mkdir(parents=True, exist_ok=True)
24# Cost database file
25COST_DB = COST_DIR / "usage.json"
28# Provider pricing (per 1K tokens) - Updated Dec 2024
29PROVIDER_COSTS = {
30 "openai": {
31 "gpt-4": {"input": 0.03, "output": 0.06},
32 "gpt-4-turbo": {"input": 0.01, "output": 0.03},
33 "gpt-4-turbo-preview": {"input": 0.01, "output": 0.03},
34 "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
35 "gpt-3.5-turbo-16k": {"input": 0.001, "output": 0.002},
36 "o1-preview": {"input": 0.015, "output": 0.06},
37 "o1-mini": {"input": 0.003, "output": 0.012},
38 },
39 "anthropic": {
40 "claude-3-5-sonnet-20241022": {"input": 0.003, "output": 0.015},
41 "claude-3-5-sonnet-20240620": {"input": 0.003, "output": 0.015},
42 "claude-3-opus-20240229": {"input": 0.015, "output": 0.075},
43 "claude-3-sonnet-20240229": {"input": 0.003, "output": 0.015},
44 "claude-3-haiku-20240307": {"input": 0.00025, "output": 0.00125},
45 },
46 "google": {
47 "gemini-pro": {"input": 0.00025, "output": 0.0005},
48 "gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
49 "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
50 "gemini-2.0-flash": {"input": 0.0001, "output": 0.0004},
51 },
52 "groq": {
53 # Groq is free for many models
54 "llama3-8b-8192": {"input": 0.00005, "output": 0.00008},
55 "llama3-70b-8192": {"input": 0.00059, "output": 0.00079},
56 "mixtral-8x7b-32768": {"input": 0.00024, "output": 0.00024},
57 "gemma-7b-it": {"input": 0.00007, "output": 0.00007},
58 },
59 "deepseek": {
60 "deepseek-chat": {"input": 0.00014, "output": 0.00028},
61 "deepseek-coder": {"input": 0.00014, "output": 0.00028},
62 },
63 "perplexity": {
64 "sonar": {"input": 0.001, "output": 0.001},
65 "sonar-pro": {"input": 0.003, "output": 0.015},
66 },
67 "openrouter": {
68 # Common OpenRouter models
69 "openai/gpt-4": {"input": 0.03, "output": 0.06},
70 "openai/gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
71 "anthropic/claude-3-5-sonnet": {"input": 0.003, "output": 0.015},
72 "meta-llama/llama-3-70b": {"input": 0.00059, "output": 0.00079},
73 }
74}
77@dataclass
78class APICall:
79 """Record of a single API call"""
80 timestamp: str
81 provider: str
82 model: str
83 input_tokens: int
84 output_tokens: int
85 total_tokens: int
86 input_cost: float
87 output_cost: float
88 total_cost: float
89 duration_ms: float
90 session_id: Optional[str] = None
92 def to_dict(self) -> Dict:
93 return asdict(self)
96class CostTracker:
97 """
98 Track API costs and token usage.
100 Features:
101 - Per-call cost tracking
102 - Session-level aggregation
103 - Daily/weekly/monthly analytics
104 - Budget alerts
105 - Cost projections
106 - Export to CSV/JSON
107 """
109 def __init__(self):
110 self.calls: List[APICall] = []
111 self.load()
113 def load(self):
114 """Load cost history from disk"""
115 if COST_DB.exists():
116 try:
117 with open(COST_DB, 'r') as f:
118 data = json.load(f)
119 self.calls = [APICall(**call) for call in data.get("calls", [])]
120 logger.debug(f"Loaded {len(self.calls)} cost records")
121 except Exception as e:
122 logger.error(f"Failed to load cost database: {e}")
123 self.calls = []
125 def save(self):
126 """Save cost history to disk"""
127 try:
128 with open(COST_DB, 'w') as f:
129 json.dump(
130 {"calls": [call.to_dict() for call in self.calls]},
131 f,
132 indent=2
133 )
134 logger.debug(f"Saved {len(self.calls)} cost records")
135 except Exception as e:
136 logger.error(f"Failed to save cost database: {e}")
138 def calculate_cost(
139 self,
140 provider: str,
141 model: str,
142 input_tokens: int,
143 output_tokens: int
144 ) -> Tuple[float, float, float]:
145 """
146 Calculate cost for an API call.
148 Args:
149 provider: Provider name
150 model: Model name
151 input_tokens: Number of input tokens
152 output_tokens: Number of output tokens
154 Returns:
155 (input_cost, output_cost, total_cost) in USD
156 """
157 # Get pricing for provider/model
158 provider_pricing = PROVIDER_COSTS.get(provider.lower(), {})
159 model_pricing = provider_pricing.get(model, {})
161 # Default pricing if not found (use low estimates)
162 input_price_per_1k = model_pricing.get("input", 0.001)
163 output_price_per_1k = model_pricing.get("output", 0.002)
165 # Calculate costs
166 input_cost = (input_tokens / 1000) * input_price_per_1k
167 output_cost = (output_tokens / 1000) * output_price_per_1k
168 total_cost = input_cost + output_cost
170 return input_cost, output_cost, total_cost
172 def track_call(
173 self,
174 provider: str,
175 model: str,
176 input_tokens: int,
177 output_tokens: int,
178 duration_ms: float,
179 session_id: Optional[str] = None
180 ) -> APICall:
181 """
182 Track an API call.
184 Args:
185 provider: Provider name
186 model: Model name
187 input_tokens: Number of input tokens
188 output_tokens: Number of output tokens
189 duration_ms: Call duration in milliseconds
190 session_id: Optional session identifier
192 Returns:
193 APICall record
194 """
195 input_cost, output_cost, total_cost = self.calculate_cost(
196 provider, model, input_tokens, output_tokens
197 )
199 call = APICall(
200 timestamp=datetime.now().isoformat(),
201 provider=provider,
202 model=model,
203 input_tokens=input_tokens,
204 output_tokens=output_tokens,
205 total_tokens=input_tokens + output_tokens,
206 input_cost=input_cost,
207 output_cost=output_cost,
208 total_cost=total_cost,
209 duration_ms=duration_ms,
210 session_id=session_id
211 )
213 self.calls.append(call)
214 self.save()
216 # Log for analytics
217 logger.api_call(provider, model, call.total_tokens, total_cost, duration_ms)
219 return call
221 def get_session_costs(self, session_id: str) -> Dict:
222 """Get aggregated costs for a session"""
223 session_calls = [c for c in self.calls if c.session_id == session_id]
225 if not session_calls:
226 return {
227 "session_id": session_id,
228 "total_calls": 0,
229 "total_tokens": 0,
230 "total_cost": 0.0
231 }
233 return {
234 "session_id": session_id,
235 "total_calls": len(session_calls),
236 "total_tokens": sum(c.total_tokens for c in session_calls),
237 "input_tokens": sum(c.input_tokens for c in session_calls),
238 "output_tokens": sum(c.output_tokens for c in session_calls),
239 "total_cost": sum(c.total_cost for c in session_calls),
240 "providers": list(set(c.provider for c in session_calls)),
241 "models": list(set(c.model for c in session_calls)),
242 }
244 def get_period_costs(self, days: int = 30) -> Dict:
245 """Get aggregated costs for a time period"""
246 cutoff = datetime.now() - timedelta(days=days)
247 period_calls = [
248 c for c in self.calls
249 if datetime.fromisoformat(c.timestamp) > cutoff
250 ]
252 if not period_calls:
253 return {
254 "period_days": days,
255 "total_calls": 0,
256 "total_tokens": 0,
257 "total_cost": 0.0
258 }
260 # Aggregate by provider
261 by_provider = defaultdict(lambda: {"calls": 0, "tokens": 0, "cost": 0.0})
262 for call in period_calls:
263 by_provider[call.provider]["calls"] += 1
264 by_provider[call.provider]["tokens"] += call.total_tokens
265 by_provider[call.provider]["cost"] += call.total_cost
267 # Aggregate by model
268 by_model = defaultdict(lambda: {"calls": 0, "tokens": 0, "cost": 0.0})
269 for call in period_calls:
270 model_key = f"{call.provider}/{call.model}"
271 by_model[model_key]["calls"] += 1
272 by_model[model_key]["tokens"] += call.total_tokens
273 by_model[model_key]["cost"] += call.total_cost
275 return {
276 "period_days": days,
277 "start_date": cutoff.isoformat(),
278 "end_date": datetime.now().isoformat(),
279 "total_calls": len(period_calls),
280 "total_tokens": sum(c.total_tokens for c in period_calls),
281 "total_cost": sum(c.total_cost for c in period_calls),
282 "avg_cost_per_call": sum(c.total_cost for c in period_calls) / len(period_calls),
283 "by_provider": dict(by_provider),
284 "by_model": dict(by_model),
285 }
287 def get_daily_costs(self, days: int = 7) -> List[Dict]:
288 """Get daily cost breakdown"""
289 daily_costs = []
291 for i in range(days):
292 day = datetime.now() - timedelta(days=i)
293 day_start = day.replace(hour=0, minute=0, second=0, microsecond=0)
294 day_end = day_start + timedelta(days=1)
296 day_calls = [
297 c for c in self.calls
298 if day_start <= datetime.fromisoformat(c.timestamp) < day_end
299 ]
301 daily_costs.append({
302 "date": day_start.isoformat()[:10],
303 "calls": len(day_calls),
304 "tokens": sum(c.total_tokens for c in day_calls),
305 "cost": sum(c.total_cost for c in day_calls)
306 })
308 return list(reversed(daily_costs))
310 def estimate_cost(
311 self,
312 provider: str,
313 model: str,
314 estimated_tokens: int
315 ) -> float:
316 """
317 Estimate cost for a future API call.
319 Args:
320 provider: Provider name
321 model: Model name
322 estimated_tokens: Estimated total tokens
324 Returns:
325 Estimated cost in USD
326 """
327 # Assume 60/40 split (input/output)
328 input_tokens = int(estimated_tokens * 0.6)
329 output_tokens = int(estimated_tokens * 0.4)
331 _, _, total_cost = self.calculate_cost(
332 provider, model, input_tokens, output_tokens
333 )
335 return total_cost
337 def check_budget(self, budget_usd: float, period_days: int = 30) -> Dict:
338 """
339 Check if usage is within budget.
341 Args:
342 budget_usd: Budget in USD
343 period_days: Period to check
345 Returns:
346 Budget status dict
347 """
348 period_costs = self.get_period_costs(period_days)
349 total_cost = period_costs["total_cost"]
351 return {
352 "budget": budget_usd,
353 "spent": total_cost,
354 "remaining": budget_usd - total_cost,
355 "percent_used": (total_cost / budget_usd * 100) if budget_usd > 0 else 0,
356 "over_budget": total_cost > budget_usd,
357 "period_days": period_days
358 }
360 def export_csv(self, output_path: Path):
361 """Export cost data to CSV"""
362 import csv
364 with open(output_path, 'w', newline='') as f:
365 if not self.calls:
366 return
368 fieldnames = list(self.calls[0].to_dict().keys())
369 writer = csv.DictWriter(f, fieldnames=fieldnames)
371 writer.writeheader()
372 for call in self.calls:
373 writer.writerow(call.to_dict())
375 logger.info(f"Exported {len(self.calls)} records to {output_path}")
377 def get_stats(self) -> Dict:
378 """Get overall statistics"""
379 if not self.calls:
380 return {
381 "total_calls": 0,
382 "total_tokens": 0,
383 "total_cost": 0.0
384 }
386 return {
387 "total_calls": len(self.calls),
388 "total_tokens": sum(c.total_tokens for c in self.calls),
389 "total_cost": sum(c.total_cost for c in self.calls),
390 "avg_tokens_per_call": sum(c.total_tokens for c in self.calls) / len(self.calls),
391 "avg_cost_per_call": sum(c.total_cost for c in self.calls) / len(self.calls),
392 "providers_used": list(set(c.provider for c in self.calls)),
393 "models_used": list(set(c.model for c in self.calls)),
394 "first_call": min(c.timestamp for c in self.calls),
395 "last_call": max(c.timestamp for c in self.calls),
396 }
399# Global cost tracker instance
400cost_tracker = CostTracker()