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, Dict, List, Optional, Tuple 

27 

28 

29# --------------------------------------------------------------------------- 

30# Cache Entry 

31# --------------------------------------------------------------------------- 

32 

33 

34@dataclass 

35class CacheEntry: 

36 """A single cache entry.""" 

37 key: str 

38 value: Any 

39 size_bytes: int 

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

41 ttl_seconds: Optional[float] = 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) -> Optional[CacheEntry]: 

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: Optional[float] = 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] = {"default": EvictionPolicy(max_entries, max_size_mb)} 

174 

175 self._max_entries = max_entries 

176 self._max_size_mb = max_size_mb 

177 

178 self._stats: Dict[str, Any] = { 

179 "hits": 0, 

180 "misses": 0, 

181 "total_requests": 0, 

182 "bytes_saved_est": 0.0, 

183 } 

184 self._lock = threading.Lock() 

185 

186 # -- Key generation -- 

187 

188 @staticmethod 

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

190 """Generate a deterministic cache key.""" 

191 normalized = json.dumps({"model": model, "prompt": prompt.strip(), **kwargs}, sort_keys=True) 

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

193 

194 @staticmethod 

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

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

197 if isinstance(value, str): 

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

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

200 

201 # -- Core operations -- 

202 

203 def get(self, model: str, prompt: str, **kwargs) -> Optional[Any]: 

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

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

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

207 

208 with self._lock: 

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

210 

211 policy = self._policies.get(policy_key) 

212 if policy is None: 

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

214 return None 

215 

216 entry = policy.get(key) 

217 if entry is None: 

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

219 return None 

220 

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

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

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

224 return entry.value 

225 

226 def put( 

227 self, 

228 model: str, 

229 prompt: str, 

230 response: Any, 

231 ttl_seconds: Optional[float] = None, 

232 **kwargs, 

233 ) -> None: 

234 """Cache a response.""" 

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

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

237 

238 with self._lock: 

239 if policy_key not in self._policies: 

240 self._policies[policy_key] = EvictionPolicy( 

241 self._max_entries, self._max_size_mb 

242 ) 

243 

244 entry = CacheEntry( 

245 key=key, 

246 value=response, 

247 size_bytes=self._estimate_bytes(response), 

248 ttl_seconds=ttl_seconds or self._ttl, 

249 ) 

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

251 

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

253 """Invalidate a specific cache entry.""" 

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

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

256 

257 with self._lock: 

258 policy = self._policies.get(policy_key) 

259 if policy is None: 

260 return False 

261 policy.remove(key) 

262 return True 

263 

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

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

266 with self._lock: 

267 if model in self._policies: 

268 self._policies[model].clear() 

269 

270 def clear(self) -> None: 

271 """Clear entire cache.""" 

272 with self._lock: 

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

274 policy.clear() 

275 

276 # -- Statistics -- 

277 

278 def get_stats(self) -> Dict[str, Any]: 

279 """Return cache statistics.""" 

280 with self._lock: 

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

282 stats = dict(self._stats) 

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

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

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

286 stats["total_size_mb"] = round( 

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

288 ) 

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

290 return stats 

291 

292 def get_per_model_stats(self) -> Dict[str, Dict[str, Any]]: 

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

294 with self._lock: 

295 return { 

296 model: { 

297 "entries": policy.size, 

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

299 } 

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

301 } 

302 

303 def warmup(self, entries: List[Tuple[str, str, Any, Optional[float]]]) -> int: 

304 """ 

305 Pre-warm the cache with known frequent prompts. 

306 

307 Returns number of entries loaded. 

308 """ 

309 count = 0 

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

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

312 count += 1 

313 return count