Coverage for agentos/tools/serial_cache.py: 29%

1""" 

2Serialization & Caching for AgentOS. 

3 

4Serializer — adaptive serializer with JSON/msgpack/pickle auto-detection. 

5TTLCache — thread-safe time-to-live cache with LRU/LFU eviction. 

6SmartCache — compute-on-miss cache wrapping TTL cache with serializer. 

7""" 

8 

9import json 

10import pickle 

11import threading 

12import time 

13from collections import OrderedDict 

14from collections.abc import Callable 

15from dataclasses import dataclass, field 

16from enum import Enum 

17from typing import Any, Generic, TypeVar 

18 

19T = TypeVar("T") 

20 

21 

22# ============================================================================ 

23# Serializer 

24# ============================================================================ 

25 

26 

27class SerialFormat(Enum): 

28 JSON = "json" 

29 PICKLE = "pickle" 

30 MSGPACK = "msgpack" 

31 AUTO = "auto" 

32 

33 def detect(self, data: bytes) -> "SerialFormat": 

34 if self != SerialFormat.AUTO: 

35 return self 

36 if data[:2] in (b"\x80\x03", b"\x80\x04", b"\x80\x05"): 

37 return SerialFormat.PICKLE 

38 if data[:1] == b"{" or data[:1] == b"[": 

39 return SerialFormat.JSON 

40 # Try msgpack header (0x80-0x8f for fixmap, 0x90-0x9f for fixarray, 0xdc/0xdd/0xde/0xdf, etc.) 

41 if len(data) > 0 and data[0] in range(0x80, 0x100): 

42 try: 

43 import msgpack 

44 

45 msgpack.unpackb(data) 

46 return SerialFormat.MSGPACK 

47 except Exception: 

48 pass 

49 raise ValueError("Cannot auto-detect serialization format") 

50 

51 

52class Serializer: 

53 """Adaptive serializer with format auto-detection and compression support.""" 

54 

55 def __init__(self, fmt: SerialFormat = SerialFormat.JSON): 

56 self._fmt = fmt 

57 self._total_serialized: int = 0 

58 self._total_deserialized: int = 0 

59 

60 def dumps(self, obj: Any, use_msgpack: bool = False) -> bytes: 

61 fmt = SerialFormat.MSGPACK if use_msgpack else self._fmt 

62 if fmt == SerialFormat.AUTO: 

63 fmt = SerialFormat.JSON 

64 

65 if fmt == SerialFormat.JSON: 

66 data = json.dumps(obj, ensure_ascii=False, default=str) 

67 self._total_serialized += 1 

68 return data.encode("utf-8") 

69 

70 elif fmt == SerialFormat.PICKLE: 

71 data = pickle.dumps(obj) 

72 self._total_serialized += 1 

73 return data 

74 

75 elif fmt == SerialFormat.MSGPACK: 

76 import msgpack 

77 

78 data = msgpack.packb(obj, default=str) 

79 self._total_serialized += 1 

80 return data 

81 

82 raise ValueError(f"Unsupported format: {fmt}") 

83 

84 def loads(self, data: bytes, fmt: SerialFormat | None = None) -> Any: 

85 if fmt is None: 

86 fmt = SerialFormat.AUTO 

87 

88 fmt = fmt.detect(data) 

89 

90 if fmt == SerialFormat.JSON: 

91 result = json.loads(data.decode("utf-8")) 

92 self._total_deserialized += 1 

93 return result 

94 

95 elif fmt == SerialFormat.PICKLE: 

96 result = pickle.loads(data) 

97 self._total_deserialized += 1 

98 return result 

99 

100 elif fmt == SerialFormat.MSGPACK: 

101 import msgpack 

102 

103 result = msgpack.unpackb(data) 

104 self._total_deserialized += 1 

105 return result 

106 

107 raise ValueError(f"Unsupported format: {fmt}") 

108 

109 @property 

110 def stats(self) -> dict[str, Any]: 

111 return { 

112 "format": self._fmt.value if isinstance(self._fmt, SerialFormat) else self._fmt, 

113 "total_serialized": self._total_serialized, 

114 "total_deserialized": self._total_deserialized, 

115 } 

116 

117 

118# ============================================================================ 

119# EvictionPolicy 

120# ============================================================================ 

121 

122 

123class EvictionPolicy(Enum): 

124 LRU = "lru" 

125 LFU = "lfu" 

126 TTL_ONLY = "ttl_only" 

127 

128 

129@dataclass 

130class _CacheEntry(Generic[T]): 

131 value: T 

132 expires_at: float 

133 access_count: int = 0 

134 last_access: float = field(default_factory=time.monotonic) 

135 

136 

137# ============================================================================ 

138# TTLCache 

139# ============================================================================ 

140 

141 

142class TTLCache(Generic[T]): 

143 """Thread-safe TTL cache with configurable eviction policy (LRU/LFU). 

144 

145 Entries expire after ttl_seconds. On maxsize overflow, evicts based on policy. 

146 """ 

147 

148 def __init__( 

149 self, 

150 max_size: int = 1000, 

151 ttl: float = 300.0, 

152 policy: EvictionPolicy = EvictionPolicy.LRU, 

153 ): 

154 self._max_size = max_size 

155 self._ttl = ttl 

156 self._policy = policy 

157 self._data: OrderedDict[str, _CacheEntry[T]] = OrderedDict() 

158 self._lock = threading.RLock() 

159 self._hits: int = 0 

160 self._misses: int = 0 

161 self._evictions: int = 0 

162 

163 def get(self, key: str) -> T | None: 

