Coverage for src/local_deep_research/metrics/pricing/pricing_cache.py: 100%

1""" 

2Pricing Cache System 

3 

4Caches pricing data to avoid repeated API calls and improve performance. 

5Uses bounded TTLCache to prevent memory leaks. 

6""" 

7 

8from typing import Any, Dict, Optional 

9 

10from cachetools import TTLCache # type: ignore[import-untyped] 

11from loguru import logger 

12 

13 

14class PricingCache: 

15 """Cache for LLM pricing data.""" 

16 

17 def __init__(self, cache_dir: Optional[str] = None, cache_ttl: int = 3600): 

18 """ 

19 Initialize pricing cache. 

20 

21 Args: 

22 cache_dir: Directory to store cache files (DEPRECATED - no longer used) 

23 cache_ttl: Cache time-to-live in seconds (default: 1 hour) 

24 """ 

25 self.cache_ttl = cache_ttl 

26 # Bounded TTLCache to prevent memory leaks 

27 self._cache: TTLCache = TTLCache(maxsize=500, ttl=cache_ttl) 

28 logger.info("PricingCache initialized with bounded TTLCache") 

29 

30 def get(self, key: str) -> Optional[Any]: 

31 """Get cached pricing data. TTLCache handles expiration automatically.""" 

32 return self._cache.get(key) 

33 

34 def set(self, key: str, data: Any): 

35 """Set cached pricing data.""" 

36 self._cache[key] = data 

37 

38 def get_model_pricing(self, model_name: str) -> Optional[Dict[str, float]]: 

39 """Get cached pricing for a specific model.""" 

40 return self.get(f"model:{model_name}") 

41 

42 def set_model_pricing(self, model_name: str, pricing: Dict[str, float]): 

43 """Cache pricing for a specific model.""" 

44 self.set(f"model:{model_name}", pricing) 

45 

46 def get_all_pricing(self) -> Optional[Dict[str, Dict[str, float]]]: 

47 """Get cached pricing for all models.""" 

48 return self.get("all_models") 

49 

50 def set_all_pricing(self, pricing: Dict[str, Dict[str, float]]): 

51 """Cache pricing for all models.""" 

52 self.set("all_models", pricing) 

53 

54 def clear(self): 

55 """Clear all cached data.""" 

56 self._cache.clear() 

57 logger.info("Pricing cache cleared") 

58 

59 def clear_expired(self): 

60 """Remove expired cache entries. TTLCache handles this automatically via expire().""" 

61 self._cache.expire() 

62 logger.debug("Expired cache entries cleared") 

63 

64 def get_cache_stats(self) -> Dict[str, Any]: 

65 """Get cache statistics.""" 

66 # TTLCache automatically evicts expired entries on access 

67 self._cache.expire() 

68 return { 

69 "total_entries": len(self._cache), 

70 "max_entries": self._cache.maxsize, 

71 "cache_type": "TTLCache", 

72 "cache_ttl": self.cache_ttl, 

73 }