Coverage for src/local_deep_research/metrics/pricing/cost_calculator.py: 99%
94 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Cost Calculator
4Calculates LLM usage costs based on token usage and pricing data.
5Integrates with pricing fetcher and cache systems.
6"""
8from typing import Any, Dict, List, Optional
10from loguru import logger
12from .pricing_cache import PricingCache
13from .pricing_fetcher import PricingFetcher
16class CostCalculator:
17 """Calculates LLM usage costs."""
19 def __init__(self, cache_dir: Optional[str] = None):
20 self.cache = PricingCache(cache_dir)
21 self.pricing_fetcher = None
23 async def __aenter__(self):
24 self.pricing_fetcher = PricingFetcher()
25 await self.pricing_fetcher.__aenter__()
26 return self
28 async def __aexit__(self, exc_type, exc_val, exc_tb):
29 if self.pricing_fetcher: 29 ↛ exitline 29 didn't return from function '__aexit__' because the condition on line 29 was always true
30 await self.pricing_fetcher.__aexit__(exc_type, exc_val, exc_tb)
32 async def get_model_pricing(
33 self, model_name: str, provider: str = None
34 ) -> Dict[str, float]:
35 """Get pricing for a model and provider (cached or fetched)."""
36 # Create cache key that includes provider
37 cache_key = f"{provider}:{model_name}" if provider else model_name
39 # Try cache first
40 cached_pricing = self.cache.get(f"model:{cache_key}")
41 if cached_pricing:
42 return cached_pricing
44 # Fetch from API
45 if self.pricing_fetcher:
46 pricing = await self.pricing_fetcher.get_model_pricing(
47 model_name, provider
48 )
49 if pricing:
50 self.cache.set(f"model:{cache_key}", pricing)
51 return pricing
53 # No pricing found
54 logger.warning(
55 f"No pricing found for {model_name} (provider: {provider})"
56 )
57 return None
59 async def calculate_cost(
60 self,
61 model_name: str,
62 prompt_tokens: int,
63 completion_tokens: int,
64 provider: str = None,
65 ) -> Dict[str, float]:
66 """
67 Calculate cost for a single LLM call.
69 Returns:
70 Dict with prompt_cost, completion_cost, total_cost
71 """
72 pricing = await self.get_model_pricing(model_name, provider)
74 # If no pricing found, return zero cost
75 if pricing is None:
76 return {
77 "prompt_cost": 0.0,
78 "completion_cost": 0.0,
79 "total_cost": 0.0,
80 "pricing_used": None,
81 "error": "No pricing data available for this model",
82 }
84 # Convert tokens to thousands for pricing calculation
85 prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
86 completion_cost = (completion_tokens / 1000) * pricing["completion"]
87 total_cost = prompt_cost + completion_cost
89 return {
90 "prompt_cost": round(prompt_cost, 6),
91 "completion_cost": round(completion_cost, 6),
92 "total_cost": round(total_cost, 6),
93 "pricing_used": pricing,
94 }
96 async def calculate_batch_costs(
97 self, usage_records: List[Dict[str, Any]]
98 ) -> List[Dict[str, Any]]:
99 """
100 Calculate costs for multiple usage records.
102 Expected record format:
103 {
104 "model_name": str,
105 "provider": str (optional),
106 "prompt_tokens": int,
107 "completion_tokens": int,
108 "research_id": int (optional),
109 "timestamp": datetime (optional)
110 }
111 """
112 results = []
114 for record in usage_records:
115 try:
116 cost_data = await self.calculate_cost(
117 record["model_name"],
118 record["prompt_tokens"],
119 record["completion_tokens"],
120 record.get("provider"),
121 )
123 result = {**record, **cost_data}
124 results.append(result)
126 except Exception:
127 logger.exception(
128 f"Failed to calculate cost for record {record}"
129 )
130 # Add record with zero cost on error
131 results.append(
132 {
133 **record,
134 "prompt_cost": 0.0,
135 "completion_cost": 0.0,
136 "total_cost": 0.0,
137 "pricing_used": None,
138 "error": "Cost calculation failed",
139 }
140 )
142 return results
144 def calculate_cost_sync(
145 self,
146 model_name: str,
147 prompt_tokens: int,
148 completion_tokens: int,
149 provider: Optional[str] = None,
150 ) -> Dict[str, float]:
151 """
152 Synchronous cost calculation using provider-scoped cache entries and
153 the same static fallback as the async path. No network I/O is used.
154 """
155 # Use cached pricing only
156 cache_key = f"{provider}:{model_name}" if provider else model_name
157 pricing = self.cache.get_model_pricing(cache_key)
158 if not pricing:
159 fetcher = PricingFetcher()
160 pricing = fetcher.get_static_model_pricing(model_name, provider)
162 # If no pricing found, return zero cost
163 if not pricing:
164 return {
165 "prompt_cost": 0.0,
166 "completion_cost": 0.0,
167 "total_cost": 0.0,
168 "pricing_used": None,
169 "error": "No pricing data available for this model",
170 }
172 prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
173 completion_cost = (completion_tokens / 1000) * pricing["completion"]
174 total_cost = prompt_cost + completion_cost
176 return {
177 "prompt_cost": round(prompt_cost, 6),
178 "completion_cost": round(completion_cost, 6),
179 "total_cost": round(total_cost, 6),
180 "pricing_used": pricing,
181 }
183 async def get_research_cost_summary(
184 self, usage_records: List[Dict[str, Any]]
185 ) -> Dict[str, Any]:
186 """Get cost summary for research session(s).
188 ``model_breakdown`` retains totals grouped by model name.
189 ``provider_breakdown`` lists each distinct model/provider pair;
190 separate fields keep names containing separators unambiguous and
191 allow the summary to be serialized as JSON.
193 Cost totals cover only successfully priced usage; the existing
194 averages divide those totals by all calls/tokens. ``unpriced_calls``
195 and ``unpriced_tokens`` identify incomplete totals. A known zero
196 price (for example, Ollama) is fully priced.
197 """
198 costs = await self.calculate_batch_costs(usage_records)
199 return summarize_costs(costs)
202def summarize_costs(costs: List[Dict[str, Any]]) -> Dict[str, Any]:
203 """Aggregate calculated usage records from either pricing path."""
204 total_cost = sum(c["total_cost"] for c in costs)
205 total_prompt_cost = sum(c["prompt_cost"] for c in costs)
206 total_completion_cost = sum(c["completion_cost"] for c in costs)
208 total_prompt_tokens = sum(c["prompt_tokens"] for c in costs)
209 total_completion_tokens = sum(c["completion_tokens"] for c in costs)
210 total_tokens = total_prompt_tokens + total_completion_tokens
212 # Keep the existing model totals and a separate provider breakdown.
213 model_costs = {}
214 provider_costs = {}
215 unpriced_calls = 0
216 unpriced_tokens = 0
217 for cost in costs:
218 model = cost["model_name"]
219 if model not in model_costs:
220 model_costs[model] = {
221 "total_cost": 0.0,
222 "prompt_tokens": 0,
223 "completion_tokens": 0,
224 "calls": 0,
225 }
227 provider = cost.get("provider")
228 key = (model, provider)
229 if key not in provider_costs:
230 provider_costs[key] = {
231 "model_name": model,
232 "provider": provider,
233 "total_cost": 0.0,
234 "prompt_tokens": 0,
235 "completion_tokens": 0,
236 "calls": 0,
237 "unpriced_calls": 0,
238 "unpriced_tokens": 0,
239 }
241 for bucket in (model_costs[model], provider_costs[key]):
242 bucket["total_cost"] += cost["total_cost"]
243 bucket["prompt_tokens"] += cost["prompt_tokens"]
244 bucket["completion_tokens"] += cost["completion_tokens"]
245 bucket["calls"] += 1
247 if cost.get("pricing_used") is None:
248 tokens = cost["prompt_tokens"] + cost["completion_tokens"]
249 unpriced_calls += 1
250 unpriced_tokens += tokens
251 provider_costs[key]["unpriced_calls"] += 1
252 provider_costs[key]["unpriced_tokens"] += tokens
254 for bucket in provider_costs.values():
255 bucket["total_cost"] = round(bucket["total_cost"], 6)
257 return {
258 "total_cost": round(total_cost, 6),
259 "prompt_cost": round(total_prompt_cost, 6),
260 "completion_cost": round(total_completion_cost, 6),
261 "total_tokens": total_tokens,
262 "prompt_tokens": total_prompt_tokens,
263 "completion_tokens": total_completion_tokens,
264 "total_calls": len(costs),
265 "model_breakdown": model_costs,
266 "provider_breakdown": list(provider_costs.values()),
267 "pricing_complete": unpriced_calls == 0,
268 "unpriced_calls": unpriced_calls,
269 "unpriced_tokens": unpriced_tokens,
270 "avg_cost_per_call": round(total_cost / len(costs), 6)
271 if costs
272 else 0.0,
273 "cost_per_token": (
274 round(total_cost / total_tokens, 8) if total_tokens > 0 else 0.0
275 ),
276 }