Coverage for agentos/memory/pyramid.py: 37%

1""" 

2Memory Pyramid for NexusAgent. 

3 

4Multi-layer memory management system inspired by human memory: 

5- Working Memory: Current task context (short-term) 

6- Episodic Memory: Past experiences and events 

7- Semantic Memory: Facts and knowledge (long-term) 

8- Procedural Memory: Skills and procedures 

9""" 

10 

11from __future__ import annotations 

12 

13import time 

14import uuid 

15from dataclasses import dataclass, field 

16from enum import StrEnum 

17from typing import Any 

18 

19 

20class MemoryType(StrEnum): 

21 """Types of memory in the pyramid.""" 

22 

23 WORKING = "working" # Current task context 

24 EPISODIC = "episodic" # Past experiences 

25 SEMANTIC = "semantic" # Facts and knowledge 

26 PROCEDURAL = "procedural" # Skills and procedures 

27 

28 

29class MemoryLayer(StrEnum): 

30 """Memory layers (L1=fast, L2=persistent).""" 

31 

32 L1 = "l1" # Fast, in-memory 

33 L2 = "l2" # Persistent, file-based 

34 

35 

36@dataclass 

37class MemoryItem: 

38 """ 

39 Single memory item. 

40 

41 Attributes: 

42 id: Unique identifier 

43 type: Memory type 

44 layer: Memory layer (L1/L2) 

45 content: Memory content 

46 metadata: Additional metadata 

47 created_at: Creation timestamp 

48 accessed_at: Last access timestamp 

49 access_count: Number of accesses 

50 importance: Importance score (0-1) 

51 """ 

52 

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

54 type: MemoryType = MemoryType.WORKING 

55 layer: MemoryLayer = MemoryLayer.L1 

56 content: Any = None 

57 metadata: dict[str, Any] = field(default_factory=dict) 

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

59 accessed_at: float = field(default_factory=time.time) 

60 access_count: int = 0 

61 importance: float = 0.5 

62 

63 def access(self) -> None: 

64 """Mark as accessed.""" 

65 self.accessed_at = time.time() 

66 self.access_count += 1 

67 

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

69 """Convert to dict.""" 

70 return { 

71 "id": self.id, 

72 "type": self.type.value, 

73 "layer": self.layer.value, 

74 "content": self.content, 

75 "metadata": self.metadata, 

76 "created_at": self.created_at, 

77 "accessed_at": self.accessed_at, 

78 "access_count": self.access_count, 

79 "importance": self.importance, 

80 } 

81 

82 @classmethod 

83 def from_dict(cls, data: dict[str, Any]) -> MemoryItem: 

84 """Create from dict.""" 

85 return cls( 

86 id=data.get("id", uuid.uuid4().hex[:12]), 

87 type=MemoryType(data.get("type", "working")), 

88 layer=MemoryLayer(data.get("layer", "l1")), 

89 content=data.get("content"), 

90 metadata=data.get("metadata", {}), 

91 created_at=data.get("created_at", time.time()), 

92 accessed_at=data.get("accessed_at", time.time()), 

93 access_count=data.get("access_count", 0), 

94 importance=data.get("importance", 0.5), 

95 ) 

96 

97 

98class MemoryPyramid: 

99 """ 

100 Multi-layer memory management system. 

101 

102 Organizes memories into types (working/episodic/semantic/procedural) 

103 and layers (L1=fast/L2=persistent). 

104 

105 Usage: 

106 pyramid = MemoryPyramid() 

107 pyramid.store("user_preference", {"theme": "dark"}, MemoryType.SEMANTIC) 

108 prefs = pyramid.recall("user_preference") 

109 """ 

110 

111 def __init__(self, max_working: int = 100, max_episodic: int = 1000): 

112 """ 

113 Initialize memory pyramid. 

114 

115 Args: 

116 max_working: Max items in working memory 

117 max_episodic: Max items in episodic memory 

118 """ 

119 self.max_working = max_working 

120 self.max_episodic = max_episodic 

121 

122 # Memory storage by type 

123 self._memories: dict[MemoryType, dict[str, MemoryItem]] = { 

124 MemoryType.WORKING: {}, 

125 MemoryType.EPISODIC: {}, 

126 MemoryType.SEMANTIC: {}, 

127 MemoryType.PROCEDURAL: {}, 

128 } 

129 

130 # Index for fast lookup 

131 self._index: dict[str, MemoryItem] = {} 

132 

133 def store( 

134 self, 

135 key: str, 

136 content: Any, 

137 memory_type: MemoryType = MemoryType.WORKING, 

138 layer: MemoryLayer = MemoryLayer.L1, 

139 importance: float = 0.5, 

140 **metadata, 

141 ) -> MemoryItem: 

142 """ 

143 Store a memory item. 

144 

145 Args: 

146 key: Memory key 

147 content: Memory content 

148 memory_type: Type of memory 

149 layer: Memory layer 

150 importance: Importance score (0-1) 

151 **metadata: Additional metadata 

152 

153 Returns: 

154 Created MemoryItem 

155 """ 

156 # Check capacity for working memory 

157 if memory_type == MemoryType.WORKING: 

158 if len(self._memories[MemoryType.WORKING]) >= self.max_working: 

