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

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 

18from collections import defaultdict 

19 

20 

21class MemoryType(str, Enum): 

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

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(str, Enum): 

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

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

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

33 

34 

35@dataclass 

36class MemoryItem: 

37 """ 

38 Single memory item. 

39 

40 Attributes: 

41 id: Unique identifier 

42 type: Memory type 

43 layer: Memory layer (L1/L2) 

44 content: Memory content 

45 metadata: Additional metadata 

46 created_at: Creation timestamp 

47 accessed_at: Last access timestamp 

48 access_count: Number of accesses 

49 importance: Importance score (0-1) 

50 """ 

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

52 type: MemoryType = MemoryType.WORKING 

53 layer: MemoryLayer = MemoryLayer.L1 

54 content: Any = None 

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

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

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

58 access_count: int = 0 

59 importance: float = 0.5 

60 

61 def access(self) -> None: 

62 """Mark as accessed.""" 

63 self.accessed_at = time.time() 

64 self.access_count += 1 

65 

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

67 """Convert to dict.""" 

68 return { 

69 "id": self.id, 

70 "type": self.type.value, 

71 "layer": self.layer.value, 

72 "content": self.content, 

73 "metadata": self.metadata, 

74 "created_at": self.created_at, 

75 "accessed_at": self.accessed_at, 

76 "access_count": self.access_count, 

77 "importance": self.importance, 

78 } 

79 

80 @classmethod 

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

82 """Create from dict.""" 

83 return cls( 

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

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

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

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

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

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

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

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

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

93 ) 

94 

95 

96class MemoryPyramid: 

97 """ 

98 Multi-layer memory management system. 

99 

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

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

102 

103 Usage: 

104 pyramid = MemoryPyramid() 

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

106 prefs = pyramid.recall("user_preference") 

107 """ 

108 

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

110 """ 

111 Initialize memory pyramid. 

112 

113 Args: 

114 max_working: Max items in working memory 

115 max_episodic: Max items in episodic memory 

116 """ 

117 self.max_working = max_working 

118 self.max_episodic = max_episodic 

119 

120 # Memory storage by type 

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

122 MemoryType.WORKING: {}, 

123 MemoryType.EPISODIC: {}, 

124 MemoryType.SEMANTIC: {}, 

125 MemoryType.PROCEDURAL: {}, 

126 } 

127 

128 # Index for fast lookup 

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

130 

131 def store( 

132 self, 

133 key: str, 

134 content: Any, 

135 memory_type: MemoryType = MemoryType.WORKING, 

136 layer: MemoryLayer = MemoryLayer.L1, 

137 importance: float = 0.5, 

138 **metadata 

139 ) -> MemoryItem: 

140 """ 

141 Store a memory item. 

142 

143 Args: 

144 key: Memory key 

145 content: Memory content 

146 memory_type: Type of memory 

147 layer: Memory layer 

148 importance: Importance score (0-1) 

149 **metadata: Additional metadata 

150 

151 Returns: 

152 Created MemoryItem 

153 """ 

154 # Check capacity for working memory 

155 if memory_type == MemoryType.WORKING: 

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

157 self._evict_working() 

158 

159 # Check capacity for episodic memory 

160 if memory_type == MemoryType.EPISODIC: 

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

162 self._evict_episodic() 

163 

164 # Create memory item 

165 item = MemoryItem( 

166 type=memory_type, 

167 layer=layer, 

168 content=content, 

169 metadata=metadata, 

170 importance=importance, 

171 ) 

172 

173 # Store 

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

175 self._index[key] = item 

176 

177 return item 

178 

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

180 """ 

181 Recall a memory item. 

182 

183 Args: 

184 key: Memory key 

185 

186 Returns: 

187 MemoryItem if found, None otherwise 

188 """ 

189 item = self._index.get(key) 

190 if item: 

191 item.access() 

192 return item 

193 

194 def search( 

195 self, 

196 memory_type: Optional[MemoryType] = None, 

197 limit: int = 10, 

198 ) -> list[MemoryItem]: 

199 """ 

200 Search memories. 

201 

202 Args: 

203 memory_type: Filter by type (None = all) 

204 limit: Max results 

205 

206 Returns: 

207 List of MemoryItem, sorted by importance 

208 """ 

209 if memory_type: 

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

211 else: 

212 items = [] 

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

214 items.extend(mems.values()) 

215 

216 # Sort by importance (descending) 

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

218 

219 return items[:limit] 

220 

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

222 """ 

223 Forget a memory item. 

224 

225 Args: 

226 key: Memory key 

227 

228 Returns: 

229 True if forgotten, False if not found 

230 """ 

231 item = self._index.get(key) 

232 if not item: 

233 return False 

234 

235 # Remove from storage 

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

237 del self._index[key] 

238 

239 return True 

240 

241 def _evict_working(self) -> None: 

242 """Evict least important working memories.""" 

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

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

245 

246 # Remove bottom 20% 

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

248 for item in to_remove: 

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

250 

251 def _evict_episodic(self) -> None: 

252 """Evict least important episodic memories.""" 

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

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

255 

256 # Remove bottom 20% 

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

258 for item in to_remove: 

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

260 

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

262 """ 

263 Get memory statistics. 

264 

265 Returns: 

266 Dict with memory counts by type 

267 """ 

268 return { 

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

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

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

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

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

274 } 

275 

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

277 """ 

278 Clear memories. 

279 

280 Args: 

281 memory_type: Type to clear (None = all) 

282 """ 

283 if memory_type: 

284 self._memories[memory_type].clear() 

285 # Rebuild index 

286 self._index.clear() 

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

288 for item in mems.values(): 

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

290 else: 

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

292 mems.clear() 

293 self._index.clear() 

294 

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

296 

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

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

299 return { 

300 "max_working": self.max_working, 

301 "max_episodic": self.max_episodic, 

302 "memories": { 

303 mt.value: { 

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

305 } 

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