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

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 _load_cache(self): 

31 """DEPRECATED: No longer loads from disk.""" 

32 pass 

33 

34 def _save_cache(self): 

35 """DEPRECATED: No longer saves to disk.""" 

36 pass 

37 

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

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

40 return self._cache.get(key) 

41 

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

43 """Set cached pricing data.""" 

44 self._cache[key] = data 

45 

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

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

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

49 

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

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

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

53 

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

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

56 return self.get("all_models") 

57 

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

59 """Cache pricing for all models.""" 

60 self.set("all_models", pricing) 

61 

62 def clear(self): 

63 """Clear all cached data.""" 

64 self._cache.clear() 

65 logger.info("Pricing cache cleared") 

66 

67 def clear_expired(self): 

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

69 self._cache.expire() 

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

71 

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

73 """Get cache statistics.""" 

74 # TTLCache automatically evicts expired entries on access 

75 self._cache.expire() 

76 return { 

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

78 "max_entries": self._cache.maxsize, 

79 "cache_type": "TTLCache", 

80 "cache_ttl": self.cache_ttl, 

81 }