Coverage for agentos/swarm/agent_memory.py: 23%
315 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +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 ────────────────────────────────────────────────
39class WorkingMemory:
40 """Ultra-fast, small-capacity memory for the current task.
42 Holds task description, active goals, intermediate results.
43 Max entries enforced — oldest evicted on overflow.
44 """
46 def __init__(self, max_entries: int = 20):
47 self.max_entries = max_entries
48 self._entries: OrderedDict[str, MemoryEntry] = OrderedDict()
49 self.task_goal: str = ""
50 self.active_subtask: str = ""
51 self.scratchpad: dict[str, Any] = {}
53 def add(self, entry: MemoryEntry) -> None:
54 self._entries[entry.id] = entry
55 while len(self._entries) > self.max_entries:
56 self._entries.popitem(last=False)
58 def set_task(self, goal: str, subtask: str = "") -> None:
59 self.task_goal = goal
60 self.active_subtask = subtask or goal
62 def get_all(self) -> list[MemoryEntry]:
63 return list(self._entries.values())
65 def get_last(self, n: int = 5) -> list[MemoryEntry]:
66 return list(self._entries.values())[-n:]
68 def clear(self) -> None:
69 self._entries.clear()
70 self.scratchpad.clear()
72 def to_context(self, max_tokens: int = 500) -> str:
73 """Serialize working memory as context string for LLM."""
74 parts = []
75 if self.task_goal:
76 parts.append(f"[Task] {self.task_goal}")
77 if self.active_subtask and self.active_subtask != self.task_goal:
78 parts.append(f"[SubTask] {self.active_subtask}")
79 for entry in list(self._entries.values())[-5:]:
80 content = entry.summary or entry.content
81 if len(content) > 200:
82 content = content[:197] + "..."
83 parts.append(f"[{entry.role}] {content}")
84 result = "\n".join(parts)
85 if self._estimate_tokens(result) > max_tokens:
86 # Truncate from front
87 lines = result.split("\n")
88 while lines and self._estimate_tokens("\n".join(lines)) > max_tokens:
89 lines.pop(0)
90 result = "\n".join(lines)
91 return result
93 def _estimate_tokens(self, text: str) -> int:
94 """Rough token estimation: ~4 chars per token."""
95 return max(1, len(text) // 4)
98# ── Short-Term Memory ─────────────────────────────────────────────
100class ShortTermMemory:
101 """Sliding window of recent conversation rounds.
103 Auto-summarizes old rounds to maintain a compact window.
104 Supports importance-based retention.
105 """
107 def __init__(
108 self,
109 max_rounds: int = 50,
110 auto_summarize: bool = True,
111 summarize_threshold: int = 20, # Summarize when rounds > threshold
112 keep_recent: int = 10, # Keep N most recent rounds raw
113 ):
114 self.max_rounds = max_rounds
115 self.auto_summarize = auto_summarize
116 self.summarize_threshold = summarize_threshold
117 self.keep_recent = keep_recent
119 self._rounds: deque[list[MemoryEntry]] = deque()
120 self._summaries: list[str] = [] # Compressed old rounds
121 self.total_rounds = 0
123 def add_round(self, entries: list[MemoryEntry]) -> None:
124 """Add a full conversation round."""
125 self._rounds.append(entries)
126 self.total_rounds += 1
128 # Enforce max rounds
129 while len(self._rounds) > self.max_rounds:
130 evicted = self._rounds.popleft()
131 if self.auto_summarize:
132 summary = self._summarize_round(evicted)
133 if summary:
134 self._summaries.append(summary)
136 # Auto-summarize middle rounds when over threshold
137 if self.auto_summarize and len(self._rounds) > self.summarize_threshold:
138 self._compress_middle()
140 def _compress_middle(self) -> None:
141 """Compress rounds between recent keepers and front."""
142 keep_count = min(self.keep_recent, len(self._rounds))
143 recent = list(self._rounds)[-keep_count:]
144 middle = list(self._rounds)[:-keep_count] if keep_count < len(self._rounds) else []
146 if not middle:
147 return
149 # Summarize middle rounds
150 for entries in middle:
151 summary = self._summarize_round(entries)
152 if summary:
153 self._summaries.append(summary)
155 # Replace deque with only recent rounds
156 self._rounds = deque(recent)
158 def _summarize_round(self, entries: list[MemoryEntry]) -> str:
159 """Create a compressed summary of a round."""
160 if not entries:
161 return ""
163 # Collect key content
164 parts = []
165 for entry in entries:
166 content = entry.content
167 if len(content) > 100:
168 content = content[:97] + "..."
169 parts.append(f"{entry.role}: {content}")
171 if not parts:
172 return ""
174 timestamp = entries[0].timestamp if entries else time.time()
175 return f"[Round@{timestamp:.0f}] " + " | ".join(parts)
177 def get_context(
178 self,
179 include_summaries: bool = True,
180 max_rounds: int = 15,
181 ) -> list[MemoryEntry]:
182 """Get flattened context entries."""
183 flat: list[MemoryEntry] = []
185 # Add summaries as system entries
186 if include_summaries:
187 for summary in self._summaries[-3:]: # Keep last 3 summaries
188 flat.append(MemoryEntry(
189 content=f"[History Summary] {summary}",
190 role="system",
191 importance=0.3,
192 ))
194 # Add recent rounds
195 recent_rounds = list(self._rounds)[-max_rounds:]
196 for entries in recent_rounds:
197 flat.extend(entries)
199 return flat
201 def clear(self) -> None:
202 self._rounds.clear()
203 self._summaries.clear()
204 self.total_rounds = 0
207# ── Long-Term Memory ──────────────────────────────────────────────
209class LongTermMemory:
210 """Vector-based semantic memory for historical knowledge retrieval.
212 Stores important memories with embeddings. Supports cosine-similarity search.
213 Falls back to keyword search when no embeddings available.
214 """
216 def __init__(
217 self,
218 max_entries: int = 10000,
219 importance_threshold: float = 0.4, # Only store entries above this importance
220 persist_path: str = "",
221 ):
222 self.max_entries = max_entries
223 self.importance_threshold = importance_threshold
224 self.persist_path = persist_path
226 self._entries: dict[str, MemoryEntry] = {}
227 self._embeddings: dict[str, list[float]] = {} # entry_id → embedding
229 self._embedder: Any = None # Lazy-loaded embedder
231 def add(self, entry: MemoryEntry) -> None:
232 """Store a memory entry. Only stores if importance >= threshold."""
233 if entry.importance < self.importance_threshold:
234 return
236 self._entries[entry.id] = entry
237 if entry.embedding:
238 self._embeddings[entry.id] = entry.embedding
240 # Evict oldest if over capacity
241 while len(self._entries) > self.max_entries:
242 oldest_id = min(
243 self._entries.keys(),
244 key=lambda k: self._entries[k].timestamp,
245 )
246 del self._entries[oldest_id]
247 self._embeddings.pop(oldest_id, None)
249 def search(
250 self,
251 query: str,
252 top_k: int = 5,
253 query_embedding: list[float] | None = None,
254 ) -> list[MemoryEntry]:
255 """Semantic search over stored memories.
257 Uses cosine similarity if embeddings available, else keyword overlap.
258 """
259 if self._embeddings and query_embedding:
260 return self._vector_search(query_embedding, top_k)
261 return self._keyword_search(query, top_k)
263 def search_keywords(
264 self,
265 keywords: list[str],
266 top_k: int = 5,
267 ) -> list[MemoryEntry]:
268 """Search memories by keyword overlap."""
269 results = []
270 for entry in self._entries.values():
271 content_lower = entry.content.lower()
272 score = sum(1 for kw in keywords if kw.lower() in content_lower)
273 if score > 0:
274 results.append((score, entry))
275 results.sort(key=lambda x: (-x[0], -x[1].importance))
276 return [entry for _, entry in results[:top_k]]
278 def search_by_timerange(
279 self,
280 start: float,
281 end: float | None = None,
282 top_k: int = 10,
283 ) -> list[MemoryEntry]:
284 """Search memories by time range."""
285 end = end or time.time()
286 results = [
287 entry for entry in self._entries.values()
288 if start <= entry.timestamp <= end
289 ]
290 results.sort(key=lambda e: e.timestamp, reverse=True)
291 return results[:top_k]
293 def _vector_search(
294 self,
295 query_emb: list[float],
296 top_k: int,
297 ) -> list[MemoryEntry]:
298 """Cosine similarity search."""
299 scores = []
300 for eid, emb in self._embeddings.items():
301 sim = self._cosine_similarity(query_emb, emb)
302 scores.append((sim, eid))
303 scores.sort(reverse=True)
304 return [self._entries[eid] for _, eid in scores[:top_k] if eid in self._entries]
306 def _keyword_search(self, query: str, top_k: int) -> list[MemoryEntry]:
307 """Fallback keyword overlap search."""
308 query_words = set(query.lower().split())
309 if not query_words:
310 return []
311 return self.search_keywords(list(query_words), top_k)
313 def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
314 """Cosine similarity between two vectors."""
315 if len(a) != len(b):
316 return 0.0
317 dot = sum(x * y for x, y in zip(a, b))
318 norm_a = sum(x * x for x in a) ** 0.5
319 norm_b = sum(y * y for y in b) ** 0.5
320 if norm_a == 0 or norm_b == 0:
321 return 0.0
322 return dot / (norm_a * norm_b)
324 def export_important(self, top_k: int = 20) -> list[MemoryEntry]:
325 """Export most important entries."""
326 entries = sorted(
327 self._entries.values(),
328 key=lambda e: (e.importance, e.timestamp),
329 reverse=True,
330 )
331 return entries[:top_k]
333 def clear(self) -> None:
334 self._entries.clear()
335 self._embeddings.clear()
337 def save(self, path: str = "") -> None:
338 """Persist to disk (without embeddings)."""
339 save_path = path or self.persist_path
340 if not save_path:
341 return
343 data = []
344 for entry in self._entries.values():
345 data.append({
346 "id": entry.id,
347 "content": entry.content,
348 "role": entry.role,
349 "timestamp": entry.timestamp,
350 "importance": entry.importance,
351 "metadata": entry.metadata,
352 })
354 with open(save_path, "w") as f:
355 json.dump(data, f, ensure_ascii=False, indent=2)
357 def load(self, path: str = "") -> int:
358 """Load from disk."""
359 load_path = path or self.persist_path
360 if not load_path:
361 return 0
363 try:
364 with open(load_path) as f:
365 data = json.load(f)
366 except (FileNotFoundError, json.JSONDecodeError):
367 return 0
369 count = 0
370 for item in data:
371 entry = MemoryEntry(
372 id=item.get("id", uuid.uuid4().hex[:8]),
373 content=item.get("content", ""),
374 role=item.get("role", "system"),
375 timestamp=item.get("timestamp", time.time()),
376 importance=item.get("importance", 0.5),
377 metadata=item.get("metadata", {}),
378 )
379 self._entries[entry.id] = entry
380 count += 1
382 return count
385# ── Context Window Manager ────────────────────────────────────────
387@dataclass
388class ContextBudget:
389 """Token budget for context window management."""
391 total_tokens: int = 4096 # Max total tokens
392 system_reserved: int = 512 # Reserved for system prompt
393 working_memory_budget: int = 640 # Budget for working memory
394 short_term_budget: int = 1536 # Budget for short-term memory
395 long_term_budget: int = 512 # Budget for injected long-term memories
396 query_budget: int = 896 # Budget for current query
397 safety_margin: int = 128 # Safety margin
400class ContextWindowManager:
401 """Manages token budgets and assembles context windows.
403 Automatically trims/compresses content to fit within token budgets.
404 Handles the three-tier memory system's context assembly.
405 """
407 def __init__(self, budget: ContextBudget | None = None):
408 self.budget = budget or ContextBudget()
410 def assemble(
411 self,
412 working: WorkingMemory,
413 short_term: ShortTermMemory,
414 long_term: LongTermMemory,
415 current_query: str = "",
416 retrieval_query: str = "",
417 ) -> str:
418 """Assemble a full context window from all memory tiers.
420 Returns a string ready to prepend to the LLM prompt.
421 """
422 sections = []
424 # 1. Working memory context
425 wm_ctx = working.to_context(max_tokens=self.budget.working_memory_budget)
426 if wm_ctx:
427 sections.append(("Working Memory", wm_ctx, self.budget.working_memory_budget))
429 # 2. Short-term memory context
430 st_entries = short_term.get_context(include_summaries=True, max_rounds=15)
431 st_ctx = self._entries_to_context(st_entries, self.budget.short_term_budget)
432 if st_ctx:
433 sections.append(("Recent History", st_ctx, self.budget.short_term_budget))
435 # 3. Long-term memory (semantic retrieval)
436 lt_entries = []
437 if retrieval_query:
438 lt_entries = long_term.search(retrieval_query, top_k=5)
439 else:
440 lt_entries = long_term.export_important(top_k=5)
442 if lt_entries:
443 lt_ctx = self._entries_to_context(lt_entries, self.budget.long_term_budget)
444 if lt_ctx:
445 sections.append(("Relevant Memories", lt_ctx, self.budget.long_term_budget))
447 # 4. Current query
448 query_section = current_query
449 if query_section:
450 est_tokens = self._estimate_tokens(query_section)
451 if est_tokens > self.budget.query_budget:
452 query_section = self._truncate_text(query_section, self.budget.query_budget)
453 sections.append(("Current Task", query_section, self.budget.query_budget))
455 # Assemble final context
456 final_parts = []
457 for name, content, _ in sections:
458 final_parts.append(f"--- {name} ---\n{content}")
460 return "\n\n".join(final_parts)
462 def fit_to_budget(self, text: str, max_tokens: int) -> str:
463 """Trim text to fit within token budget."""
464 if self._estimate_tokens(text) <= max_tokens:
465 return text
466 return self._truncate_text(text, max_tokens)
468 def _entries_to_context(self, entries: list[MemoryEntry], max_tokens: int) -> str:
469 """Convert memory entries to context string, fitting budget."""
470 if not entries:
471 return ""
473 lines = []
474 token_count = 0
476 for entry in entries:
477 content = entry.summary or entry.content
478 line = f"[{entry.role}] {content}"
479 line_tokens = self._estimate_tokens(line)
481 if token_count + line_tokens > max_tokens:
482 # Try truncated version
483 available = max_tokens - token_count - 10
484 if available > 20:
485 truncated = content[:available * 4]
486 line = f"[{entry.role}] {truncated}..."
487 token_count += self._estimate_tokens(line)
488 break
490 lines.append(line)
491 token_count += line_tokens
493 return "\n".join(lines)
495 def _truncate_text(self, text: str, max_tokens: int) -> str:
496 """Truncate text from the beginning to fit token budget."""
497 # Estimate char budget: ~4 chars per token
498 char_budget = max_tokens * 4
499 if len(text) <= char_budget:
500 return text
502 # Keep last char_budget characters for relevance
503 return "...(truncated) " + text[-char_budget:]
505 def _estimate_tokens(self, text: str) -> int:
506 """Rough token estimation."""
507 return max(1, len(text) // 4)
510# ── Unified Agent Memory ──────────────────────────────────────────
512class AgentMemory:
513 """Unified memory system combining all three tiers + context management.
515 High-level API for agent memory operations:
516 - Remember conversation rounds
517 - Retrieve relevant history
518 - Assemble context window
520 Usage:
521 memory = AgentMemory()
522 memory.add_round([user_msg, assistant_msg])
523 context = memory.get_context(query="What files did I create yesterday?")
524 """
526 def __init__(
527 self,
528 working_max: int = 20,
529 short_term_max_rounds: int = 50,
530 long_term_max: int = 10000,
531 budget: ContextBudget | None = None,
532 ):
533 self.working = WorkingMemory(max_entries=working_max)
534 self.short_term = ShortTermMemory(max_rounds=short_term_max_rounds)
535 self.long_term = LongTermMemory(max_entries=long_term_max)
536 self.window_manager = ContextWindowManager(budget=budget)
538 def add_round(
539 self,
540 entries: list[MemoryEntry],
541 importance: float = 0.5,
542 ) -> None:
543 """Add a full conversation round to memory."""
544 self.short_term.add_round(entries)
546 # Store important entries to long-term
547 for entry in entries:
548 if entry.importance >= 0.4:
549 self.long_term.add(entry)
551 def set_task(self, goal: str, subtask: str = "") -> None:
552 """Set current task context in working memory."""
553 self.working.set_task(goal, subtask)
555 def remember(
556 self,
557 content: str,
558 role: str = "system",
559 importance: float = 0.5,
560 ttl: float = 0.0,
561 metadata: dict | None = None,
562 ) -> MemoryEntry:
563 """Store a single memory entry."""
564 entry = MemoryEntry(
565 content=content,
566 role=role,
567 importance=importance,
568 ttl=ttl,
569 metadata=metadata or {},
570 )
571 self.working.add(entry)
572 if importance >= 0.6:
573 self.long_term.add(entry)
574 return entry
576 def recall(
577 self,
578 query: str,
579 top_k: int = 5,
580 include_short_term: bool = True,
581 include_long_term: bool = True,
582 ) -> list[MemoryEntry]:
583 """Search across memory tiers."""
584 results: list[MemoryEntry] = []
586 if include_short_term:
587 st_entries = self.short_term.get_context(include_summaries=True)
588 # Simple keyword filter on short-term
589 query_words = set(query.lower().split())
590 for entry in st_entries:
591 score = sum(1 for w in query_words if w in entry.content.lower())
592 if score > 0:
593 results.append(entry)
595 if include_long_term:
596 lt_results = self.long_term.search(query, top_k=top_k)
597 for entry in lt_results:
598 if entry not in results:
599 results.append(entry)
601 # Sort by importance then timestamp
602 results.sort(key=lambda e: (e.importance, e.timestamp), reverse=True)
603 return results[:top_k]
605 def get_context(self, query: str = "") -> str:
606 """Assemble full context window for the current task."""
607 return self.window_manager.assemble(
608 working=self.working,
609 short_term=self.short_term,
610 long_term=self.long_term,
611 current_query=query,
612 retrieval_query=query,
613 )
615 def clear_working(self) -> None:
616 self.working.clear()
618 def clear_all(self) -> None:
619 self.working.clear()
620 self.short_term.clear()
621 self.long_term.clear()