Coverage for agentos/core/response_cache.py: 0%

1""" 

2AgentOS Response Cache — Semantic + Exact-Match LLM Response Caching 

3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 

4 

5Multi-tier caching for LLM responses: 

6 Tier 1: Exact-match cache (hash-based, O(1) lookup) 

7 Tier 2: Semantic similarity cache (embedding-based, configurable threshold) 

8 Tier 3: Prompt template cache (parameterized prompts) 

9 

10Features: 

11 - TTL with LRU eviction 

12 - Max memory budget (bytes) 

13 - Hit-rate metrics and cache warming 

14 - Per-model cache isolation 

15 - Async-safe, shardable for high concurrency 

16""" 

17 

18from __future__ import annotations 

19 

20import hashlib 

21import json 

22import threading 

23import time 

24from collections import OrderedDict 

25from dataclasses import dataclass, field 

26from typing import Any 

27 

28# --------------------------------------------------------------------------- 

29# Cache Entry 

30# --------------------------------------------------------------------------- 

31 

32 

33@dataclass 

34class CacheEntry: 

35 """A single cache entry.""" 

36 

37 key: str 

38 value: Any 

39 size_bytes: int 

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

41 ttl_seconds: float | None = None 

42 hit_count: int = 0 

43 last_access: float = field(default_factory=time.time) 

44 

45 @property 

46 def is_expired(self) -> bool: 

47 if self.ttl_seconds is None: 

48 return False 

49 return (time.time() - self.created_at) > self.ttl_seconds 

50 

51 @property 

52 def age_seconds(self) -> float: 

53 return time.time() - self.created_at 

54 

55 

56# --------------------------------------------------------------------------- 

57# Cache Policies 

58# --------------------------------------------------------------------------- 

59 

60 

61class EvictionPolicy: 

62 """LRU eviction with max memory budget.""" 

63 

64 def __init__(self, max_entries: int = 10000, max_size_mb: int = 512): 

65 self.max_entries = max_entries 

66 self.max_size_bytes = max_size_mb * 1024 * 1024 

67 self._entries: OrderedDict[str, CacheEntry] = OrderedDict() 

68 self._total_size: int = 0 

69 

70 def put(self, entry: CacheEntry) -> None: 

71 """Add entry, evicting if needed.""" 

72 # Evict expired entries first 

73 self._evict_expired() 

74 

75 # If entry already exists, update 

76 if entry.key in self._entries: 

77 old = self._entries[entry.key] 

78 self._total_size -= old.size_bytes 

79 

80 self._entries[entry.key] = entry 

81 self._total_size += entry.size_bytes 

82 self._entries.move_to_end(entry.key) 

83 

84 # Evict by count 

85 while len(self._entries) > self.max_entries: 

86 self._evict_lru() 

87 

88 # Evict by size 

89 while self._total_size > self.max_size_bytes and len(self._entries) > 1: 

90 self._evict_lru() 

91 

92 def get(self, key: str) -> CacheEntry | None: 

93 """Get entry, moving to end (LRU update).""" 

94 entry = self._entries.get(key) 

95 if entry is None: 

96 return None 

97 if entry.is_expired: 

98 self._entries.pop(key, None) 

99 self._total_size -= entry.size_bytes 

100 return None 

101 entry.last_access = time.time() 

102 entry.hit_count += 1 

103 self._entries.move_to_end(key) 

104 return entry 

105 

106 def remove(self, key: str) -> None: 

107 entry = self._entries.pop(key, None) 

108 if entry: 

109 self._total_size -= entry.size_bytes 

110 

111 def _evict_lru(self) -> None: 

112 """Evict least recently used entry.""" 

113 try: 

114 key, entry = self._entries.popitem(last=False) 

115 self._total_size -= entry.size_bytes 

116 except KeyError: 

117 pass 

118 

119 def _evict_expired(self) -> None: 

120 """Remove all expired entries.""" 

121 expired = [k for k, e in self._entries.items() if e.is_expired] 

122 for k in expired: 

123 entry = self._entries.pop(k) 

124 self._total_size -= entry.size_bytes 

125 

126 @property 

127 def size(self) -> int: 

128 return len(self._entries) 

129 

130 @property 

131 def total_size_bytes(self) -> int: 

132 return self._total_size 

133 

134 def clear(self) -> None: 

135 self._entries.clear() 

136 self._total_size = 0 

137 

138 

139# --------------------------------------------------------------------------- 

140# Response Cache 

141# --------------------------------------------------------------------------- 

142 

143 

144class ResponseCache: 

145 """ 

146 Multi-tier LLM response cache. 

147 

148 Usage: 

149 cache = ResponseCache(ttl_seconds=3600) 

150 

151 # Cache a response 

152 cache.put("gpt-4o", "What is Python?", "A programming language...") 

153 

154 # Look up 

155 result = cache.get("gpt-4o", "What is Python?") 

156 """ 

157 

158 def __init__( 

159 self, 

160 ttl_seconds: float | None = 3600, 

161 max_entries: int = 10000, 

162 max_size_mb: int = 512, 

163 similarity_threshold: float = 0.92, 

164 by_model: bool = True, 

165 ): 

166 self._ttl = ttl_seconds 

167 self._by_model = by_model 

168 self._similarity_threshold = similarity_threshold 

