Coverage for agentos/swarm/agent_memory.py: 23%
315 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1"""
2v1.9.7: Agent Memory System — layered memory with context window management.
4Three-tier memory architecture:
5- WorkingMemory: current task context, small capacity, fast access
6- ShortTermMemory: recent N conversation rounds, sliding window with summarization
7- LongTermMemory: vector-based semantic retrieval for historical knowledge
9ContextWindowManager: auto-trim/compress context to fit token budgets.
10"""
12from __future__ import annotations
14import json
15import time
16import uuid
17from collections import OrderedDict, deque
18from dataclasses import dataclass, field
19from typing import Any
22@dataclass
23class MemoryEntry:
24 """A single memory entry."""
26 id: str = field(default_factory=lambda: uuid.uuid4().hex[:8])
27 content: str = ""
28 role: str = "system" # system, user, assistant, tool
29 timestamp: float = field(default_factory=time.time)
30 importance: float = 0.5 # 0.0-1.0
31 ttl: float = 0.0 # Time-to-live in seconds, 0 = never expire
32 metadata: dict[str, Any] = field(default_factory=dict)
33 embedding: list[float] | None = None # For long-term vector search
34 summary: str = "" # Compressed version for context window
37# ── Working Memory ────────────────────────────────────────────────
40class WorkingMemory:
41 """Ultra-fast, small-capacity memory for the current task.
43 Holds task description, active goals, intermediate results.
44 Max entries enforced — oldest evicted on overflow.
45 """
47 def __init__(self, max_entries: int = 20):
48 self.max_entries = max_entries
49 self._entries: OrderedDict[str, MemoryEntry] = OrderedDict()
50 self.task_goal: str = ""
51 self.active_subtask: str = ""
52 self.scratchpad: dict[str, Any] = {}
54 def add(self, entry: MemoryEntry) -> None:
55 self._entries[entry.id] = entry
56 while len(self._entries) > self.max_entries:
57 self._entries.popitem(last=False)
59 def set_task(self, goal: str, subtask: str = "") -> None:
60 self.task_goal = goal
61 self.active_subtask = subtask or goal
63 def get_all(self) -> list[MemoryEntry]:
64 return list(self._entries.values())
66 def get_last(self, n: int = 5) -> list[MemoryEntry]:
67 return list(self._entries.values())[-n:]
69 def clear(self) -> None:
70 self._entries.clear()
71 self.scratchpad.clear()
73 def to_context(self, max_tokens: int = 500) -> str:
74 """Serialize working memory as context string for LLM."""
75 parts = []
76 if self.task_goal:
77 parts.append(f"[Task] {self.task_goal}")
78 if self.active_subtask and self.active_subtask != self.task_goal:
79 parts.append(f"[SubTask] {self.active_subtask}")
80 for entry in list(self._entries.values())[-5:]:
81 content = entry.summary or entry.content
82 if len(content) > 200:
83 content = content[:197] + "..."
84 parts.append(f"[{entry.role}] {content}")
85 result = "\n".join(parts)
86 if self._estimate_tokens(result) > max_tokens:
87 # Truncate from front
88 lines = result.split("\n")
89 while lines and self._estimate_tokens("\n".join(lines)) > max_tokens:
90 lines.pop(0)
91 result = "\n".join(lines)
92 return result
94 def _estimate_tokens(self, text: str) -> int:
95 """Rough token estimation: ~4 chars per token."""
96 return max(1, len(text) // 4)
99# ── Short-Term Memory ─────────────────────────────────────────────
102class ShortTermMemory:
103 """Sliding window of recent conversation rounds.
105 Auto-summarizes old rounds to maintain a compact window.
106 Supports importance-based retention.
107 """
109 def __init__(
110 self,
111 max_rounds: int = 50,
112 auto_summarize: bool = True,
113 summarize_threshold: int = 20, # Summarize when rounds > threshold
114 keep_recent: int = 10, # Keep N most recent rounds raw
115 ):
116 self.max_rounds = max_rounds
117 self.auto_summarize = auto_summarize
118 self.summarize_threshold = summarize_threshold
119 self.keep_recent = keep_recent
121 self._rounds: deque[list[MemoryEntry]] = deque()
122 self._summaries: list[str] = [] # Compressed old rounds
123 self.total_rounds = 0
125 def add_round(self, entries: list[MemoryEntry]) -> None:
126 """Add a full conversation round."""
127 self._rounds.append(entries)
128 self.total_rounds += 1
130 # Enforce max rounds
131 while len(self._rounds) > self.max_rounds:
132 evicted = self._rounds.popleft()
133 if self.auto_summarize:
134 summary = self._summarize_round(evicted)
135 if summary:
136 self._summaries.append(summary)
138 # Auto-summarize middle rounds when over threshold
139 if self.auto_summarize and len(self._rounds) > self.summarize_threshold:
140 self._compress_middle()
142 def _compress_middle(self) -> None:
143 """Compress rounds between recent keepers and front."""
144 keep_count = min(self.keep_recent, len(self._rounds))
145 recent = list(self._rounds)[-keep_count:]
146 middle = list(self._rounds)[:-keep_count] if keep_count < len(self._rounds) else []
148 if not middle:
149 return
151 # Summarize middle rounds
152 for entries in middle:
153 summary = self._summarize_round(entries)
154 if summary:
155 self._summaries.append(summary)
157 # Replace deque with only recent rounds
158 self._rounds = deque(recent)
160 def _summarize_round(self, entries: list[MemoryEntry]) -> str:
161 """Create a compressed summary of a round."""
162 if not entries:
163 return ""
165 # Collect key content
166 parts = []
167 for entry in entries:
168 content = entry.content
169 if len(content) > 100:
170 content = content[:97] + "..."
171 parts.append(f"{entry.role}: {content}")
173 if not parts:
174 return ""
176 timestamp = entries[0].timestamp if entries else time.time()
177 return f"[Round@{timestamp:.0f}] " + " | ".join(parts)
179 def get_context(
180 self,
181 include_summaries: bool = True,
182 max_rounds: int = 15,
183 ) -> list[MemoryEntry]:
184 """Get flattened context entries."""
185 flat: list[MemoryEntry] = []
187 # Add summaries as system entries
188 if include_summaries:
189 for summary in self._summaries[-3:]: # Keep last 3 summaries
190 flat.append(
191 MemoryEntry(
192 content=f"[History Summary] {summary}",
193 role="system",
194 importance=0.3,
195 )
196 )
198 # Add recent rounds
199 recent_rounds = list(self._rounds)[-max_rounds:]
200 for entries in recent_rounds:
201 flat.extend(entries)
203 return flat
205 def clear(self) -> None:
206 self._rounds.clear()
207 self._summaries.clear()
208 self.total_rounds = 0
211# ── Long-Term Memory ──────────────────────────────────────────────
214class LongTermMemory:
215 """Vector-based semantic memory for historical knowledge retrieval.
217 Stores important memories with embeddings. Supports cosine-similarity search.
218 Falls back to keyword search when no embeddings available.
219 """
221 def __init__(
222 self,
223 max_entries: int = 10000,
224 importance_threshold: float = 0.4, # Only store entries above this importance
225 persist_path: str = "",
226 ):
227 self.max_entries = max_entries
228 self.importance_threshold = importance_threshold
229 self.persist_path = persist_path
231 self._entries: dict[str, MemoryEntry] = {}
232 self._embeddings: dict[str, list[float]] = {} # entry_id → embedding
234 self._embedder: Any = None # Lazy-loaded embedder
236 def add(self, entry: MemoryEntry) -> None:
237 """Store a memory entry. Only stores if importance >= threshold."""
238 if entry.importance < self.importance_threshold:
239 return
241 self._entries[entry.id] = entry
242 if entry.embedding:
243 self._embeddings[entry.id] = entry.embedding
245 # Evict oldest if over capacity
246 while len(self._entries) > self.max_entries:
247 oldest_id = min(
248 self._entries.keys(),
249 key=lambda k: self._entries[k].timestamp,
250 )
251 del self._entries[oldest_id]
252 self._embeddings.pop(oldest_id, None)
254 def search(
255 self,
256 query: str,
257 top_k: int = 5,
258 query_embedding: list[float] | None = None,
259 ) -> list[MemoryEntry]:
260 """Semantic search over stored memories.
262 Uses cosine similarity if embeddings available, else keyword overlap.
263 """
264 if self._embeddings and query_embedding:
265 return self._vector_search(query_embedding, top_k)
266 return self._keyword_search(query, top_k)
268 def search_keywords(
269 self,
270 keywords: list[str],
271 top_k: int = 5,
272 ) -> list[MemoryEntry]:
273 """Search memories by keyword overlap."""
274 results = []
275 for entry in self._entries.values():
276 content_lower = entry.content.lower()
277 score = sum(1 for kw in keywords if kw.lower() in content_lower)
278 if score > 0:
279 results.append((score, entry))
280 results.sort(key=lambda x: (-x[0], -x[1].importance))
281 return [entry for _, entry in results[:top_k]]
283 def search_by_timerange(
284 self,
285 start: float,
286 end: float | None = None,
287 top_k: int = 10,
288 ) -> list[MemoryEntry]:
289 """Search memories by time range."""
290 end = end or time.time()
291 results = [entry for entry in self._entries.values() if start <= entry.timestamp <= end]
292 results.sort(key=lambda e: e.timestamp, reverse=True)
293 return results[:top_k]
295 def _vector_search(
296 self,
297 query_emb: list[float],
298 top_k: int,
299 ) -> list[MemoryEntry]:
300 """Cosine similarity search."""
301 scores = []
302 for eid, emb in self._embeddings.items():
303 sim = self._cosine_similarity(query_emb, emb)
304 scores.append((sim, eid))
305 scores.sort(reverse=True)
306 return [self._entries[eid] for _, eid in scores[:top_k] if eid in self._entries]
308 def _keyword_search(self, query: str, top_k: int) -> list[MemoryEntry]:
309 """Fallback keyword overlap search."""
310 query_words = set(query.lower().split())
311 if not query_words:
312 return []
313 return self.search_keywords(list(query_words), top_k)
315 def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
316 """Cosine similarity between two vectors."""
317 if len(a) != len(b):
318 return 0.0
319 dot = sum(x * y for x, y in zip(a, b))
320 norm_a = sum(x * x for x in a) ** 0.5
321 norm_b = sum(y * y for y in b) ** 0.5
322 if norm_a == 0 or norm_b == 0:
323 return 0.0
324 return dot / (norm_a * norm_b)
326 def export_important(self, top_k: int = 20) -> list[MemoryEntry]:
327 """Export most important entries."""
328 entries = sorted(
329 self._entries.values(),
330 key=lambda e: (e.importance, e.timestamp),
331 reverse=True,
332 )
333 return entries[:top_k]
335 def clear(self) -> None:
336 self._entries.clear()
337 self._embeddings.clear()
339 def save(self, path: str = "") -> None:
340 """Persist to disk (without embeddings)."""
341 save_path = path or self.persist_path
342 if not save_path:
343 return
345 data = []
346 for entry in self._entries.values():
347 data.append(
348 {
349 "id": entry.id,
350 "content": entry.content,
351 "role": entry.role,
352 "timestamp": entry.timestamp,
353 "importance": entry.importance,
354 "metadata": entry.metadata,
355 }
356 )
358 with open(save_path, "w") as f:
359 json.dump(data, f, ensure_ascii=False, indent=2)
361 def load(self, path: str = "") -> int:
362 """Load from disk."""
363 load_path = path or self.persist_path
364 if not load_path:
365 return 0
367 try:
368 with open(load_path) as f:
369 data = json.load(f)
370 except (FileNotFoundError, json.JSONDecodeError):
371 return 0
373 count = 0
374 for item in data:
375 entry = MemoryEntry(
376 id=item.get("id", uuid.uuid4().hex[:8]),
377 content=item.get("content", ""),
378 role=item.get("role", "system"),
379 timestamp=item.get("timestamp", time.time()),
380 importance=item.get("importance", 0.5),
381 metadata=item.get("metadata", {}),
382 )
383 self._entries[entry.id] = entry
384 count += 1
386 return count
389# ── Context Window Manager ────────────────────────────────────────
392@dataclass
393class ContextBudget:
394 """Token budget for context window management."""
396 total_tokens: int = 4096 # Max total tokens
397 system_reserved: int = 512 # Reserved for system prompt
398 working_memory_budget: int = 640 # Budget for working memory
399 short_term_budget: int = 1536 # Budget for short-term memory
400 long_term_budget: int = 512 # Budget for injected long-term memories
401 query_budget: int = 896 # Budget for current query
402 safety_margin: int = 128 # Safety margin
405class ContextWindowManager:
406 """Manages token budgets and assembles context windows.
408 Automatically trims/compresses content to fit within token budgets.
409 Handles the three-tier memory system's context assembly.
410 """
412 def __init__(self, budget: ContextBudget | None = None):
413 self.budget = budget or ContextBudget()
415 def assemble(
416 self,
417 working: WorkingMemory,
418 short_term: ShortTermMemory,
419 long_term: LongTermMemory,
420 current_query: str = "",
421 retrieval_query: str = "",
422 ) -> str:
423 """Assemble a full context window from all memory tiers.
425 Returns a string ready to prepend to the LLM prompt.
426 """
427 sections = []
429 # 1. Working memory context
430 wm_ctx = working.to_context(max_tokens=self.budget.working_memory_budget)
431 if wm_ctx:
432 sections.append(("Working Memory", wm_ctx, self.budget.working_memory_budget))
434 # 2. Short-term memory context
435 st_entries = short_term.get_context(include_summaries=True, max_rounds=15)
436 st_ctx = self._entries_to_context(st_entries, self.budget.short_term_budget)
437 if st_ctx:
438 sections.append(("Recent History", st_ctx, self.budget.short_term_budget))
440 # 3. Long-term memory (semantic retrieval)
441 lt_entries = []
442 if retrieval_query:
443 lt_entries = long_term.search(retrieval_query, top_k=5)
444 else:
445 lt_entries = long_term.export_important(top_k=5)
447 if lt_entries:
448 lt_ctx = self._entries_to_context(lt_entries, self.budget.long_term_budget)
449 if lt_ctx:
450 sections.append(("Relevant Memories", lt_ctx, self.budget.long_term_budget))
452 # 4. Current query
453 query_section = current_query
454 if query_section:
455 est_tokens = self._estimate_tokens(query_section)
456 if est_tokens > self.budget.query_budget:
457 query_section = self._truncate_text(query_section, self.budget.query_budget)
458 sections.append(("Current Task", query_section, self.budget.query_budget))
460 # Assemble final context
461 final_parts = []
462 for name, content, _ in sections:
463 final_parts.append(f"--- {name} ---\n{content}")
465 return "\n\n".join(final_parts)
467 def fit_to_budget(self, text: str, max_tokens: int) -> str:
468 """Trim text to fit within token budget."""
469 if self._estimate_tokens(text) <= max_tokens:
470 return text
471 return self._truncate_text(text, max_tokens)
473 def _entries_to_context(self, entries: list[MemoryEntry], max_tokens: int) -> str:
474 """Convert memory entries to context string, fitting budget."""
475 if not entries:
476 return ""
478 lines = []
479 token_count = 0
481 for entry in entries:
482 content = entry.summary or entry.content
483 line = f"[{entry.role}] {content}"
484 line_tokens = self._estimate_tokens(line)
486 if token_count + line_tokens > max_tokens:
487 # Try truncated version
488 available = max_tokens - token_count - 10
489 if available > 20:
490 truncated = content[: available * 4]
491 line = f"[{entry.role}] {truncated}..."
492 token_count += self._estimate_tokens(line)
493 break
495 lines.append(line)
496 token_count += line_tokens
498 return "\n".join(lines)
500 def _truncate_text(self, text: str, max_tokens: int) -> str:
501 """Truncate text from the beginning to fit token budget."""
502 # Estimate char budget: ~4 chars per token
503 char_budget = max_tokens * 4
504 if len(text) <= char_budget:
505 return text
507 # Keep last char_budget characters for relevance
508 return "...(truncated) " + text[-char_budget:]
510 def _estimate_tokens(self, text: str) -> int:
511 """Rough token estimation."""
512 return max(1, len(text) // 4)
515# ── Unified Agent Memory ──────────────────────────────────────────
518class AgentMemory:
519 """Unified memory system combining all three tiers + context management.
521 High-level API for agent memory operations:
522 - Remember conversation rounds
523 - Retrieve relevant history
524 - Assemble context window
526 Usage:
527 memory = AgentMemory()
528 memory.add_round([user_msg, assistant_msg])
529 context = memory.get_context(query="What files did I create yesterday?")
530 """
532 def __init__(
533 self,
534 working_max: int = 20,
535 short_term_max_rounds: int = 50,
536 long_term_max: int = 10000,
537 budget: ContextBudget | None = None,
538 ):
539 self.working = WorkingMemory(max_entries=working_max)
540 self.short_term = ShortTermMemory(max_rounds=short_term_max_rounds)
541 self.long_term = LongTermMemory(max_entries=long_term_max)
542 self.window_manager = ContextWindowManager(budget=budget)
544 def add_round(
545 self,
546 entries: list[MemoryEntry],
547 importance: float = 0.5,
548 ) -> None:
549 """Add a full conversation round to memory."""
550 self.short_term.add_round(entries)
552 # Store important entries to long-term
553 for entry in entries:
554 if entry.importance >= 0.4:
555 self.long_term.add(entry)
557 def set_task(self, goal: str, subtask: str = "") -> None:
558 """Set current task context in working memory."""
559 self.working.set_task(goal, subtask)
561 def remember(
562 self,
563 content: str,
564 role: str = "system",
565 importance: float = 0.5,
566 ttl: float = 0.0,
567 metadata: dict | None = None,
568 ) -> MemoryEntry:
569 """Store a single memory entry."""
570 entry = MemoryEntry(
571 content=content,
572 role=role,
573 importance=importance,
574 ttl=ttl,
575 metadata=metadata or {},
576 )
577 self.working.add(entry)
578 if importance >= 0.6:
579 self.long_term.add(entry)
580 return entry
582 def recall(
583 self,
584 query: str,
585 top_k: int = 5,
586 include_short_term: bool = True,
587 include_long_term: bool = True,
588 ) -> list[MemoryEntry]:
589 """Search across memory tiers."""
590 results: list[MemoryEntry] = []
592 if include_short_term:
593 st_entries = self.short_term.get_context(include_summaries=True)
594 # Simple keyword filter on short-term
595 query_words = set(query.lower().split())
596 for entry in st_entries:
597 score = sum(1 for w in query_words if w in entry.content.lower())
598 if score > 0:
599 results.append(entry)
601 if include_long_term:
602 lt_results = self.long_term.search(query, top_k=top_k)
603 for entry in lt_results:
604 if entry not in results:
605 results.append(entry)
607 # Sort by importance then timestamp
608 results.sort(key=lambda e: (e.importance, e.timestamp), reverse=True)
609 return results[:top_k]
611 def get_context(self, query: str = "") -> str:
612 """Assemble full context window for the current task."""
613 return self.window_manager.assemble(
614 working=self.working,
615 short_term=self.short_term,
616 long_term=self.long_term,
617 current_query=query,
618 retrieval_query=query,
619 )
621 def clear_working(self) -> None:
622 self.working.clear()
624 def clear_all(self) -> None:
625 self.working.clear()
626 self.short_term.clear()
627 self.long_term.clear()