Coverage for agentos/llm/smart_cache.py: 38%

1"""Smart Cache — LLM response caching with exact and fuzzy matching. 

2 

3Reduces API costs by caching LLM responses. Supports: 

4 - Exact match: identical prompt → cached response 

5 - Fuzzy match: semantically similar prompts → cached response (via embeddings) 

6 - TTL-based expiration and LRU eviction 

7 - Cost savings tracking 

8""" 

9 

10from __future__ import annotations 

11 

12import hashlib 

13import json 

14import time 

15from collections import OrderedDict 

16from dataclasses import dataclass, field 

17from typing import Any, Optional 

18 

19 

20__all__ = [ 

21 "CacheConfig", 

22 "CacheStats", 

23 "SmartCache", 

24 "CacheEntry", 

25] 

26 

27 

28# ── Config & Stats ───────────────────────────────────────────────── 

29 

30@dataclass 

31class CacheConfig: 

32 """Configuration for SmartCache. 

33 

34 Attributes: 

35 max_entries: Maximum number of cached entries (LRU eviction). 

36 ttl_seconds: Time-to-live in seconds (0 = no expiry). 

37 enable_fuzzy: Enable semantic similarity matching. 

38 fuzzy_threshold: Similarity threshold for fuzzy matching (0-1). 

39 """ 

40 max_entries: int = 1000 

41 ttl_seconds: int = 3600 # 1 hour default 

42 enable_fuzzy: bool = False 

43 fuzzy_threshold: float = 0.85 

44 

45 

46@dataclass 

47class CacheStats: 

48 """Cache performance and cost savings statistics. 

49 

50 Attributes: 

51 hits: Number of cache hits (exact). 

52 fuzzy_hits: Number of fuzzy match hits. 

53 misses: Number of cache misses. 

54 total_cost_saved_usd: Estimated total API cost saved via caching. 

55 entries: Current number of cached entries. 

56 evictions: Total evicted entries (LRU + TTL). 

57 """ 

58 hits: int = 0 

59 fuzzy_hits: int = 0 

60 misses: int = 0 

61 total_cost_saved_usd: float = 0.0 

62 entries: int = 0 

63 evictions: int = 0 

64 

65 @property 

66 def hit_rate(self) -> float: 

67 total = self.hits + self.fuzzy_hits + self.misses 

68 if total == 0: 

69 return 0.0 

70 return (self.hits + self.fuzzy_hits) / total 

71 

72 def summary(self) -> dict: 

73 return { 

74 "hits": self.hits, 

75 "fuzzy_hits": self.fuzzy_hits, 

76 "misses": self.misses, 

77 "hit_rate": round(self.hit_rate, 4), 

78 "total_cost_saved_usd": round(self.total_cost_saved_usd, 6), 

79 "entries": self.entries, 

80 "evictions": self.evictions, 

81 } 

82 

83 

84# ── Cache Entry ─────────────────────────────────────────────────── 

85 

86@dataclass 

87class CacheEntry: 

88 """A single cached LLM response. 

89 

90 Attributes: 

91 key: Cache key (usually hash of the prompt/model). 

92 prompt: Original prompt text. 

93 response: Cached LLM response. 

94 model: Model name used. 

95 cost_usd: Estimated cost of the original API call (savings). 

96 tokens: Token count of the response. 

97 created_at: Unix timestamp when cached. 

98 ttl: Time-to-live in seconds. 

99 hit_count: Number of times this entry was used. 

100 """ 

101 key: str 

102 prompt: str 

103 response: Any 

104 model: str = "" 

105 cost_usd: float = 0.0 

106 tokens: int = 0 

107 created_at: float = field(default_factory=time.time) 

108 ttl: int = 3600 

109 hit_count: int = 0 

110 

111 @property 

112 def expired(self) -> bool: 

113 """Check if this entry has exceeded its TTL.""" 

114 if self.ttl <= 0: 

115 return False 

116 return time.time() - self.created_at > self.ttl 

117 

118 @property 

119 def age_seconds(self) -> float: 

120 return time.time() - self.created_at 

121 

122 

123# ── Smart Cache ────────────────────────────────────────────────── 

124 

125class SmartCache: 

126 """LLM response cache with exact + fuzzy matching and cost tracking. 

127 

128 Usage: 

129 cache = SmartCache(config=CacheConfig(max_entries=500, ttl_seconds=7200)) 

130 cache.set(prompt="What is quantum computing?", response="...", model="gpt-4o", cost_usd=0.005) 

131 cached = cache.get(prompt="What is quantum computing?", model="gpt-4o") 

132 if cached: 

133 print(f"Cache hit! Saved ${cached.cost_usd}") 

134 

135 The cache uses an LRU eviction policy with TTL-based expiration. 

136 Keys are derived from a hash of (prompt + model) for exact matching. 

137 """ 

138 

139 def __init__(self, config: Optional[CacheConfig] = None): 

140 self._config = config or CacheConfig() 

141 self._cache: OrderedDict[str, CacheEntry] = OrderedDict() 

142 self.stats = CacheStats() 

143 

144 # ── Core API ────────────────────────────────────────────────── 

145 

146 def get( 

147 self, 

148 prompt: str, 

149 model: str = "", 

150 use_fuzzy: bool = False, 

151 ) -> Optional[CacheEntry]: 

