Working Memory and Context Windows

Duration: 15 min

Working Memory and Context Windows

Duration: 15 min

Overview

Working memory is the agent's active workspace - the information it can directly reason about right now. In LLM-based agents, this maps to the context window: the tokens currently fed to the model. This module explores how to manage context windows effectively, techniques for fitting more relevant information into limited space, and strategies for deciding what stays and what gets evicted.

Understanding context window management is essential because it directly determines the quality of agent responses. An agent with a poorly managed context window will hallucinate, forget instructions, or lose track of multi-step tasks.

Key Concepts & Foundations

  • What: Working memory is the subset of all agent knowledge actively available for reasoning in a single inference call
  • Why: Context windows are finite (4K to 200K tokens); what you include determines response quality
  • How: Use sliding windows, priority queues, and dynamic allocation to maximize context utility

Context Window Anatomy

A typical agent context window contains:

1. System prompt - Agent identity, capabilities, and rules (fixed) 2. Retrieved context - RAG results, tool outputs, memory recalls (dynamic) 3. Conversation history - Recent messages (sliding window) 4. Current task - The immediate user request (fixed per turn)

Token Budget Allocation

With a 128K context window, a practical allocation might be:

  • System prompt: 2K tokens (fixed overhead)
  • Long-term memory recalls: 4K tokens (retrieved per turn)
  • Conversation history: 8K tokens (sliding window)
  • Tool results: 4K tokens (current turn)
  • User message + response space: remaining tokens

The Recency-Relevance Trade-off

  • Recency bias - Keep the most recent messages (simple, but loses important earlier context)
  • Relevance bias - Keep semantically similar messages (complex, but preserves important context)
  • Hybrid - Keep recent messages + retrieve relevant older ones

Hands-On Implementation

from dataclasses import dataclass
from typing import Optional

@dataclass
class Message:
    role: str  # system, user, assistant, tool
    content: str
    token_count: int = 0
    priority: float = 1.0  # higher = more important to keep
    def __post_init__(self):
        # Rough token estimate: ~4 chars per token
        if self.token_count == 0:
            self.token_count = len(self.content) // 4

class ContextWindowManager:
    """Manages what fits in the agent's context window."""
    def __init__(self, max_tokens: int = 128000):
        self.max_tokens = max_tokens
        self.system_prompt: Optional[Message] = None
        self.messages: list[Message] = []
        self.reserved_tokens = 4000  # Reserve for model response
    @property
    def available_tokens(self) -> int:
        used = self.reserved_tokens
        if self.system_prompt:
            used += self.system_prompt.token_count
        used += sum(m.token_count for m in self.messages)
        return self.max_tokens - used
    def set_system_prompt(self, content: str):
        self.system_prompt = Message(role="system", content=content, priority=10.0)
    def add_message(self, role: str, content: str, priority: float = 1.0):
        msg = Message(role=role, content=content, priority=priority)
        self.messages.append(msg)
        self._enforce_limit()
    def _enforce_limit(self):
        """Evict lowest-priority messages when over budget."""
        while self.available_tokens < 0 and self.messages:
            # Find lowest priority message (but never remove the last user message)
            candidates = self.messages[:-1] if len(self.messages) > 1 else []
            if not candidates:
                break
            lowest = min(candidates, key=lambda m: m.priority)
            self.messages.remove(lowest)
    def build_prompt(self) -> list[dict]:
        """Build the final prompt for the LLM."""
        prompt = []
        if self.system_prompt:
            prompt.append({"role": "system", "content": self.system_prompt.content})
        for msg in self.messages:
            prompt.append({"role": msg.role, "content": msg.content})
        return prompt
    def get_stats(self) -> dict:
        total_used = (
            (self.system_prompt.token_count if self.system_prompt else 0)
            + sum(m.token_count for m in self.messages)
            + self.reserved_tokens
        )
        return {
            "total_capacity": self.max_tokens,
            "used_tokens": total_used,
            "available_tokens": self.available_tokens,
            "message_count": len(self.messages),
            "utilization": f"{total_used / self.max_tokens * 100:.1f}%"
        }


Demo: Managing a conversation with limited context

ctx = ContextWindowManager(max_tokens=2000)  # Small window for demo
ctx.set_system_prompt("You are a helpful coding assistant.")
Simulate a multi-turn conversation

ctx.add_message("user", "What is a Python decorator?", priority=0.5)
ctx.add_message("assistant", "A decorator is a function that wraps another function to extend its behavior without modifying it.", priority=0.5)
ctx.add_message("user", "Show me an example", priority=0.5)
ctx.add_message("assistant", "Here is a simple timing decorator that measures execution time of any function.", priority=0.5)

This high-priority message should survive eviction

ctx.add_message("user", "Remember: I always want type hints in code examples", priority=5.0)
ctx.add_message("user", "Now explain context managers", priority=1.0)
print("Context stats:", ctx.get_stats())
print(f"Messages retained: {len(ctx.messages)}")
for msg in ctx.messages:
    print(f"  [{msg.role}] (priority={msg.priority}) {msg.content[:60]}...")

Advanced Techniques

When managing working memory in production agents:

1. Dynamic Budget Allocation - Adjust token budgets per section based on task complexity 2. Attention Anchoring - Place critical instructions at the start and end of context (LLMs attend more to these positions) 3. Progressive Summarization - As conversation grows, summarize older turns rather than dropping them entirely 4. Context Window Monitoring - Track utilization and alert when approaching limits 5. Multi-Model Routing - Use smaller models for simple tasks (saves context for complex ones)

Quiz

Q1: What happens when an agent's context window is full?

  • A) The model automatically summarizes old messages
  • B) Older or lower-priority information must be evicted or summarized ✓
  • C) The model crashes
  • D) Tokens become cheaper

Q2: Why should you reserve tokens in the context window?

  • A) To save money
  • B) To leave space for the model's response generation ✓
  • C) To make the prompt look shorter
  • D) To comply with API limits

Q3: What is the recency-relevance trade-off?

  • A) Choosing between fast and slow models
  • B) Deciding whether to keep recent messages or semantically relevant older ones ✓
  • C) Picking between short and long prompts
  • D) Trading accuracy for speed

← Back ▶ Run in Colab