159 self._evict_working() 

160 

161 # Check capacity for episodic memory 

162 if memory_type == MemoryType.EPISODIC: 

163 if len(self._memories[MemoryType.EPISODIC]) >= self.max_episodic: 

164 self._evict_episodic() 

165 

166 # Create memory item 

167 item = MemoryItem( 

168 type=memory_type, 

169 layer=layer, 

170 content=content, 

171 metadata=metadata, 

172 importance=importance, 

173 ) 

174 

175 # Store 

176 self._memories[memory_type][key] = item 

177 self._index[key] = item 

178 

179 return item 

180 

181 def recall(self, key: str) -> MemoryItem | None: 

182 """ 

183 Recall a memory item. 

184 

185 Args: 

186 key: Memory key 

187 

188 Returns: 

189 MemoryItem if found, None otherwise 

190 """ 

191 item = self._index.get(key) 

192 if item: 

193 item.access() 

194 return item 

195 

196 def search( 

197 self, 

198 memory_type: MemoryType | None = None, 

199 limit: int = 10, 

200 ) -> list[MemoryItem]: 

201 """ 

202 Search memories. 

203 

204 Args: 

205 memory_type: Filter by type (None = all) 

206 limit: Max results 

207 

208 Returns: 

209 List of MemoryItem, sorted by importance 

210 """ 

211 if memory_type: 

212 items = list(self._memories[memory_type].values()) 

213 else: 

214 items = [] 

215 for mems in self._memories.values(): 

216 items.extend(mems.values()) 

217 

218 # Sort by importance (descending) 

219 items.sort(key=lambda x: x.importance, reverse=True) 

220 

221 return items[:limit] 

222 

223 def forget(self, key: str) -> bool: 

224 """ 

225 Forget a memory item. 

226 

227 Args: 

228 key: Memory key 

229 

230 Returns: 

231 True if forgotten, False if not found 

232 """ 

233 item = self._index.get(key) 

234 if not item: 

235 return False 

236 

237 # Remove from storage 

238 del self._memories[item.type][key] 

239 del self._index[key] 

240 

241 return True 

242 

243 def _evict_working(self) -> None: 

244 """Evict least important working memories.""" 

245 items = list(self._memories[MemoryType.WORKING].values()) 

246 items.sort(key=lambda x: x.importance) 

247 

248 # Remove bottom 20% 

249 to_remove = items[: len(items) // 5 + 1] 

250 for item in to_remove: 

251 self.forget(item.metadata.get("key", "")) 

252 

253 def _evict_episodic(self) -> None: 

254 """Evict least important episodic memories.""" 

255 items = list(self._memories[MemoryType.EPISODIC].values()) 

256 items.sort(key=lambda x: x.importance) 

257 

258 # Remove bottom 20% 

259 to_remove = items[: len(items) // 5 + 1] 

260 for item in to_remove: 

261 self.forget(item.metadata.get("key", "")) 

262 

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

264 """ 

265 Get memory statistics. 

266 

267 Returns: 

268 Dict with memory counts by type 

269 """ 

270 return { 

271 "working": len(self._memories[MemoryType.WORKING]), 

272 "episodic": len(self._memories[MemoryType.EPISODIC]), 

273 "semantic": len(self._memories[MemoryType.SEMANTIC]), 

274 "procedural": len(self._memories[MemoryType.PROCEDURAL]), 

275 "total": sum(len(m) for m in self._memories.values()), 

276 } 

277 

278 def clear(self, memory_type: MemoryType | None = None) -> None: 

279 """ 

280 Clear memories. 

281 

282 Args: 

283 memory_type: Type to clear (None = all) 

284 """ 

285 if memory_type: 

286 self._memories[memory_type].clear() 

287 # Rebuild index 

288 self._index.clear() 

289 for mems in self._memories.values(): 

290 for item in mems.values(): 

291 self._index[item.metadata.get("key", item.id)] = item 

292 else: 

293 for mems in self._memories.values(): 

294 mems.clear() 

295 self._index.clear() 

296 

297 # ── Persistence (v1.14.9) ──────────────── 

298 

299 def get_state(self) -> dict[str, Any]: 

300 """Export full memory state for persistence.""" 

301 return { 

302 "max_working": self.max_working, 

303 "max_episodic": self.max_episodic, 

304 "memories": { 

305 mt.value: {key: item.to_dict() for key, item in mems.items()} 

306 for mt, mems in self._memories.items() 

307 }, 

308 } 

309 

310 def restore_state(self, state: dict[str, Any]) -> None: 

311 """Restore memory state from a persisted snapshot.""" 

312 self.max_working = state.get("max_working", self.max_working) 

313 self.max_episodic = state.get("max_episodic", self.max_episodic) 

314 self._memories = {mt: {} for mt in MemoryType} 

315 self._index.clear() 

316 

317 memories_data = state.get("memories", {}) 

318 for mt_str, items_dict in memories_data.items(): 

319 try: 

320 mt = MemoryType(mt_str) 

321 except ValueError: 

322 continue 

323 for key, item_data in items_dict.items(): 

324 item = MemoryItem.from_dict(item_data) 

325 self._memories[mt][key] = item 

326 self._index[key] = item