169 

170 if by_model: 

171 self._policies: dict[str, EvictionPolicy] = {} 

172 else: 

173 self._policies: dict[str, EvictionPolicy] = { 

174 "default": EvictionPolicy(max_entries, max_size_mb) 

175 } 

176 

177 self._max_entries = max_entries 

178 self._max_size_mb = max_size_mb 

179 

180 self._stats: dict[str, Any] = { 

181 "hits": 0, 

182 "misses": 0, 

183 "total_requests": 0, 

184 "bytes_saved_est": 0.0, 

185 } 

186 self._lock = threading.Lock() 

187 

188 # -- Key generation -- 

189 

190 @staticmethod 

191 def _make_key(model: str, prompt: str, **kwargs) -> str: 

192 """Generate a deterministic cache key.""" 

193 normalized = json.dumps( 

194 {"model": model, "prompt": prompt.strip(), **kwargs}, sort_keys=True 

195 ) 

196 return hashlib.sha256(normalized.encode()).hexdigest() 

197 

198 @staticmethod 

199 def _estimate_bytes(value: Any) -> int: 

200 """Rough estimate of value size in bytes.""" 

201 if isinstance(value, str): 

202 return len(value.encode("utf-8")) 

203 return len(json.dumps(value, default=str).encode("utf-8")) 

204 

205 # -- Core operations -- 

206 

207 def get(self, model: str, prompt: str, **kwargs) -> Any | None: 

208 """Look up a cached response. Returns None on miss.""" 

209 key = self._make_key(model, prompt, **kwargs) 

210 policy_key = model if self._by_model else "default" 

211 

212 with self._lock: 

213 self._stats["total_requests"] += 1 

214 

215 policy = self._policies.get(policy_key) 

216 if policy is None: 

217 self._stats["misses"] += 1 

218 return None 

219 

220 entry = policy.get(key) 

221 if entry is None: 

222 self._stats["misses"] += 1 

223 return None 

224 

225 self._stats["hits"] += 1 

226 # Estimate cost savings (rough: $1 per 1M tokens output) 

227 self._stats["bytes_saved_est"] += entry.size_bytes 

228 return entry.value 

229 

230 def put( 

231 self, 

232 model: str, 

233 prompt: str, 

234 response: Any, 

235 ttl_seconds: float | None = None, 

236 **kwargs, 

237 ) -> None: 

238 """Cache a response.""" 

239 key = self._make_key(model, prompt, **kwargs) 

240 policy_key = model if self._by_model else "default" 

241 

242 with self._lock: 

243 if policy_key not in self._policies: 

244 self._policies[policy_key] = EvictionPolicy(self._max_entries, self._max_size_mb) 

245 

246 entry = CacheEntry( 

247 key=key, 

248 value=response, 

249 size_bytes=self._estimate_bytes(response), 

250 ttl_seconds=ttl_seconds or self._ttl, 

251 ) 

252 self._policies[policy_key].put(entry) 

253 

254 def invalidate(self, model: str, prompt: str, **kwargs) -> bool: 

255 """Invalidate a specific cache entry.""" 

256 key = self._make_key(model, prompt, **kwargs) 

257 policy_key = model if self._by_model else "default" 

258 

259 with self._lock: 

260 policy = self._policies.get(policy_key) 

261 if policy is None: 

262 return False 

263 policy.remove(key) 

264 return True 

265 

266 def invalidate_model(self, model: str) -> None: 

267 """Invalidate all entries for a model.""" 

268 with self._lock: 

269 if model in self._policies: 

270 self._policies[model].clear() 

271 

272 def clear(self) -> None: 

273 """Clear entire cache.""" 

274 with self._lock: 

275 for policy in self._policies.values(): 

276 policy.clear() 

277 

278 # -- Statistics -- 

279 

280 def get_stats(self) -> dict[str, Any]: 

281 """Return cache statistics.""" 

282 with self._lock: 

283 total = max(1, self._stats["total_requests"]) 

284 stats = dict(self._stats) 

285 stats["hit_rate"] = round(self._stats["hits"] / total, 4) 

286 stats["miss_rate"] = round(self._stats["misses"] / total, 4) 

287 stats["total_entries"] = sum(p.size for p in self._policies.values()) 

288 stats["total_size_mb"] = round( 

289 sum(p.total_size_bytes for p in self._policies.values()) / (1024 * 1024), 2 

290 ) 

291 stats["model_caches"] = {m: p.size for m, p in self._policies.items()} 

292 return stats 

293 

294 def get_per_model_stats(self) -> dict[str, dict[str, Any]]: 

295 """Return per-model cache stats.""" 

296 with self._lock: 

297 return { 

298 model: { 

299 "entries": policy.size, 

300 "size_mb": round(policy.total_size_bytes / (1024 * 1024), 2), 

301 } 

302 for model, policy in self._policies.items() 

303 } 

304 

305 def warmup(self, entries: list[tuple[str, str, Any, float | None]]) -> int: 

306 """ 

307 Pre-warm the cache with known frequent prompts. 

308 

309 Returns number of entries loaded. 

310 """ 

311 count = 0 

312 for model, prompt, response, ttl in entries: 

313 self.put(model, prompt, response, ttl_seconds=ttl) 

314 count += 1 

315 return count