1"""

2 Memory Pyramid for NexusAgent.

3

4 Multi-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

11 from __future__ import annotations

12

13 import time

14 import uuid

15 from dataclasses import dataclass, field

16 from enum import Enum

17 from typing import Any, Optional

18 from collections import defaultdict

19

20

21 class 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

29 class 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

36 class 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

96 class 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