152 """Look up a cached response for the given prompt. 

153 

154 Args: 

155 prompt: The prompt text to look up. 

156 model: Model name for key disambiguation. 

157 use_fuzzy: Enable fuzzy/semantic matching (requires embeddings). 

158 

159 Returns: 

160 CacheEntry if found, None otherwise. 

161 """ 

162 # 1. Exact match 

163 key = self._make_key(prompt, model) 

164 if key in self._cache: 

165 entry = self._cache[key] 

166 # Check TTL 

167 if entry.expired: 

168 del self._cache[key] 

169 self.stats.evictions += 1 

170 self.stats.entries = len(self._cache) 

171 else: 

172 # LRU: move to end 

173 self._cache.move_to_end(key) 

174 entry.hit_count += 1 

175 self.stats.hits += 1 

176 return entry 

177 

178 # 2. Fuzzy match (optional) 

179 if use_fuzzy and self._config.enable_fuzzy: 

180 entry = self._fuzzy_lookup(prompt, model) 

181 if entry: 

182 self.stats.fuzzy_hits += 1 

183 return entry 

184 

185 # 3. Miss 

186 self.stats.misses += 1 

187 return None 

188 

189 def set( 

190 self, 

191 prompt: str, 

192 response: Any, 

193 model: str = "", 

194 cost_usd: float = 0.0, 

195 tokens: int = 0, 

196 ) -> str: 

197 """Cache a response. 

198 

199 Args: 

200 prompt: The original prompt. 

201 response: The LLM response to cache. 

202 model: Model name used. 

203 cost_usd: Cost of the original call (for savings tracking). 

204 tokens: Token count of the response. 

205 

206 Returns: 

207 The cache key. 

208 """ 

209 key = self._make_key(prompt, model) 

210 

211 # Update existing entry if present 

212 if key in self._cache: 

213 entry = self._cache[key] 

214 entry.response = response 

215 entry.cost_usd = cost_usd 

216 entry.tokens = tokens 

217 entry.created_at = time.time() 

218 self._cache.move_to_end(key) 

219 return key 

220 

221 # Evict if over capacity 

222 while len(self._cache) >= self._config.max_entries: 

223 self._evict_one() 

224 

225 entry = CacheEntry( 

226 key=key, 

227 prompt=prompt, 

228 response=response, 

229 model=model, 

230 cost_usd=cost_usd, 

231 tokens=tokens, 

232 ttl=self._config.ttl_seconds, 

233 ) 

234 self._cache[key] = entry 

235 self.stats.entries = len(self._cache) 

236 return key 

237 

238 def clear(self): 

239 """Clear all cached entries.""" 

240 self._cache.clear() 

241 self.stats.entries = 0 

242 

243 def gc(self) -> int: 

244 """Garbage collect expired entries. Returns number of entries evicted.""" 

245 count = 0 

246 expired_keys = [k for k, v in self._cache.items() if v.expired] 

247 for k in expired_keys: 

248 del self._cache[k] 

249 count += 1 

250 self.stats.evictions += count 

251 self.stats.entries = len(self._cache) 

252 return count 

253 

254 # ── Info ───────────────────────────────────────────────────── 

255 

256 @property 

257 def size(self) -> int: 

258 return len(self._cache) 

259 

260 @property 

261 def config(self) -> CacheConfig: 

262 return self._config 

263 

264 def contains(self, prompt: str, model: str = "") -> bool: 

265 """Check if a prompt is cached (exact match, ignores TTL).""" 

266 return self._make_key(prompt, model) in self._cache 

267 

268 # ── Internal ───────────────────────────────────────────────── 

269 

270 def _make_key(self, prompt: str, model: str = "") -> str: 

271 """Generate a deterministic cache key from prompt + model.""" 

272 content = f"{model or 'default'}:{prompt}" 

273 return hashlib.sha256(content.encode()).hexdigest()[:32] 

274 

275 def _fuzzy_lookup(self, prompt: str, model: str = "") -> Optional[CacheEntry]: 

276 """Semantic similarity lookup using simple keyword overlap. 

277 

278 For production, replace with embedding-based similarity (cosine). 

279 """ 

280 prompt_words = set(prompt.lower().split()) 

281 if not prompt_words: 

282 return None 

283 

284 best_score = 0.0 

285 best_entry: Optional[CacheEntry] = None 

286 best_key = "" 

287 

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

289 if entry.expired: 

290 continue 

291 if model and entry.model and entry.model != model: 

292 continue 

293 

294 entry_words = set(entry.prompt.lower().split()) 

295 if not entry_words: 

296 continue 

297 

298 # Jaccard similarity 

299 intersection = prompt_words & entry_words 

300 union = prompt_words | entry_words 

301 score = len(intersection) / len(union) if union else 0.0 

302 

303 if score > best_score and score >= self._config.fuzzy_threshold: 

304 best_score = score 

305 best_entry = entry 

306 best_key = key 

307 

308 if best_entry: 

309 self._cache.move_to_end(best_key) 

310 best_entry.hit_count += 1 

311 return best_entry 

312 

313 return None 

314 

315 def _evict_one(self): 

316 """Evict the oldest entry (LRU front).""" 

317 if self._cache: 

318 self._cache.popitem(last=False) 

319 self.stats.evictions += 1