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

1""" 

2Pricing Cache System 

3 

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

5Includes cache expiration and refresh mechanisms. 

6""" 

7 

8import time 

9from typing import Any, Dict, Optional 

10 

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 # In-memory cache only - no file operations 

27 self._cache = {} 

28 logger.info("PricingCache initialized with in-memory storage only") 

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 _is_expired(self, timestamp: float) -> bool: 

39 """Check if cache entry is expired.""" 

40 return (time.time() - timestamp) > self.cache_ttl 

41 

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

43 """Get cached pricing data.""" 

44 if key not in self._cache: 

45 return None 

46 

47 entry = self._cache[key] 

48 if self._is_expired(entry["timestamp"]): 

49 # Remove expired entry 

50 del self._cache[key] 

51 return None 

52 

53 return entry["data"] 

54 

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

56 """Set cached pricing data.""" 

57 self._cache[key] = {"data": data, "timestamp": time.time()} 

58 

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

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

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

62 

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

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

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

66 

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

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

69 return self.get("all_models") 

70 

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

72 """Cache pricing for all models.""" 

73 self.set("all_models", pricing) 

74 

75 def clear(self): 

76 """Clear all cached data.""" 

77 self._cache = {} 

78 logger.info("Pricing cache cleared") 

79 

80 def clear_expired(self): 

81 """Remove expired cache entries.""" 

82 expired_keys = [] 

83 for key, entry in self._cache.items(): 

84 if self._is_expired(entry["timestamp"]): 

85 expired_keys.append(key) 

86 

87 for key in expired_keys: 

88 del self._cache[key] 

89 

90 if expired_keys: 

91 logger.info(f"Removed {len(expired_keys)} expired cache entries") 

92 

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

94 """Get cache statistics.""" 

95 total_entries = len(self._cache) 

96 expired_count = 0 

97 

98 for entry in self._cache.values(): 

99 if self._is_expired(entry["timestamp"]): 

100 expired_count += 1 

101 

102 return { 

103 "total_entries": total_entries, 

104 "expired_entries": expired_count, 

105 "valid_entries": total_entries - expired_count, 

106 "cache_type": "in-memory", 

107 "cache_ttl": self.cache_ttl, 

108 }