Coverage for agentos/memory/pager.py: 28%

1""" 

2Virtual Memory Pager — Letta-style context memory swapping. 

3 

4Inspired by OS virtual memory: when the agent's working context fills up, 

5pages out old/low-importance episodic memories to a persistent swap store, 

6and intelligently pages them back in when relevant to the current task. 

7 

8Architecture: 

9 L1 (in-memory) ──page_out──→ SwapStore (disk/DB) 

10 SwapStore ──page_in───→ L1 (promoted back) 

11 

12Key features: 

13- Importance-weighted eviction: least important + oldest first 

14- Smart recall: semantic search in swap for context-relevant pages 

15- Page compaction: batch-compress multiple items into summary pages 

16- Statistics: track page hits/misses for tuning 

17""" 

18 

19from __future__ import annotations 

20 

21import json 

22import os 

23import time 

24import uuid 

25from collections.abc import Callable 

26from dataclasses import dataclass, field 

27from typing import Any 

28 

29from agentos.memory.pyramid import MemoryItem, MemoryLayer, MemoryPyramid, MemoryType 

30 

31# ── Data Structures ────────────────────────────────────────────── 

32 

33 

34@dataclass 

35class MemoryPage: 

36 """A compressed page of evicted memories, like a virtual memory page.""" 

37 

38 id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) 

39 items: list[dict[str, Any]] = field(default_factory=list) # serialized MemoryItem dicts 

40 summary: str = "" # LLM-generated summary of page contents 

41 keywords: list[str] = field(default_factory=list) 

42 importance_avg: float = 0.0 

43 item_count: int = 0 

44 evicted_at: float = field(default_factory=time.time) 

45 evicted_from: str = "l1" # which layer it was evicted from 

46 

47 def to_dict(self) -> dict[str, Any]: 

48 return { 

49 "id": self.id, 

50 "items": self.items, 

51 "summary": self.summary, 

52 "keywords": self.keywords, 

53 "importance_avg": self.importance_avg, 

54 "item_count": self.item_count, 

55 "evicted_at": self.evicted_at, 

56 "evicted_from": self.evicted_from, 

57 } 

58 

59 @classmethod 

60 def from_dict(cls, d: dict[str, Any]) -> MemoryPage: 

61 return cls( 

62 id=d.get("id", uuid.uuid4().hex[:12]), 

63 items=d.get("items", []), 

64 summary=d.get("summary", ""), 

65 keywords=d.get("keywords", []), 

66 importance_avg=d.get("importance_avg", 0.0), 

67 item_count=d.get("item_count", 0), 

68 evicted_at=d.get("evicted_at", time.time()), 

69 evicted_from=d.get("evicted_from", "l1"), 

70 ) 

71 

72 

73@dataclass 

74class PagerStats: 

75 """Pager performance statistics.""" 

76 

77 total_page_outs: int = 0 

78 total_page_ins: int = 0 

79 total_items_evicted: int = 0 

80 total_items_recalled: int = 0 

81 page_hits: int = 0 # page-in found relevant data 

82 page_misses: int = 0 # page-in found nothing relevant 

83 last_page_out_at: float = 0.0 

84 last_page_in_at: float = 0.0 

85 

86 @property 

87 def hit_rate(self) -> float: 

88 total = self.page_hits + self.page_misses 

89 return self.page_hits / total if total > 0 else 0.0 

90 

91 

92# ── Swap Store Backend ──────────────────────────────────────────── 

93 

94 

95class SwapStore: 

96 """ 

97 Persistent storage for paged-out memories. 

98 

99 Default: file-based JSON store. Can be swapped for SQLite/Postgres. 

100 """ 

101 

102 def __init__(self, path: str = ""): 

103 self.path = path or self._default_path() 

104 self._pages: dict[str, MemoryPage] = {} 

105 self._keyword_index: dict[str, set[str]] = {} # keyword → page_ids 

106 os.makedirs(os.path.dirname(self.path), exist_ok=True) 

107 self._load() 

108 

109 @staticmethod 

110 def _default_path() -> str: 

111 return os.path.join(os.path.expanduser("~"), ".agentos", "memory_swap.json") 

112 

113 def store(self, page: MemoryPage) -> None: 

114 self._pages[page.id] = page 

115 for kw in page.keywords: 

116 self._keyword_index.setdefault(kw.lower(), set()).add(page.id) 

117 self._flush() 

118 

119 def search(self, query_keywords: list[str], limit: int = 5) -> list[MemoryPage]: 

120 """Keyword-based search for relevant pages.""" 

121 scored: dict[str, float] = {} 

122 for kw in query_keywords: 

123 kw_lower = kw.lower() 

124 for page_id in self._keyword_index.get(kw_lower, set()): 

125 scored[page_id] = scored.get(page_id, 0) + 1.0 

126 

127 # Sort by score desc, then by importance_avg desc 

128 ranked = sorted( 

129 scored.items(), 

130 key=lambda x: (x[1], self._pages.get(x[0], MemoryPage()).importance_avg), 

131 reverse=True, 

132 ) 

