Memory Management

Memory Management for Long-Running Agents

Overview

Long-running agents face a fundamental problem: they can't remember everything, but forgetting the wrong thing breaks their usefulness. This skill covers memory architectures that balance context retention, token budget, and retrieval accuracy for agents that run for hours, days, or continuously.

Core Concepts

The Memory Problem

Memory Tiers

Step-by-Step Implementation

Step 1: Build a Tiered Memory System

from dataclasses import dataclass, field
from typing import Optional
import json
import time

@dataclass
class MemoryEntry:
    content: str
    timestamp: float = None
    importance: float = 0.5  # 0.0 (trivial) to 1.0 (critical)
    tags: list[str] = field(default_factory=list)
    token_count: int = 0
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = time.time()

class TieredMemory:
    """Multi-tier memory with automatic promotion and demotion."""
    
    def __init__(self, llm, vector_store, max_context_tokens: int = 8000):
        self.llm = llm
        self.vector_store = vector_store
        self.max_context_tokens = max_context_tokens
        
        # L1: Working context (in-memory)
        self.working_memory: list[MemoryEntry] = []
        self.current_tokens = 0
        
        # L2: Recent history buffer
        self.recent_buffer: list[MemoryEntry] = []
        self.buffer_size = 50
        
        # L3: Episodic memory
        self.episodes: list[MemoryEntry] = []
        
        # L4: Semantic memory (vector DB)
        # Initialized externally
    
    async def remember(self, content: str, importance: float = 0.5, 
                       tags: list[str] = None):
        """Store a new memory across tiers."""
        entry = MemoryEntry(
            content=content,
            importance=importance,
            tags=tags or [],
            token_count=self._count_tokens(content)
        )
        
        # Always add to working memory
        self.working_memory.append(entry)
        self.current_tokens += entry.token_count
        
        # If important, store in episodic + semantic
        if importance > 0.7:
            self.episodes.append(entry)
            await self.vector_store.store(entry)
        
        # Trim if needed
        await self._trim_working_memory()

Step 2: Implement Context Window Management

class ContextManager:
    """Optimize what stays in the context window."""
    
    def __init__(self, tiered_memory: TieredMemory, 
                 summarizer, max_tokens: int = 8000):
        self.memory = tiered_memory
        self.summarizer = summarizer
        self.max_tokens = max_tokens
        self.reserved_tokens = 2000  # Reserve for new input/output
    
    async def build_context(self, task: str, top_k: int = 5) -> list[dict]:
        """Build the optimal context for a task."""
        
        available_tokens = self.max_tokens - self.reserved_tokens
        
        # 1. Start with high-importance working memory
        context = []
        tokens_used = 0
        
        working = sorted(
            self.memory.working_memory,
            key=lambda e: e.importance,
            reverse=True
        )
        
        for entry in working:
            if tokens_used + entry.token_count > available_tokens:
                break
            context.append({"role": "system", "content": entry.content})
            tokens_used += entry.token_count
        
        # 2. Add semantically relevant memories
        relevant = await self.memory.vector_store.search(task, k=top_k)
        for mem in relevant:
            if tokens_used + mem.token_count > available_tokens:
                break
            context.append({"role": "system", "content": mem.content})
            tokens_used += mem.token_count
        
        # 3. If we had to drop items, add a summary
        if len(context) < len(working):
            summary = await self._get_summary()
            context.insert(0, {"role": "system", 
                               "content": f"[Summary of earlier context]: {summary}"})
        
        return context
    
    async def _get_summary(self) -> str:
        """Summarize what was excluded from context."""
        excluded = self.memory.working_memory[
            len(self.memory.working_memory) - 10:
        ]
        texts = [e.content for e in excluded]
        return await self.summarizer.summarize("\n".join(texts))
    
    async def _trim_working_memory(self):
        """Reduce working memory when over capacity."""
        while self.memory.current_tokens > self.max_tokens * 0.8:
            # Remove lowest-importance items
            self.memory.working_memory.sort(
                key=lambda e: e.importance
            )
            removed = self.memory.working_memory.pop(0)
            self.memory.current_tokens -= removed.token_count

Step 3: Memory Summarization Strategies

