Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/pricing/manager.py: 28%

1"""Pricing manager for LLM models with builder pattern. 

2 

3This module provides a safe, intuitive API for configuring pricing sources 

4with validation and factory methods for common use cases. 

5 

6Example: 

7 >>> # Simple usage (95% of cases) 

8 >>> pricing = PricingManager.from_defaults() 

9 >>> model_pricing = await pricing.get_pricing("gpt-4-turbo") 

10 >>> 

11 >>> # Custom sources 

12 >>> pricing = ( 

13 ... PricingManager.builder() 

14 ... .add_json_source("custom.json") 

15 ... .add_api_source("https://api.example.com/pricing") 

16 ... .with_cache_ttl(3600) 

17 ... .build() 

18 ... ) 

19 

20""" 

21 

22from __future__ import annotations 

23 

24import asyncio 

25from collections.abc import Sequence 

26from datetime import UTC, datetime 

27from pathlib import Path 

28 

29from lexigram.ai.llm.pricing.sources import ( 

30 AbstractPricingSource, 

31 APIPricingSource, 

32 JSONFilePricingSource, 

33 StaticPricingSource, 

34) 

35from lexigram.ai.llm.pricing.types import ModelPricing 

36from lexigram.logging import ( 

37 get_logger, 

38) 

39 

40logger = get_logger(__name__) 

41 

42 

43class PricingCache: 

44 """Cache for pricing data with TTL. 

45 

46 Attributes: 

47 ttl: Time-to-live in seconds. 

48 cache: In-memory cache storage. 

49 

50 """ 

51 

52 def __init__(self, ttl: int = 86400): 

53 """Initialize pricing cache. 

54 

55 Args: 

56 ttl: Cache TTL in seconds (default: 24 hours). 

57 

58 """ 

59 self.ttl = ttl 

60 self._cache: dict[str, ModelPricing] = {} 

61 self._lock = asyncio.Lock() 

62 

63 async def get(self, model: str) -> ModelPricing | None: 

64 """Get pricing from cache. 

65 

66 Args: 

67 model: Model identifier. 

68 

69 Returns: 

70 Cached pricing if valid, None otherwise. 

71 

72 """ 

73 async with self._lock: 

74 if model not in self._cache: 

75 return None 

76 

77 pricing = self._cache[model] 

78 age = (datetime.now(UTC) - pricing.last_updated).total_seconds() 

79 

80 if age > self.ttl: 

81 del self._cache[model] 

82 return None 

83 

84 return pricing 

85 

86 async def set(self, model: str, pricing: ModelPricing) -> None: 

87 """Set pricing in cache. 

88 

89 Args: 

90 model: Model identifier. 

91 pricing: Pricing data. 

92 

93 """ 

94 async with self._lock: 

95 self._cache[model] = pricing 

96 

97 async def clear(self) -> None: 

98 """Clear all cached pricing.""" 

99 async with self._lock: 

100 self._cache.clear() 

101 

102 

103class PricingManager: 

104 """Manages pricing data from multiple sources with caching. 

105 

106 Sources are queried in order until pricing is found. Typical hierarchy: 

107 1. JSON file (fastest, most reliable) 

108 2. API endpoints (for updates) 

109 3. Static fallback (hardcoded) 

110 

111 Attributes: 

112 sources: List of pricing sources in priority order. 

113 cache: Pricing cache instance. 

114 enable_fuzzy_match: Whether to enable fuzzy model name matching. 

115 

116 Example: 

117 >>> # Use defaults 

118 >>> manager = PricingManager.from_defaults() 

119 >>> 

120 >>> # Custom configuration 

121 >>> manager = ( 

122 ... PricingManager.builder() 

123 ... .add_json_source("pricing.json") 

124 ... .add_api_source("https://api.example.com/pricing") 

125 ... .with_cache_ttl(3600) 

126 ... .enable_fuzzy_matching() 

127 ... .build() 

128 ... ) 

129 >>> 

130 >>> pricing = await manager.get_pricing("gpt-4-turbo") 

131 

132 """ 