164 with self._lock: 

165 entry = self._data.get(key) 

166 if entry is None: 

167 self._misses += 1 

168 return None 

169 

170 if time.monotonic() > entry.expires_at: 

171 del self._data[key] 

172 self._misses += 1 

173 self._evictions += 1 

174 return None 

175 

176 entry.access_count += 1 

177 entry.last_access = time.monotonic() 

178 # Move to end for LRU ordering 

179 self._data.move_to_end(key) 

180 self._hits += 1 

181 return entry.value 

182 

183 def set(self, key: str, value: T, ttl: float | None = None) -> None: 

184 with self._lock: 

185 if key in self._data: 

186 self._data.pop(key) 

187 

188 if len(self._data) >= self._max_size: 

189 self._evict_one() 

190 

191 self._data[key] = _CacheEntry( 

192 value=value, 

193 expires_at=time.monotonic() + (ttl if ttl is not None else self._ttl), 

194 ) 

195 self._data.move_to_end(key) 

196 

197 def _evict_one(self) -> None: 

198 if not self._data: 

199 return 

200 

201 if self._policy == EvictionPolicy.TTL_ONLY: 

202 # Remove oldest (first inserted) 

203 self._data.popitem(last=False) 

204 self._evictions += 1 

205 return 

206 

207 if self._policy == EvictionPolicy.LRU: 

208 # First item is least recently used (get moves items to end) 

209 self._data.popitem(last=False) 

210 self._evictions += 1 

211 return 

212 

213 if self._policy == EvictionPolicy.LFU: 

214 # Find item with lowest access count 

215 victim_key = min(self._data, key=lambda k: self._data[k].access_count) 

216 del self._data[victim_key] 

217 self._evictions += 1 

218 

219 def delete(self, key: str) -> bool: 

220 with self._lock: 

221 if key in self._data: 

222 del self._data[key] 

223 return True 

224 return False 

225 

226 def clear(self) -> None: 

227 with self._lock: 

228 self._data.clear() 

229 

230 def cleanup(self) -> int: 

231 """Remove all expired entries. Returns count removed.""" 

232 now = time.monotonic() 

233 count = 0 

234 with self._lock: 

235 expired = [k for k, v in self._data.items() if now > v.expires_at] 

236 for k in expired: 

237 del self._data[k] 

238 count += 1 

239 self._evictions += count 

240 return count 

241 

242 @property 

243 def size(self) -> int: 

244 with self._lock: 

245 return len(self._data) 

246 

247 @property 

248 def stats(self) -> dict[str, Any]: 

249 with self._lock: 

250 return { 

251 "size": len(self._data), 

252 "max_size": self._max_size, 

253 "ttl": self._ttl, 

254 "policy": self._policy.value, 

255 "hits": self._hits, 

256 "misses": self._misses, 

257 "evictions": self._evictions, 

258 "hit_rate": round(self._hits / max(1, self._hits + self._misses), 3), 

259 } 

260 

261 

262# ============================================================================ 

263# SmartCache 

264# ============================================================================ 

265 

266 

267class SmartCache(Generic[T]): 

268 """Compute-on-miss cache combining TTLCache with Serializer. 

269 

270 Provides get_or_compute() — key misses trigger the factory function, 

271 result stored in cache automatically. Supports serialization for persistence. 

272 """ 

273 

274 def __init__( 

275 self, 

276 max_size: int = 1000, 

277 ttl: float = 300.0, 

278 policy: EvictionPolicy = EvictionPolicy.LRU, 

279 ): 

280 self._cache = TTLCache[T](max_size=max_size, ttl=ttl, policy=policy) 

281 self._serializer = Serializer() 

282 

283 def get(self, key: str) -> T | None: 

284 return self._cache.get(key) 

285 

286 def get_or_compute(self, key: str, factory: Callable[[], T], ttl: float | None = None) -> T: 

287 """Get from cache or compute via factory and cache the result.""" 

288 value = self._cache.get(key) 

289 if value is not None: 

290 return value 

291 value = factory() 

292 self._cache.set(key, value, ttl=ttl) 

293 return value 

294 

295 def set(self, key: str, value: T, ttl: float | None = None) -> None: 

296 self._cache.set(key, value, ttl=ttl) 

297 

298 def delete(self, key: str) -> bool: 

299 return self._cache.delete(key) 

300 

301 def clear(self) -> None: 

302 self._cache.clear() 

303 

304 def dump(self) -> bytes: 

305 """Serialize entire cache state.""" 

306 with self._cache._lock: 

307 entries = { 

308 k: { 

309 "value": v.value, 

310 "expires_at": v.expires_at, 

311 "access_count": v.access_count, 

312 "last_access": v.last_access, 

313 } 

314 for k, v in self._cache._data.items() 

315 } 

316 return self._serializer.dumps(entries) 

317 

318 def load(self, data: bytes) -> int: 

319 """Restore cache from serialized data. Returns number of entries loaded.""" 

320 now = time.monotonic() 

321 entries = self._serializer.loads(data) 

322 count = 0 

323 for k, v in entries.items(): 

324 if v["expires_at"] > now: 

325 self._cache._data[k] = _CacheEntry( 

326 value=v["value"], 

327 expires_at=v["expires_at"], 

328 access_count=v["access_count"], 

329 last_access=v["last_access"], 

330 ) 

331 count += 1 

332 return count 

333 

334 @property 

335 def size(self) -> int: 

336 return self._cache.size 

337 

338 @property 

339 def stats(self) -> dict[str, Any]: 

340 return self._cache.stats