Working Memory: Managing Agent Context Windows
Source: mortalapps.com- The AI agent context window functions as its working memory, a finite buffer for current task-relevant information.
- It solves the problem of Large Language Models (LLMs) being stateless by providing a mechanism to inject conversational history, tool definitions, and dynamic data for coherent multi-turn interactions.
- In production, mismanaging context windows leads to high inference costs, truncated responses, hallucinations due to lost context, and degraded agent performance.
- Engineers must strategically allocate token budgets across system prompts, tool descriptions, conversation history, and retrieved data to ensure agents maintain operational effectiveness.
- Context eviction strategies are critical for handling long-running conversations or complex tasks that exceed the LLM's fixed token limit.
Why This Matters
The AI agent context window is the operational memory for Large Language Models (LLMs), analogous to a CPU's L1 cache or a human's short-term working memory. LLMs are inherently stateless; each API call is independent, lacking memory of prior interactions. This approach was invented to overcome this limitation, enabling LLMs to engage in multi-turn conversations, follow complex instructions, and utilize external tools by providing all necessary information within a single, bounded input. Without effective context window management, agents cannot maintain conversational coherence, execute multi-step plans, or leverage tools effectively. In production, ignoring this leads to critical failures: agents forget previous turns, hallucinate due to missing context, or fail to use tools correctly. This directly impacts user experience, increases operational costs through inefficient token usage, and necessitates frequent human intervention. Engineers must understand context windows to build robust, reliable, and cost-efficient agent systems. This approach is essential for any agent requiring statefulness or tool interaction, contrasting with stateless, single-turn LLM calls where context management is minimal.
Core Concepts
The context window serves as the primary working memory for an AI agent, a finite buffer where all relevant information for the current inference call is assembled.
- Context Window: A fixed-size input buffer, measured in tokens, that an LLM can process in a single inference request. It's the agent's immediate operational memory, holding the current state of interaction and necessary instructions.
- Token Budget: The maximum number of tokens an LLM can accept in its input context window. Exceeding this budget requires truncation or summarization of input data.
- System Prompt: The initial, immutable instructions provided to the LLM, defining its persona, capabilities, and constraints. It occupies a critical, often prioritized, portion of the token budget.
- Tool Descriptions (Function Signatures): Declarations of external functions or APIs the agent can call. These descriptions, including their names, parameters, and usage instructions, are injected into the context window to enable tool use.
- Conversation History: A sequence of past user queries and agent responses. This history is dynamically added to the context window to maintain conversational coherence and state over multiple turns.
- Tool Outputs: The results returned from executing an external tool. These outputs are injected back into the context window, allowing the LLM to interpret the results and continue its reasoning process.
- Context Eviction: The process of removing older or less relevant information from the context window when the token budget is approached or exceeded. Common strategies include simple truncation (FIFO) or more sophisticated methods like summarization or relevance-based filtering.
- Tokenization: The process of breaking down raw text into smaller units called tokens. The exact number of tokens for a given text can vary significantly between different LLM models and their respective tokenizers, directly impacting context window utilization.
How It Works
An AI agent's interaction with its context window follows a dynamic, iterative process, managing information flow for each LLM inference call.
1. Initialization and Input Reception
When an agent starts or receives a new user query, the process begins. The agent's orchestrator prepares the initial context, including the static system prompt and any predefined tool descriptions. For subsequent turns, the new user input is added to a temporary buffer.
2. Context Assembly and Prioritization
The agent's context manager dynamically assembles the input for the LLM. This involves retrieving various components and prioritizing them based on their criticality to the current task:
- System Prompt: Always included, as it defines the agent's core behavior.
- Tool Descriptions: Included if the agent is designed to use tools, enabling the LLM to understand available functions.
- Conversation History: Retrieved from an external memory store (e.g., a database or vector store). Recent turns are typically prioritized over older ones.
- Retrieved Documents: If a Retrieval-Augmented Generation (RAG) component is active, relevant documents are fetched based on the current query and added.
- Tool Outputs: If the previous turn involved tool execution, the results are injected to inform the LLM's next action.
- Current User Input: The most recent query from the user.
3. Tokenization and Budget Check
Once assembled, the entire input string is tokenized using the LLM's specific tokenizer. The total token count is then compared against the LLM's maximum context window size (token budget). This step is critical for preventing API errors and managing costs.
4. Context Eviction (if necessary)
If the token count exceeds the budget, an eviction strategy is triggered. Common approaches include:
- First-In, First-Out (FIFO): Oldest conversation turns are removed first.
- Summarization: Older parts of the conversation or less critical retrieved documents are summarized into fewer tokens.
- Relevance-based Eviction: Less relevant pieces of information (e.g., based on embedding similarity to the current query) are removed.
- Hybrid: A combination of the above, often prioritizing system prompts and current task details.
This process continues until the total token count fits within the budget. Failure to evict effectively results in a context_window_exceeded error from the LLM API, halting agent execution.
5. LLM Inference
The truncated and tokenized input is sent to the LLM. The LLM processes this information, generates a response, or decides to call a tool.
6. Output Processing and Memory Update
The agent receives the LLM's output. If it's a direct response, it's passed to the user. If it's a tool call, the tool is executed, and its output is prepared for the next turn's context. The current user input and the agent's response (or tool call/output) are then stored in the agent's long-term memory for future retrieval, ensuring state persistence beyond the immediate context window.
Architecture
The conceptual architecture for managing an agent's working memory, the context window, involves several interacting components. At its core is the Agent Orchestrator, which coordinates the entire agent workflow. This orchestrator initiates the process by receiving user input and delegating to a Context Manager. The Context Manager is responsible for assembling the prompt for the LLM Provider.
Data flows into the Context Manager from multiple sources. A Memory Store (e.g., a database or vector store) provides historical conversation turns and potentially long-term knowledge, which are retrieved and passed to the Context Manager. A Tool Registry exposes available tool definitions and their schemas, which the Context Manager incorporates into the prompt. Additionally, any Retriever components (e.g., for RAG) feed relevant documents into the Context Manager based on the current query.
The Context Manager aggregates these inputs, tokenizes them, applies eviction strategies if the Token Budget is exceeded, and then constructs the final prompt. This prompt, a single string of tokens, flows as input to the LLM Provider. The LLM Provider processes this input and returns an LLM Response (either a natural language output or a tool call instruction) back to the Agent Orchestrator. The Orchestrator then either presents the response to the user, executes the tool via the Tool Registry, or updates the Memory Store with the latest interaction, completing the cycle.
Context Window as a Finite Resource: The Scratchpad Analogy
Consider the context window as a fixed-size scratchpad or a CPU's register set. It's the only place the LLM can 'see' and 'think' about information for a given inference step. Unlike human memory, which can fluidly access vast amounts of information, an LLM's 'attention' is strictly limited to what's present in this window. This finite nature necessitates careful management, as every piece of information - system instructions, tool definitions, conversation history, retrieved documents, and the current query - competes for space. Exceeding this limit results in truncation, leading to a loss of information that can manifest as forgetting, hallucination, or an inability to complete complex tasks.
Token Budget Allocation Strategies
Effective context window management is largely about strategic token budget allocation. A typical allocation might reserve tokens for:
- System Prompt (High Priority): Often fixed, defining the agent's core identity and rules. Should be concise but comprehensive.
- Tool Descriptions (High Priority for Tool-Using Agents): The schemas and descriptions of callable functions. These are critical for the LLM to understand its capabilities. For agents with many tools, this can consume a significant portion of the budget.
- Conversation History (Dynamic Priority): Past user queries and agent responses. The most recent turns are generally more relevant. Strategies include keeping a fixed number of turns or summarizing older turns.
- Retrieved Documents (Dynamic Priority): Information fetched from external knowledge bases (e.g., RAG). The number and length of these documents must be carefully controlled to avoid overwhelming the context.
- Current User Input (Highest Priority): The immediate query or instruction from the user. This is always included.
Engineers must define a clear hierarchy for these components. For instance, a critical system prompt should never be truncated, while older conversation turns might be the first candidates for eviction.
Dynamic Context Eviction Algorithms
When the token budget is approached, an eviction algorithm determines what information to discard. Simple methods include:
- First-In, First-Out (FIFO): The oldest messages in the conversation history are removed first. This is straightforward but can remove context that becomes relevant again later.
- Least Recently Used (LRU): Similar to cache eviction, items not accessed recently are removed. This is harder to implement for LLM context as 'access' is conceptual.
- Summarization: Instead of removing, older parts of the conversation or less critical documents are summarized by a smaller LLM call or a heuristic. This preserves the gist but loses detail.
- Relevance-based Eviction: Using embedding similarity, information least relevant to the current query or the overall system prompt is removed. This requires an additional embedding model call and can add latency.
- Hybrid Approaches: Combining these, e.g., always keeping the system prompt and last N turns, then summarizing or evicting based on relevance for the remaining budget.
Choosing an eviction strategy involves tradeoffs between computational cost, latency, and the risk of losing critical context. For mission-critical agents, summarization or relevance-based methods, despite their overhead, often yield better performance than simple truncation.
Impact of Tokenization on Context Length
Tokenization is not a simple character count. Different LLMs use different tokenizers (e.g., BPE, SentencePiece), and the same text can yield varying token counts. Special tokens, like those for tool calls or system messages, also consume budget. This non-uniformity means that a fixed character limit is an unreliable proxy for token budget. Developers must use the specific tokenizer provided by their LLM vendor to accurately measure context usage. Furthermore, some models have specific token limits for different sections of the prompt (e.g., user vs. system messages), requiring even more granular management.
Managing Tool Use within Context Constraints
Tool use significantly impacts context window management. Each tool's description, including its name, parameters, and a brief explanation, consumes tokens. For agents with a large number of tools, these descriptions alone can exhaust a substantial portion of the context. Strategies include:
- Dynamic Tool Loading: Only injecting tool descriptions into the context window when they are potentially relevant to the current task or user query.
- Tool Summarization: Providing concise, high-level descriptions for less frequently used tools, with detailed schemas only loaded when the LLM indicates interest.
- Tool Orchestration: Using a separate, smaller LLM or a heuristic to select a subset of tools to present to the main agent LLM, reducing the overall token footprint for tool descriptions.
- Output Pruning: Tool outputs can be verbose. Only including the most critical parts of a tool's output in the subsequent context, or summarizing it, can save tokens. This requires careful parsing and understanding of tool output structure. Failure to manage tool context can lead to the LLM 'forgetting' about available tools or being unable to interpret tool results due to truncation.
Code Example
import os
import tiktoken
import logging
from typing import List, Dict, Any
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class AgentContextManager:
def __init__(self, max_tokens: int, model_name: str = "gpt-4"):
self.max_tokens = max_tokens
self.tokenizer = tiktoken.encoding_for_model(model_name)
self.system_prompt_template = os.environ.get("AGENT_SYSTEM_PROMPT", "You are a helpful AI assistant.")
self.tool_definitions = self._load_tool_definitions()
self.conversation_history: List[Dict[str, str]] = []
def _load_tool_definitions(self) -> List[Dict[str, Any]]:
# In a real system, these would be loaded from a registry or config.
# For demonstration, hardcode a simple tool.
return [
{
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
},
"required": ["location"]
}
}
]
def _count_tokens(self, text: str) -> int:
return len(self.tokenizer.encode(text))
def _format_message(self, role: str, content: str) -> Dict[str, str]:
return {"role": role, "content": content}
def _assemble_context(self, user_input: str) -> List[Dict[str, str]]:
context_messages = []
# 1. System Prompt (highest priority)
context_messages.append(self._format_message("system", self.system_prompt_template))
# 2. Tool Definitions (high priority for tool-using agents)
# In a real scenario, this might be a JSON string of tool definitions.
# For simplicity, we'll represent it as a single string for token counting.
tool_str = "Available tools: " + str(self.tool_definitions)
context_messages.append(self._format_message("system", tool_str))
# 3. Conversation History (dynamic priority - oldest first for eviction)
# Add history in chronological order
context_messages.extend(self.conversation_history)
# 4. Current User Input (highest priority for the current turn)
context_messages.append(self._format_message("user", user_input))
return context_messages
def get_trimmed_context(self, user_input: str) -> List[Dict[str, str]]:
full_context = self._assemble_context(user_input)
current_token_count = sum(self._count_tokens(msg['content']) for msg in full_context)
logging.info(f"Initial context token count: {current_token_count}")
# Simple FIFO eviction strategy for conversation history
# We assume system prompt and tool definitions are critical and not evicted.
# This means we only evict from self.conversation_history.
# Calculate tokens for fixed parts (system prompt + tools + current input)
fixed_parts_tokens = self._count_tokens(self.system_prompt_template) + \
self._count_tokens("Available tools: " + str(self.tool_definitions)) + \
self._count_tokens(user_input)
if current_token_count > self.max_tokens:
logging.warning(f"Context window exceeded. Trimming from {current_token_count} tokens.")
# Tokens available for history after fixed parts are accounted for
available_for_history = self.max_tokens - fixed_parts_tokens
if available_for_history < 0:
logging.error("Fixed parts alone exceed max_tokens. Agent cannot function.")
# In a real system, this would raise an error or trigger a fallback.
return []
trimmed_history = []
history_tokens = 0
# Iterate history from newest to oldest to prioritize recent turns
for msg in reversed(self.conversation_history):
msg_tokens = self._count_tokens(msg['content'])
if history_tokens + msg_tokens <= available_for_history:
trimmed_history.insert(0, msg) # Add to front to maintain chronological order
history_tokens += msg_tokens
else:
logging.warning(f"Evicting old message: {msg['content'][:50]}...")
break # Stop adding if budget exceeded
# Re-assemble context with trimmed history
final_context = [
self._format_message("system", self.system_prompt_template),
self._format_message("system", tool_str) # Re-add tool string
]
final_context.extend(trimmed_history)
final_context.append(self._format_message("user", user_input))
final_token_count = sum(self._count_tokens(msg['content']) for msg in final_context)
logging.info(f"Final context token count after trimming: {final_token_count}")
return final_context
return full_context
def add_to_history(self, user_message: str, agent_response: str):
self.conversation_history.append(self._format_message("user", user_message))
self.conversation_history.append(self._format_message("assistant", agent_response))
# --- Simulation ---
if __name__ == "__main__":
# Set a dummy API key for tiktoken, though it's not strictly needed for encoding
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["AGENT_SYSTEM_PROMPT"] = "You are a helpful AI assistant that can answer questions and use tools."
# Simulate a small context window for demonstration purposes
MAX_CONTEXT_TOKENS = 100
manager = AgentContextManager(max_tokens=MAX_CONTEXT_TOKENS)
print(f"
--- Turn 1 ---")
user_query_1 = "Hello, how are you?"
context_1 = manager.get_trimmed_context(user_query_1)
print(f"Context for LLM (Turn 1): {context_1}")
manager.add_to_history(user_query_1, "I'm doing well, thank you!")
print(f"
--- Turn 2 ---")
user_query_2 = "Can you tell me the weather in New York?"
context_2 = manager.get_trimmed_context(user_query_2)
print(f"Context for LLM (Turn 2): {context_2}")
manager.add_to_history(user_query_2, "Sure, I can use the weather tool for that.")
print(f"
--- Turn 3 (Longer conversation, forcing eviction) ---")
# Add more history to force eviction
for i in range(5):
manager.add_to_history(f"User message {i}", f"Agent response {i} which is quite long and verbose to consume tokens quickly.")
user_query_3 = "What was the first thing I asked you? This is a very important question."
context_3 = manager.get_trimmed_context(user_query_3)
print(f"Context for LLM (Turn 3): {context_3}")
print(f"
--- Turn 4 (Even more history, demonstrating more aggressive eviction) ---")
for i in range(5, 10):
manager.add_to_history(f"User message {i}", f"Agent response {i} which is even longer and more verbose to consume tokens quickly.")
user_query_4 = "Summarize our conversation so far."
context_4 = manager.get_trimmed_context(user_query_4)
print(f"Context for LLM (Turn 4): {context_4}")
The output will show the constructed context for each turn. For Turn 1 and 2, the context will likely fit. For Turn 3 and 4, the `logging.warning` messages will indicate that the context window was exceeded, and older conversation history messages will be evicted (truncated) to fit within the `MAX_CONTEXT_TOKENS` limit. The final context for the LLM will contain the system prompt, tool definitions, the most recent conversation history, and the current user query, with older history removed.