133 

134 def __init__( 

135 self, 

136 sources: Sequence[AbstractPricingSource], 

137 cache_ttl: int = 86400, 

138 enable_fuzzy_match: bool = True, 

139 ): 

140 """Initialize pricing manager. 

141 

142 Args: 

143 sources: List of pricing sources in priority order. 

144 cache_ttl: Cache TTL in seconds (default: 24 hours). 

145 enable_fuzzy_match: Enable fuzzy model name matching (default: True). 

146 

147 """ 

148 if not sources: 

149 msg = "At least one pricing source is required" 

150 raise ValueError(msg) 

151 

152 self.sources = list(sources) 

153 self.cache = PricingCache(ttl=cache_ttl) 

154 self.enable_fuzzy_match = enable_fuzzy_match 

155 

156 logger.info( 

157 "PricingManager initialized with %d sources: %s", 

158 len(sources), 

159 [s.source_name for s in sources], 

160 ) 

161 

162 async def get_pricing( 

163 self, 

164 model: str, 

165 force_refresh: bool = False, 

166 ) -> ModelPricing | None: 

167 """Get pricing for a specific model. 

168 

169 Queries sources in order: 

170 1. Cache (if not force_refresh) 

171 2. Each source in priority order 

172 3. Fuzzy match if enabled 

173 4. None if not found 

174 

175 Args: 

176 model: Model identifier (e.g., "gpt-4-turbo"). 

177 force_refresh: Bypass cache and fetch fresh data. 

178 

179 Returns: 

180 ModelPricing if found, None otherwise. 

181 

182 """ 

183 model_normalized = model.lower().strip() 

184 

185 # Check cache first 

186 if not force_refresh: 

187 cached = await self.cache.get(model_normalized) 

188 if cached: 

189 logger.debug("Cache hit for %s", model) 

190 return cached 

191 

192 # Query sources in order 

193 for source in self.sources: 

194 pricing = await source.get_pricing(model_normalized) 

195 if pricing: 

196 logger.debug("Found pricing for %s in %s", model, source.source_name) 

197 await self.cache.set(model_normalized, pricing) 

198 return pricing 

199 

200 # Try fuzzy matching 

201 if self.enable_fuzzy_match: 

202 fuzzy_match = await self._fuzzy_match(model_normalized) 

203 if fuzzy_match: 

204 logger.info("Fuzzy matched %s to %s", model, fuzzy_match.model) 

205 await self.cache.set(model_normalized, fuzzy_match) 

206 return fuzzy_match 

207 

208 # Unknown model: report None (callers skip cost tracking) rather 

209 # than fabricate a price — mirrors PricingCostEstimator's 0.0 policy. 

210 logger.warning("No pricing found for %s", model) 

211 return None 

212 

213 async def _fuzzy_match(self, model: str) -> ModelPricing | None: 

214 """Try to fuzzy match model name. 

215 

216 Args: 

217 model: Normalized model name. 

218 

219 Returns: 

220 Matched pricing or None. 

221 

222 """ 

223 # Get all pricing from all sources 

224 for source in self.sources: 

225 all_pricing = await source.get_all_pricing() 

226 

227 # Try substring matching 

228 for known_model, pricing in all_pricing.items(): 

229 if model in known_model or known_model in model: 

230 return pricing 

231 

232 return None 

233 

234 async def list_models(self, provider: str | None = None) -> list[str]: 

235 """List all available models. 

236 

237 Args: 

238 provider: Filter by provider (optional). 

239 

240 Returns: 

241 List of model names. 

242 

243 """ 

244 all_models = set() 

245 

246 for source in self.sources: 

247 all_pricing = await source.get_all_pricing() 

248 for pricing in all_pricing.values(): 

249 if provider is None or pricing.provider == provider: 

250 all_models.add(pricing.model) 

251 

252 return sorted(all_models) 

