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 

3 

4Tracks token usage and costs across providers, provides 

5cost transparency, budget alerts, and usage analytics. 

6""" 

7 

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 

14 

15from .logger import get_logger 

16from .config import CONFIG_DIR 

17 

18logger = get_logger() 

19 

20# Cost tracking directory 

21COST_DIR = CONFIG_DIR / "costs" 

22COST_DIR.mkdir(parents=True, exist_ok=True) 

23 

24# Cost database file 

25COST_DB = COST_DIR / "usage.json" 

26 

27 

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} 

75 

76 

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 

91 

92 def to_dict(self) -> Dict: 

93 return asdict(self) 

94 

95 

96class CostTracker: 

97 """ 

98 Track API costs and token usage. 

99  

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 """ 

108 

109 def __init__(self): 

110 self.calls: List[APICall] = [] 

111 self.load() 

112 

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 = [] 

124 

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}") 

137 

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. 

147  

148 Args: 

149 provider: Provider name 

150 model: Model name 

151 input_tokens: Number of input tokens 

152 output_tokens: Number of output tokens 

153  

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, {}) 

160 

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) 

164 

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 

169 

170 return input_cost, output_cost, total_cost 

171 

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. 

183  

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 

191  

192 Returns: 

193 APICall record 

194 """ 

195 input_cost, output_cost, total_cost = self.calculate_cost( 

196 provider, model, input_tokens, output_tokens 

197 ) 

198 

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 ) 

212 

213 self.calls.append(call) 

214 self.save() 

215 

216 # Log for analytics 

217 logger.api_call(provider, model, call.total_tokens, total_cost, duration_ms) 

218 

219 return call 

220 

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] 

224 

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 } 

232 

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 } 

243 

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 ] 

251 

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 } 

259 

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 

266 

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 

274 

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 } 

286 

287 def get_daily_costs(self, days: int = 7) -> List[Dict]: 

288 """Get daily cost breakdown""" 

289 daily_costs = [] 

290 

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) 

295 

296 day_calls = [ 

297 c for c in self.calls 

298 if day_start <= datetime.fromisoformat(c.timestamp) < day_end 

299 ] 

300 

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 }) 

307 

308 return list(reversed(daily_costs)) 

309 

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. 

318  

319 Args: 

320 provider: Provider name 

321 model: Model name 

322 estimated_tokens: Estimated total tokens 

323  

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) 

330 

331 _, _, total_cost = self.calculate_cost( 

332 provider, model, input_tokens, output_tokens 

333 ) 

334 

335 return total_cost 

336 

337 def check_budget(self, budget_usd: float, period_days: int = 30) -> Dict: 

338 """ 

339 Check if usage is within budget. 

340  

341 Args: 

342 budget_usd: Budget in USD 

343 period_days: Period to check 

344  

345 Returns: 

346 Budget status dict 

347 """ 

348 period_costs = self.get_period_costs(period_days) 

349 total_cost = period_costs["total_cost"] 

350 

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 } 

359 

360 def export_csv(self, output_path: Path): 

361 """Export cost data to CSV""" 

362 import csv 

363 

364 with open(output_path, 'w', newline='') as f: 

365 if not self.calls: 

366 return 

367 

368 fieldnames = list(self.calls[0].to_dict().keys()) 

369 writer = csv.DictWriter(f, fieldnames=fieldnames) 

370 

371 writer.writeheader() 

372 for call in self.calls: 

373 writer.writerow(call.to_dict()) 

374 

375 logger.info(f"Exported {len(self.calls)} records to {output_path}") 

376 

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 } 

385 

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 } 

397 

398 

399# Global cost tracker instance 

400cost_tracker = CostTracker()