133 return [self._pages[pid] for pid, _ in ranked[:limit]] 

134 

135 def get(self, page_id: str) -> MemoryPage | None: 

136 return self._pages.get(page_id) 

137 

138 def remove(self, page_id: str) -> bool: 

139 page = self._pages.pop(page_id, None) 

140 if page: 

141 for kw in page.keywords: 

142 idx = self._keyword_index.get(kw.lower(), set()) 

143 idx.discard(page_id) 

144 if not idx: 

145 del self._keyword_index[kw.lower()] 

146 self._flush() 

147 return True 

148 return False 

149 

150 def list_all(self) -> list[MemoryPage]: 

151 return list(self._pages.values()) 

152 

153 def clear(self) -> None: 

154 self._pages.clear() 

155 self._keyword_index.clear() 

156 self._flush() 

157 

158 def _flush(self) -> None: 

159 try: 

160 with open(self.path, "w") as f: 

161 json.dump( 

162 { 

163 "pages": {k: v.to_dict() for k, v in self._pages.items()}, 

164 "keyword_index": {k: list(v) for k, v in self._keyword_index.items()}, 

165 }, 

166 f, 

167 indent=2, 

168 ) 

169 except Exception: 

170 pass 

171 

172 def _load(self) -> None: 

173 if not os.path.exists(self.path): 

174 return 

175 try: 

176 with open(self.path) as f: 

177 data = json.load(f) 

178 self._pages = {k: MemoryPage.from_dict(v) for k, v in data.get("pages", {}).items()} 

179 self._keyword_index = {k: set(v) for k, v in data.get("keyword_index", {}).items()} 

180 except Exception: 

181 self._pages = {} 

182 self._keyword_index = {} 

183 

184 

185# ── Memory Pager ────────────────────────────────────────────────── 

186 

187 

188class MemoryPager: 

189 """ 

190 Virtual memory pager for Agent context. 

191 

192 Pages out old/less-important memories to swap when context is full, 

193 and intelligently pages them back in when needed. 

194 

195 Usage: 

196 pager = MemoryPager(pyramid, summarizer_fn=my_summarizer) 

197 paged = await pager.page_out(ratio=0.8) # evict when 80% full 

198 recalled = await pager.page_in(["python", "error", "debug"]) 

199 """ 

200 

201 def __init__( 

202 self, 

203 pyramid: MemoryPyramid, 

204 summarizer_fn: Callable | None = None, 

205 swap_store: SwapStore | None = None, 

206 max_pages: int = 500, 

207 page_size: int = 10, # items per page 

208 eviction_ratio: float = 0.3, # evict this % of working+episodic when full 

209 ): 

210 self.pyramid = pyramid 

211 self.summarizer = summarizer_fn 

212 self.swap = swap_store or SwapStore() 

213 self.max_pages = max_pages 

214 self.page_size = page_size 

215 self.eviction_ratio = eviction_ratio 

216 self.stats = PagerStats() 

217 

218 async def page_out(self, fill_ratio: float = 0.85) -> int: 

219 """ 

220 Evict low-importance episodic+working memories to swap. 

221 

222 Called when context is approaching the token limit. 

223 

224 Args: 

225 fill_ratio: Current context fill ratio (0-1) 

226 

227 Returns: 

228 Number of items evicted 

229 """ 

230 if fill_ratio < 0.7: 

231 return 0 # Not full enough yet 

232 

233 # Determine scale: the fuller it is, the more aggressive 

234 scale = min(1.0, (fill_ratio - 0.7) / 0.3) 

235 to_evict_count = max( 

236 1, int(len(self.pyramid._memories[MemoryType.EPISODIC]) * self.eviction_ratio * scale) 

237 ) 

238 

239 # Collect candidates: episodic + old working memories 

240 candidates: list[tuple[str, MemoryItem]] = [] 

241 for key, item in self.pyramid._memories[MemoryType.EPISODIC].items(): 

242 candidates.append((key, item)) 

243 for key, item in self.pyramid._memories[MemoryType.WORKING].items(): 

244 if item.access_count < 3: # Only evict rarely-accessed working memories 

245 candidates.append((key, item)) 

246 

247 if not candidates: 

248 return 0 

249 

250 # Sort by (importance asc, age desc) — evict least important, oldest first 

251 now = time.time() 

252 candidates.sort(key=lambda x: (x[1].importance, -(now - x[1].created_at))) 

253 to_evict = candidates[:to_evict_count] 

254 

255 # Batch into pages 

256 pages_created = 0 

257 for i in range(0, len(to_evict), self.page_size): 

258 batch = to_evict[i : i + self.page_size] 

259 page = await self._create_page(batch) 

260 self.swap.store(page) 

261 pages_created += 1 

262 

263 self.stats.total_page_outs += pages_created 

264 self.stats.total_items_evicted += len(to_evict) 

265 self.stats.last_page_out_at = time.time() 

266 

267 return len(to_evict) 

268 

269 async def page_in(self, query_keywords: list[str], limit: int = 3) -> list[MemoryItem]: 