253 

254 async def clear_cache(self) -> None: 

255 """Clear pricing cache.""" 

256 await self.cache.clear() 

257 logger.info("Pricing cache cleared") 

258 

259 async def preload(self) -> dict[str, ModelPricing]: 

260 """Load all pricing from all sources into one merged map. 

261 

262 Earlier sources win on duplicate model names, mirroring 

263 :meth:`get_pricing` priority semantics. Used to build the 

264 synchronous snapshot for cost estimators. 

265 

266 Returns: 

267 Merged dictionary of model name to pricing. 

268 """ 

269 merged: dict[str, ModelPricing] = {} 

270 for source in self.sources: 

271 all_pricing = await source.get_all_pricing() 

272 for model_name, pricing in all_pricing.items(): 

273 merged.setdefault(model_name, pricing) 

274 return merged 

275 

276 @classmethod 

277 def from_defaults(cls) -> PricingManager: 

278 """Create manager with default configuration. 

279 

280 Uses LiteLLM API for dynamic, up-to-date pricing data. 

281 No static pricing files - always fetches current data. 

282 

283 Returns: 

284 PricingManager with API source. 

285 

286 Example: 

287 >>> manager = PricingManager.from_defaults() 

288 >>> pricing = await manager.get_pricing("gpt-4") 

289 

290 """ 

291 sources = [ 

292 APIPricingSource( 

293 "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json", 

294 ), 

295 ] 

296 

297 return cls(sources=sources, cache_ttl=86400, enable_fuzzy_match=True) 

298 

299 @classmethod 

300 def from_json( 

301 cls, 

302 file_path: str | Path, 

303 cache_ttl: int = 86400, 

304 ) -> PricingManager: 

305 """Create manager from JSON file only. 

306 

307 Useful for offline applications or when you want full control 

308 over pricing data. 

309 

310 Args: 

311 file_path: Path to JSON pricing file. 

312 cache_ttl: Cache TTL in seconds (default: 24 hours). 

313 

314 Returns: 

315 PricingManager with JSON source only. 

316 

317 Example: 

318 >>> manager = PricingManager.from_json("my_pricing.json") 

319 >>> pricing = await manager.get_pricing("custom-model") 

320 

321 """ 

322 if isinstance(file_path, str): 

323 file_path = Path(file_path) 

324 

325 sources = [JSONFilePricingSource(file_path)] 

326 return cls(sources=sources, cache_ttl=cache_ttl, enable_fuzzy_match=True) 

327 

328 @classmethod 

329 def from_api(cls, endpoint: str, cache_ttl: int = 86400) -> PricingManager: 

330 """Create manager from API endpoint only. 

331 

332 Args: 

333 endpoint: API endpoint URL. 

334 cache_ttl: Cache TTL in seconds (default: 24 hours). 

335 

336 Returns: 

337 PricingManager with API source only. 

338 

339 Example: 

340 >>> manager = PricingManager.from_api("https://api.example.com/pricing") 

341 >>> pricing = await manager.get_pricing("gpt-4") 

342 

343 """ 

344 sources = [APIPricingSource(endpoint)] 

345 return cls(sources=sources, cache_ttl=cache_ttl, enable_fuzzy_match=True) 

346 

347 @classmethod 

348 def builder(cls) -> PricingManagerBuilder: 

349 """Create a builder for custom configuration. 

350 

351 Returns: 

352 PricingManagerBuilder instance. 

353 

354 Example: 

355 >>> manager = ( 

356 ... PricingManager.builder() 

357 ... .add_json_source("custom.json") 

358 ... .add_api_source("https://api.example.com") 

359 ... .with_cache_ttl(3600) 

360 ... .build() 

361 ... ) 

362 

363 """ 

364 return PricingManagerBuilder() 

365 

366 

367class PricingManagerBuilder: 

