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 Enum 

17from typing import Any, Optional 

18 

19 

20class MemoryType(str, Enum): 

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

22 WORKING = "working" # Current task context 

23 EPISODIC = "episodic" # Past experiences 

24 SEMANTIC = "semantic" # Facts and knowledge 

25 PROCEDURAL = "procedural" # Skills and procedures 

26 

27 

28class MemoryLayer(str, Enum): 

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

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

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

32 

33 

34@dataclass 

35class MemoryItem: 

36 """ 

37 Single memory item. 

38 

39 Attributes: 

40 id: Unique identifier 

41 type: Memory type 

42 layer: Memory layer (L1/L2) 

43 content: Memory content 

44 metadata: Additional metadata 

45 created_at: Creation timestamp 

46 accessed_at: Last access timestamp 

47 access_count: Number of accesses 

48 importance: Importance score (0-1) 

49 """ 

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

51 type: MemoryType = MemoryType.WORKING 

52 layer: MemoryLayer = MemoryLayer.L1 

53 content: Any = None 

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

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

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

57 access_count: int = 0 

58 importance: float = 0.5 

59 

60 def access(self) -> None: 

61 """Mark as accessed.""" 

62 self.accessed_at = time.time() 

63 self.access_count += 1 

64 

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

66 """Convert to dict.""" 

67 return { 

68 "id": self.id, 

69 "type": self.type.value, 

70 "layer": self.layer.value, 

71 "content": self.content, 

72 "metadata": self.metadata, 

73 "created_at": self.created_at, 

74 "accessed_at": self.accessed_at, 

75 "access_count": self.access_count, 

76 "importance": self.importance, 

77 } 

78 

79 @classmethod 

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

81 """Create from dict.""" 

82 return cls( 

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

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

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

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

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

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

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

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

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

92 ) 

93 

94 

95class MemoryPyramid: 

96 """ 

97 Multi-layer memory management system. 

98 

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

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

101 

102 Usage: 

103 pyramid = MemoryPyramid() 

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

105 prefs = pyramid.recall("user_preference") 

106 """ 

107 

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

109 """ 

110 Initialize memory pyramid. 

111 

112 Args: 

113 max_working: Max items in working memory 

114 max_episodic: Max items in episodic memory 

115 """ 

116 self.max_working = max_working 

117 self.max_episodic = max_episodic 

118 

119 # Memory storage by type 

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

121 MemoryType.WORKING: {}, 

122 MemoryType.EPISODIC: {}, 

123 MemoryType.SEMANTIC: {}, 

124 MemoryType.PROCEDURAL: {}, 

125 } 

126 

127 # Index for fast lookup 

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

129 

130 def store( 

131 self, 

132 key: str, 

133 content: Any, 

134 memory_type: MemoryType = MemoryType.WORKING, 

135 layer: MemoryLayer = MemoryLayer.L1, 

136 importance: float = 0.5, 

137 **metadata 

138 ) -> MemoryItem: 

139 """ 

140 Store a memory item. 

141 

142 Args: 

143 key: Memory key 

144 content: Memory content 

145 memory_type: Type of memory 

146 layer: Memory layer 

147 importance: Importance score (0-1) 

148 **metadata: Additional metadata 

149 

150 Returns: 

151 Created MemoryItem 

152 """ 

153 # Check capacity for working memory 

154 if memory_type == MemoryType.WORKING: 

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

156 self._evict_working() 

157 

158 # Check capacity for episodic memory 

159 if memory_type == MemoryType.EPISODIC: 

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

161 self._evict_episodic() 

162 

163 # Create memory item 

164 item = MemoryItem( 

165 type=memory_type, 

166 layer=layer, 

167 content=content, 

168 metadata=metadata, 

169 importance=importance, 

170 ) 

171 

172 # Store 

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

174 self._index[key] = item 

175 

176 return item 

177 

178 def recall(self, key: str) -> Optional[MemoryItem]: 

179 """ 

180 Recall a memory item. 

181 

182 Args: 

183 key: Memory key 

184 

185 Returns: 

186 MemoryItem if found, None otherwise 

187 """ 

188 item = self._index.get(key) 

189 if item: 

190 item.access() 

191 return item 

192 

193 def search( 

194 self, 

195 memory_type: Optional[MemoryType] = None, 

196 limit: int = 10, 

197 ) -> list[MemoryItem]: 

198 """ 

199 Search memories. 

200 

201 Args: 

202 memory_type: Filter by type (None = all) 

203 limit: Max results 

204 

205 Returns: 

206 List of MemoryItem, sorted by importance 

207 """ 

208 if memory_type: 

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

210 else: 

211 items = [] 

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

213 items.extend(mems.values()) 

214 

215 # Sort by importance (descending) 

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

217 

218 return items[:limit] 

219 

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

221 """ 

222 Forget a memory item. 

223 

224 Args: 

225 key: Memory key 

226 

227 Returns: 

228 True if forgotten, False if not found 

229 """ 

230 item = self._index.get(key) 

231 if not item: 

232 return False 

233 

234 # Remove from storage 

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

236 del self._index[key] 

237 

238 return True 

239 

240 def _evict_working(self) -> None: 

241 """Evict least important working memories.""" 

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

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

244 

245 # Remove bottom 20% 

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

247 for item in to_remove: 

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

249 

250 def _evict_episodic(self) -> None: 

251 """Evict least important episodic memories.""" 

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

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

254 

255 # Remove bottom 20% 

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

257 for item in to_remove: 

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

259 

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

261 """ 

262 Get memory statistics. 

263 

264 Returns: 

265 Dict with memory counts by type 

266 """ 

267 return { 

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

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

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

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

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

273 } 

274 

275 def clear(self, memory_type: Optional[MemoryType] = None) -> None: 

276 """ 

277 Clear memories. 

278 

279 Args: 

280 memory_type: Type to clear (None = all) 

281 """ 

282 if memory_type: 

283 self._memories[memory_type].clear() 

284 # Rebuild index 

285 self._index.clear() 

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

287 for item in mems.values(): 

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

289 else: 

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

291 mems.clear() 

292 self._index.clear() 

293 

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

295 

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

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

298 return { 

299 "max_working": self.max_working, 

300 "max_episodic": self.max_episodic, 

301 "memories": { 

302 mt.value: { 

303 key: item.to_dict() for key, item in mems.items() 

304 } 

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

306 }, 

307 } 

308 

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

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

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

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

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

314 self._index.clear() 

315 

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

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

318 try: 

319 mt = MemoryType(mt_str) 

320 except ValueError: 

321 continue 

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

323 item = MemoryItem.from_dict(item_data) 

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

325 self._index[key] = item