270 """ 

271 Search swap for relevant memories and promote them back to L1. 

272 

273 Args: 

274 query_keywords: Keywords to search for (e.g., current task description) 

275 limit: Max pages to recall 

276 

277 Returns: 

278 List of MemoryItem that were restored 

279 """ 

280 pages = self.swap.search(query_keywords, limit=limit) 

281 

282 if not pages: 

283 self.stats.page_misses += 1 

284 return [] 

285 

286 self.stats.page_hits += 1 

287 self.stats.last_page_in_at = time.time() 

288 

289 restored_items: list[MemoryItem] = [] 

290 for page in pages: 

291 for item_dict in page.items: 

292 try: 

293 item = MemoryItem.from_dict(item_dict) 

294 # Promote back to L1 episodic memory 

295 item.layer = MemoryLayer.L1 

296 key = item_dict.get("metadata", {}).get("key", item.id) 

297 self.pyramid._memories[MemoryType.EPISODIC][key] = item 

298 self.pyramid._index[key] = item 

299 restored_items.append(item) 

300 except Exception: 

301 continue 

302 # Remove the page from swap (it's back in memory now) 

303 self.swap.remove(page.id) 

304 

305 self.stats.total_page_ins += len(pages) 

306 self.stats.total_items_recalled += len(restored_items) 

307 

308 return restored_items 

309 

310 async def _create_page(self, items: list[tuple[str, MemoryItem]]) -> MemoryPage: 

311 """Create a compressed memory page from a batch of items.""" 

312 page = MemoryPage() 

313 keywords_set: set[str] = set() 

314 total_imp = 0.0 

315 

316 for key, item in items: 

317 item_dict = { 

318 **item.to_dict(), 

319 "metadata": {**item.metadata, "key": key}, 

320 } 

321 page.items.append(item_dict) 

322 

323 # Extract keywords from content and metadata 

324 if isinstance(item.content, str): 

325 for word in item.content.lower().split()[:20]: 

326 if len(word) > 3 and word.isalpha(): 

327 keywords_set.add(word) 

328 for v in item.metadata.values(): 

329 if isinstance(v, str) and len(v) < 50: 

330 keywords_set.add(v.lower()) 

331 

332 total_imp += item.importance 

333 

334 # Remove from pyramid 

335 self.pyramid._memories[item.type].pop(key, None) 

336 self.pyramid._index.pop(key, None) 

337 

338 page.item_count = len(items) 

339 page.importance_avg = total_imp / len(items) if items else 0.0 

340 page.keywords = list(keywords_set) 

341 

342 # Generate summary via LLM if available 

343 if self.summarizer and page.items: 

344 try: 

345 contents = "\n".join( 

346 i.get("content", "") for i in page.items if isinstance(i.get("content"), str) 

347 )[:2000] 

348 page.summary = await self.summarizer(contents) 

349 except Exception: 

350 page.summary = f"{len(items)} memories, avg importance {page.importance_avg:.2f}" 

351 

352 return page 

353 

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

355 """Get comprehensive pager statistics.""" 

356 return { 

357 "total_page_outs": self.stats.total_page_outs, 

358 "total_page_ins": self.stats.total_page_ins, 

359 "total_items_evicted": self.stats.total_items_evicted, 

360 "total_items_recalled": self.stats.total_items_recalled, 

361 "page_hits": self.stats.page_hits, 

362 "page_misses": self.stats.page_misses, 

363 "hit_rate": f"{self.stats.hit_rate:.1%}", 

364 "swap_pages_stored": len(self.swap.list_all()), 

365 "last_page_out": self.stats.last_page_out_at, 

366 "last_page_in": self.stats.last_page_in_at, 

367 } 

368 

369 

370# ── Loop Integration Helper ──────────────────────────────────────── 

371 

372 

373def create_paging_callback(pager: MemoryPager) -> Callable: 

374 """ 

375 Create a callback for the agent loop's auto-paging hook. 

376 

377 Usage: 

378 pager = MemoryPager(pyramid) 

379 loop.set_auto_paging(create_paging_callback(pager)) 

380 """ 

381 

382 async def auto_paging(usage_ratio: float) -> int: 

383 evicted = await pager.page_out(usage_ratio) 

384 return evicted 

385 

386 return auto_paging 

387 

388 

389async def recall_relevant_memories( 

390 pager: MemoryPager, 

391 task_description: str, 

392 limit: int = 3, 

393) -> list[MemoryItem]: 

394 """ 

395 Recall memories relevant to a task from swap. 

396 

397 Extracts keywords from task description and pages in relevant memories. 

398 

399 Usage: 

400 items = await recall_relevant_memories(pager, "debug the database connection error") 

401 """ 

402 keywords = [w.lower() for w in task_description.split() if len(w) > 3 and w.isalpha()] 

403 return await pager.page_in(keywords, limit=limit) 

404 

405 

406# ── Auto-generated compat stubs ── 

407 

408 

409def recall_relevant_memories(*args, **kwargs): # noqa: F811 

410 pass