368 """Builder for PricingManager with validation. 

369 

370 Provides a fluent API for configuring pricing sources safely. 

371 

372 Example: 

373 >>> manager = ( 

374 ... PricingManager.builder() 

375 ... .add_json_source("pricing.json") 

376 ... .add_api_source("https://api.example.com/pricing") 

377 ... .add_fallback({"custom-model": ModelPricing(...)}) 

378 ... .with_cache_ttl(3600) 

379 ... .enable_fuzzy_matching() 

380 ... .build() 

381 ... ) 

382 

383 """ 

384 

385 def __init__(self) -> None: 

386 """Initialize builder.""" 

387 self._sources: list[AbstractPricingSource] = [] 

388 self._cache_ttl: int = 86400 

389 self._enable_fuzzy_match: bool = True 

390 

391 def add_json_source(self, file_path: str | Path) -> PricingManagerBuilder: 

392 """Add JSON file pricing source. 

393 

394 Args: 

395 file_path: Path to JSON file. 

396 

397 Returns: 

398 Self for chaining. 

399 

400 """ 

401 if isinstance(file_path, str): 

402 file_path = Path(file_path) 

403 

404 self._sources.append(JSONFilePricingSource(file_path)) 

405 return self 

406 

407 def add_api_source( 

408 self, 

409 endpoint: str, 

410 timeout: float = 10.0, 

411 ) -> PricingManagerBuilder: 

412 """Add API endpoint pricing source. 

413 

414 Args: 

415 endpoint: API endpoint URL. 

416 timeout: Request timeout in seconds (default: 10). 

417 

418 Returns: 

419 Self for chaining. 

420 

421 """ 

422 self._sources.append(APIPricingSource(endpoint, timeout)) 

423 return self 

424 

425 def add_fallback( 

426 self, 

427 pricing_map: dict[str, ModelPricing], 

428 ) -> PricingManagerBuilder: 

429 """Add static fallback pricing. 

430 

431 Args: 

432 pricing_map: Dictionary of model to pricing. 

433 

434 Returns: 

435 Self for chaining. 

436 

437 """ 

438 self._sources.append(StaticPricingSource(pricing_map)) 

439 return self 

440 

441 def add_source(self, source: AbstractPricingSource) -> PricingManagerBuilder: 

442 """Add custom pricing source. 

443 

444 Args: 

445 source: Custom AbstractPricingSource implementation. 

446 

447 Returns: 

448 Self for chaining. 

449 

450 """ 

451 self._sources.append(source) 

452 return self 

453 

454 def with_cache_ttl(self, seconds: int) -> PricingManagerBuilder: 

455 """Set cache TTL. 

456 

457 Args: 

458 seconds: Cache TTL in seconds. 

459 

460 Returns: 

461 Self for chaining. 

462 

463 Raises: 

464 ValueError: If seconds is negative. 

465 

466 """ 

467 if seconds < 0: 

468 msg = "cache TTL must be non-negative" 

469 raise ValueError(msg) 

470 

471 self._cache_ttl = seconds 

472 return self 

473 

474 def enable_fuzzy_matching(self, enabled: bool = True) -> PricingManagerBuilder: 

475 """Enable or disable fuzzy model name matching. 

476 

477 Args: 

478 enabled: Whether to enable fuzzy matching (default: True). 

479 

480 Returns: 

481 Self for chaining. 

482 

483 """ 

484 self._enable_fuzzy_match = enabled 

485 return self 

486 

487 def build(self) -> PricingManager: 

488 """Build PricingManager instance. 

489 

490 Returns: 

491 Configured PricingManager. 

492 

493 Raises: 

494 ValueError: If no sources were added. 

495 

496 """ 

497 if not self._sources: 

498 raise ValueError( 

499 "At least one pricing source must be added. " 

500 "Use add_json_source(), add_api_source(), or add_fallback()", 

501 ) 

502 

503 return PricingManager( 

504 sources=self._sources, 

505 cache_ttl=self._cache_ttl, 

506 enable_fuzzy_match=self._enable_fuzzy_match, 

507 )