class MemorySummarizer:
    """Different summarization strategies for different memory types."""
    
    def __init__(self, llm):
        self.llm = llm
    
    async def rolling_summary(self, conversation: list[str], 
                              window: int = 20) -> str:
        """Summarize recent conversation window."""
        recent = conversation[-window:]
        return await self.llm.generate(
            f"Summarize this conversation concisely, preserving key facts, "
            f"decisions, and user preferences:\n\n{chr(10).join(recent)}"
        )
    
    async def hierarchical_summary(self, episodes: list[MemoryEntry], 
                                   level: int = 1) -> str:
        """Multi-level summarization for long-running agents."""
        if len(episodes) < 10:
            # Base case: summarize directly
            texts = [e.content for e in episodes]
            return await self.llm.generate(
                f"Summarize these episodes:\n\n{chr(10).join(texts)}"
            )
        
        # Recursive: summarize groups, then summarize summaries
        groups = [
            episodes[i:i+10] 
            for i in range(0, len(episodes), 10)
        ]
        summaries = []
        for group in groups:
            summary = await self.hierarchical_summary(group, level + 1)
            summaries.append(summary)
        
        return await self.llm.generate(
            f"Synthesize these summaries into a higher-level overview:\n\n"
            f"{chr(10).join(summaries)}"
        )
    
    async def importance_weighted_summary(self, episodes: list[MemoryEntry],
                                          max_tokens: int = 500) -> str:
        """Prioritize important memories in summary."""
        # Sort by importance, keep top items
        sorted_eps = sorted(episodes, key=lambda e: e.importance, reverse=True)
        
        important = [e for e in sorted_eps if e.importance > 0.7]
        routine = [e for e in sorted_eps if e.importance <= 0.7]
        
        result = "## Key Events\n"
        result += "\n".join(e.content for e in important[:5])
        
        if routine:
            brief = await self.llm.generate(
                f"Summarize these routine events in one sentence:\n"
                f"{chr(10).join(e.content[:3] for e in routine[:10])}"
            )
            result += f"\n## Other Events\n{brief}"
        
        return result

Step 4: Memory Consolidation & GC

class MemoryConsolidator:
    """Periodically consolidate, prune, and optimize memory."""
    
    def __init__(self, memory: TieredMemory, llm, 
                 consolidation_interval: int = 3600):
        self.memory = memory
        self.llm = llm
        self.interval = consolidation_interval
        self.last_consolidation = time.time()
    
    async def consolidate_if_needed(self):
        """Run consolidation if interval has elapsed."""
        if time.time() - self.last_consolidation > self.interval:
            await self.consolidate()
            self.last_consolidation = time.time()
    
    async def consolidate(self):
        """Merge, prune, and optimize memory store."""
        
        # Phase 1: Deduplicate
        await self._deduplicate()
        
        # Phase 2: Merge related entries
        await self._merge_related()
        
        # Phase 3: Prune low-importance old entries
        await self._prune()
        
        # Phase 4: Re-index vector store
        await self._reindex()
    
    async def _deduplicate(self):
        """Remove duplicate or near-duplicate entries."""
        seen = set()
        unique = []
        for entry in self.memory.episodes:
            # Use first 100 chars as fingerprint
            fingerprint = entry.content[:100]
            if fingerprint not in seen:
                seen.add(fingerprint)
                unique.append(entry)
        self.memory.episodes = unique
    
    async def _merge_related(self):
        """Merge related memories into composite entries."""
        # Group by tags
        from collections import defaultdict
        tagged = defaultdict(list)
        for entry in self.memory.episodes:
            for tag in entry.tags:
                tagged[tag].append(entry)
        
        # Merge groups with >5 entries
        for tag, entries in tagged.items():
            if len(entries) > 5:
                merged = await self.llm.generate(
                    f"Merge these related memories into one coherent summary:\n"
                    f"{chr(10).join(e.content for e in entries)}"
                )
                # Replace with merged entry
                self.memory.episodes = [
                    e for e in self.memory.episodes 
                    if e not in entries
                ]
                self.memory.episodes.append(MemoryEntry(
                    content=merged,
                    importance=0.8,
                    tags=[tag],
                    timestamp=time.time()
                ))
    
    async def _prune(self, max_episodes: int = 1000, 
                     max_age_days: int = 30):
        """Remove old, low-importance entries."""
        now = time.time()
        day = 86400
        
        self.memory.episodes = [
            e for e in self.memory.episodes
            if (e.importance > 0.3 or 
                (now - e.timestamp) < max_age_days * day)
        ]
        
        # If still over limit, remove lowest importance
        if len(self.memory.episodes) > max_episodes:
            self.memory.episodes.sort(
                key=lambda e: (e.importance, e.timestamp),
                reverse=True
            )
            self.memory.episodes = self.memory.episodes[:max_episodes]

Step 5: Memory Retrieval with Reranking

class MemoryRetriever:
    """Retrieve relevant memories with multi-stage ranking."""
    
    def __init__(self, vector_store, llm):
        self.vector_store = vector_store
        self.llm = llm
    
    async def retrieve(self, query: str, k: int = 10, rerank_top: int = 5):
        """Retrieve and rerank memories."""
        
        # Stage 1: Quick vector search (get more than needed)
        candidates = await self.vector_store.search(query, k=k * 3)
        
        # Stage 2: Rerank with LLM
        scored = []
        for mem in candidates:
            score = await self._relevance_score(query, mem.content)
            scored.append((score, mem))
        
        scored.sort(key=lambda x: x[0], reverse=True)
        
        # Stage 3: Return top results
        return [mem for _, mem in scored[:rerank_top]]
    
    async def _relevance_score(self, query: str, memory: str) -> float:
        """Score how relevant a memory is to the query."""
        prompt = f"""Rate the relevance of this memory to the query from 0.0 to 1.0.
Only return a number, nothing else.

Query: {query}
Memory: {memory}
Relevance:"""
        
        response = await self.llm.generate(prompt, temperature=0)
        try:
            return float(response.strip())
        except ValueError:
            return 0.5  # Default on parse failure

Memory Budget Planning

Estimating Memory Costs

Optimization Levers

Trigger Phrases

Anti-Patterns