AI Agent Memory Systems: Working, Episodic, Procedural
Source: mortalapps.com- AI agent memory systems categorize and manage information across three tiers: working, episodic, and procedural memory.
- This architecture addresses LLM context window limitations, enables long-term learning, and facilitates complex task execution by providing relevant information at different temporal scales.
- In production, effective memory management is critical for maintaining agent coherence, reducing hallucinations, optimizing operational costs, and ensuring reliable performance.
- Engineers use working memory for immediate task context, episodic memory (vector databases) for historical events and knowledge, and procedural memory (tool libraries) for capabilities and actions.
- Choosing the correct memory type and retrieval strategy is a key decision point for agent design, directly impacting accuracy, latency, and resource consumption.
Why This Matters
Effective AI agent memory systems are fundamental to building intelligent, persistent, and capable agents that can operate beyond single-turn interactions. The core problem these systems solve is the inherent limitation of Large Language Models (LLMs) regarding context window size and their stateless nature. While LLMs excel at processing immediate input, they lack intrinsic mechanisms for recalling past interactions, learning from experiences, or executing complex, multi-step procedures without explicit guidance. This is where a structured approach to AI agent memory types becomes indispensable.
This multi-tiered memory approach was invented to mimic human cognitive functions, allowing agents to maintain conversational coherence, adapt to new information over time, and leverage external tools. Without a robust memory architecture, agents in production environments suffer from severe limitations: they hallucinate due to lack of relevant context, fail to learn from past mistakes, cannot execute multi-step plans reliably, and incur high operational costs by re-generating information. This leads to poor user experiences, unreliable automation, and significant engineering overhead in debugging and maintenance.
Engineers must understand these memory systems to design agents that are not only functional but also efficient and scalable. Ignoring this leads to agents that are brittle, expensive, and unable to handle real-world complexity. When to use this vs. a simpler, stateless approach? For any agent requiring persistence, learning, tool use, or multi-turn interaction beyond a trivial scope, a well-designed memory system is non-negotiable. For simple, single-shot prompt-response systems, the overhead may not be justified, but for true agentic behavior, it's foundational.
Core Concepts
AI agent memory systems are typically structured into distinct tiers, each serving a specific purpose in managing information for the agent's operation.
- Working Memory: This is the agent's short-term, immediate context. It resides within the LLM's active context window and holds information directly relevant to the current turn or sub-task. It's analogous to a human's short-term memory or a scratchpad, containing recent conversation turns, intermediate thoughts, retrieved facts, and tool outputs. Its capacity is strictly limited by the LLM's token limit and is highly transient.
- Episodic Memory: This constitutes the agent's long-term, experience-based memory. It stores past interactions, observations, learned facts, and outcomes of previous tasks. Typically implemented using vector databases, episodic memory allows for semantic retrieval of relevant historical data based on the current query or context. It's non-volatile and can grow indefinitely, providing a persistent knowledge base for the agent.
- Procedural Memory: This refers to the agent's knowledge of *how* to do things. It encompasses the agent's skill set, including available tools, functions, APIs, and internal capabilities. Procedural memory is typically represented as a library of callable functions or services with defined schemas, allowing the agent to select and execute appropriate actions to achieve its goals. It's distinct from factual knowledge, focusing on operational capabilities.
- Context Window: The fixed-size input buffer of an LLM, measured in tokens. All information passed to the LLM for a single inference call, including system prompts, user input, working memory, and retrieved episodic memory, must fit within this window. Managing this limit is central to agent design.
- Retrieval-Augmented Generation (RAG): A technique where an LLM retrieves relevant information from an external knowledge base (like episodic memory) before generating a response. This augments the LLM's internal knowledge with external, up-to-date, or specific facts, reducing hallucinations and improving factual accuracy.
- Memory Consolidation: The process of summarizing, abstracting, or transforming raw episodic memories into more compact, organized, or higher-level representations. This can involve generating summaries of conversations, extracting key facts, or forming relationships in a knowledge graph, making future retrieval more efficient and effective.
- Eviction Policies: Algorithms or strategies used to manage memory capacity, particularly for working memory or when consolidating episodic memory. These determine which pieces of information are discarded or summarized when memory limits are reached, often based on recency, frequency, or perceived importance.
How It Works
An AI agent's memory system operates as a dynamic, interconnected hierarchy, orchestrating information flow to support complex reasoning and action. The process typically unfolds in several phases:
1. Initial Input and Working Memory Initialization
Upon receiving a user query or an external event, the agent's orchestrator initializes its working memory. This includes the current user input, the system prompt defining the agent's persona and goals, and any immediate context from the previous turn. This initial working memory forms the prompt for the first LLM invocation.
2. LLM Reasoning and Initial Action Proposal
The orchestrator sends the working memory to the LLM. The LLM processes this information to understand the current task, generate internal thoughts, and propose an initial action. This action might be to directly answer the user, query episodic memory for more information, or invoke a tool from procedural memory.
3. Episodic Memory Retrieval (RAG)
If the LLM determines that more context is needed (e.g., to answer a question, understand user history, or inform a decision), the orchestrator triggers a retrieval from episodic memory. The current query or a synthesized query from working memory is embedded and used to perform a similarity search against the vector database. Relevant chunks of historical data, past interactions, or external knowledge are retrieved. These retrieved documents are then added to the working memory, potentially displacing older, less relevant information based on an eviction policy.
*Failure Case*: If retrieval yields irrelevant or insufficient information, the LLM may hallucinate or indicate it lacks knowledge. Robust systems implement re-ranking, query expansion, or even self-correction loops (e.g., CRAG, Self-RAG) to refine retrieval or flag low-confidence results.
4. Procedural Memory Invocation (Tool Use)
If the LLM's reasoning indicates that an external action is required (e.g., fetching data from an API, sending an email, performing a calculation), the orchestrator identifies and invokes a tool from procedural memory. The LLM generates the necessary parameters based on the working memory. The tool executes, and its output (success or failure) is then added back to the working memory.
*Failure Case*: Tool invocation can fail due to incorrect parameters, network issues, or external service errors. The orchestrator must handle these by logging errors, retrying, or instructing the LLM to re-plan or inform the user of the failure. The tool's error message becomes part of the working memory for the LLM to process.
5. Iterative Refinement and Response Generation
The LLM continues to iterate, using updated working memory (incorporating retrieved episodic context and tool outputs) to refine its understanding, generate further thoughts, and propose subsequent actions. This loop continues until the agent determines it has sufficient information to formulate a final response or complete its task. The final response is then generated and sent to the user.
6. Memory Update and Consolidation
After a task or conversation turn is completed, relevant parts of the working memory, the agent's final response, and any significant observations are processed for storage in episodic memory. This often involves summarizing the interaction, extracting key entities, or creating new embeddings to represent the experience. This consolidation ensures that the agent learns from its interactions and builds a persistent knowledge base for future use. Less critical working memory elements are discarded.
Architecture
The conceptual architecture for an AI agent memory system involves a central Agent Orchestrator interacting with distinct memory components and external services.
At the core is the Agent Orchestrator, which acts as the control plane. It receives user inputs, manages the LLM's inference calls, and directs the flow of information between memory types. The orchestrator is responsible for constructing prompts, parsing LLM outputs, and executing decisions.
Connected to the Orchestrator are three primary memory components:
- Working Memory Module: This component manages the active context window for the LLM. It's typically an in-memory data structure (e.g., a list of messages, a dictionary) that stores recent conversation turns, intermediate thoughts, and temporary data. The Orchestrator pushes and pulls data from here to construct LLM prompts. Data flows into the LLM and out as generated thoughts or actions.
- Episodic Memory Store: This is a persistent storage layer, often implemented as a Vector Database (e.g., Pinecone, Weaviate, Chroma) coupled with a Metadata Store (e.g., PostgreSQL, Redis) for additional context and filtering. The Orchestrator sends queries (embeddings) to the Vector Database, which returns semantically similar document chunks. These chunks flow back to the Orchestrator and are injected into Working Memory. New experiences or consolidated summaries flow from the Orchestrator to the Episodic Memory Store for long-term storage.
- Procedural Memory Registry: This component is a Tool Registry or Function Library that defines the agent's capabilities. It contains metadata (schemas, descriptions) for each tool. The Orchestrator queries this registry to understand available tools and, based on LLM output, invokes specific Tool Execution Services. Tool calls (parameters) flow from the Orchestrator to the Tool Execution Services, and results (outputs, errors) flow back to the Orchestrator, which then updates Working Memory.
External systems, such as user interfaces, third-party APIs, and data sources, interact with the Agent Orchestrator (for input/output) and the Tool Execution Services (for specific actions). Execution starts with a user input to the Orchestrator and ends with a response back to the user, with continuous loops of memory access and processing in between.
Working Memory Management Strategies
Working memory is the most volatile and constrained memory type, directly impacting an agent's immediate coherence and cost. Effective management is paramount. The primary challenge is the LLM's fixed context window. Strategies include:
- Token Budgeting: Strictly allocating tokens for different prompt components (system instructions, retrieved context, conversation history, scratchpad). For example, reserving 20% for system prompt, 30% for RAG results, and 50% for conversation history, dynamically adjusting as needed. This prevents critical instructions from being truncated.
- Sliding Window: For conversational agents, maintaining a fixed-size window of recent turns. When a new turn exceeds the limit, the oldest turns are dropped. This is simple but can lose crucial context from earlier in the conversation.
- Summarization: Periodically summarizing older conversation turns or intermediate thoughts and replacing the raw history with the summary. This compresses information, retaining key points while reducing token count. The summarization itself consumes tokens and can introduce information loss or bias. Recursive summarization can be used for very long conversations.
- Importance-Based Eviction: Assigning a score to each piece of information in working memory based on its relevance to the current goal or recent interactions. When the context window is full, the least important items are evicted. This requires an LLM call or a heuristic to determine importance, adding latency and cost.
- Tree-of-Thought or Graph-of-Thought: Structuring the agent's internal reasoning as a tree or graph, where only relevant branches or nodes are brought into the working memory for the current LLM call. This allows for complex reasoning paths without overwhelming the context window.
Deciding which strategy to use depends on the agent's task, required conversational depth, and tolerance for latency/cost. For simple chatbots, a sliding window might suffice. For complex problem-solving agents, summarization or importance-based eviction combined with graph structures is often necessary.
Episodic Memory Retrieval and Storage
Episodic memory provides the agent's long-term knowledge. Its effectiveness hinges on robust retrieval and storage mechanisms.
- Vector Embedding Models: Information (documents, conversation turns, observations) is converted into high-dimensional numerical vectors (embeddings) using models like OpenAI's
text-embedding-ada-002or various open-source models (e.g., Sentence-BERT). The quality of these embeddings directly impacts retrieval accuracy. - Similarity Search: When the agent needs to recall information, a query (e.g., the current user input, an agent's thought) is also embedded. This query embedding is then used to find the most semantically similar embeddings in the vector database. Common similarity metrics include cosine similarity, dot product, and Euclidean distance. Indexing structures like HNSW (Hierarchical Navigable Small World) or IVF (Inverted File Index) are used for efficient nearest neighbor search in large datasets.
- Hybrid Retrieval: Combining vector search with keyword-based methods (e.g., BM25) can improve recall, especially for specific entity names or rare terms that might not be well-represented by embeddings alone. The results from both methods are then fused or re-ranked.
- Memory Consolidation: To prevent episodic memory from becoming a flat, unmanageable list of events, consolidation is crucial. This can involve:
- Summarization: Periodically summarizing clusters of related events or conversations.
- Fact Extraction: Identifying and storing discrete facts from interactions, potentially in a structured format.
- Knowledge Graph Construction: Representing entities and their relationships, allowing for multi-hop reasoning and more nuanced retrieval. This moves beyond simple semantic similarity to inferential connections.
- Eviction Policies for Episodic Memory: While episodic memory is generally persistent, some systems might implement eviction for less important or outdated information, especially in resource-constrained environments or for privacy reasons. Policies could be based on age (TTL), usage frequency, or explicit importance scores.
Procedural Memory: Tool Orchestration
Procedural memory enables agents to interact with the external world. Its implementation involves defining tools, selecting them, and executing them reliably.
- Tool Definition: Tools are typically defined with clear schemas (e.g., OpenAPI specifications, Pydantic models) that describe their purpose, required parameters, and expected outputs. This allows the LLM to understand how to use them and for the orchestrator to validate calls.
- Tool Selection Logic: The agent (usually the LLM) determines which tool to use based on its current goal and the available context in working memory. This often involves the LLM generating a function call in a structured format (e.g., JSON) that matches a tool's schema. For complex scenarios, semantic search over tool descriptions can augment LLM-based selection.
- Tool Execution: The orchestrator intercepts the LLM's tool call, validates parameters against the schema, and then invokes the corresponding external service or function. This execution should be robust, including error handling, retries, and timeouts.
- Output Integration: The output of the tool execution (success message, data, or error) is captured and injected back into the agent's working memory. This allows the LLM to incorporate the results into its ongoing reasoning or to handle failures.
Inter-Memory Communication and Orchestration
The true power of a multi-tiered memory system lies in the intelligent orchestration of these components. The Agent Orchestrator acts as the conductor.
- Decision Point: The LLM, guided by its system prompt and current working memory, decides whether to:
- Generate a direct response.
- Query episodic memory (e.g., "I need more information about X").
- Invoke a tool from procedural memory (e.g., "I need to perform action Y").
- Refine its thoughts and continue internal reasoning.
- Feedback Loops: Outcomes from one memory type feed into another. For example, a successful tool execution (procedural memory) updates the working memory, and the entire interaction (working memory) is later consolidated into episodic memory. Retrieved information from episodic memory directly augments working memory for subsequent LLM calls.
- State Management: Frameworks like LangGraph explicitly model the agent's state transitions, allowing for clear definitions of when each memory type is accessed and how the agent progresses through its task. This provides explicit control over memory interactions, crucial for complex, multi-step workflows.
Tradeoffs: Latency, Cost, and Accuracy
Each memory type introduces specific tradeoffs:
- Working Memory: High accuracy (direct context), low latency (in-context), but extremely high cost (per-token LLM charges) and limited capacity. Over-reliance leads to context window exhaustion and high API bills.
- Episodic Memory: High capacity, potentially high accuracy (if retrieval is good), but introduces retrieval latency (vector DB lookup) and storage/embedding costs. Poor retrieval leads to low accuracy and hallucinations.
- Procedural Memory: Enables powerful actions and external data access, but introduces tool execution latency, potential for external service failures, and development/maintenance costs for tools. Incorrect tool selection or execution can lead to critical errors.
Optimizing an agent's memory system involves balancing these factors. For instance, aggressive summarization reduces working memory cost but can decrease accuracy. Extensive episodic memory retrieval improves factual grounding but increases latency. A well-designed system intelligently queries the minimal necessary memory for the current task.
Code Example
import os
import time
import logging
from typing import List, Dict, Any, Optional
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class WorkingMemory:
def __init__(self, max_tokens: int = 4096):
self.context: List[Dict[str, str]] = []
self.max_tokens = max_tokens
def add_message(self, role: str, content: str):
self.context.append({"role": role, "content": content})
self.trim_context() # Ensure context stays within limits
def get_context(self) -> List[Dict[str, str]]:
return self.context
def trim_context(self):
# A simplistic token trimming: remove oldest messages if over limit
# In a real system, this would involve actual token counting and summarization
current_length = sum(len(msg['content'].split()) for msg in self.context)
while current_length > self.max_tokens * 0.75 and len(self.context) > 1:
logging.info(f"Trimming working memory: current length {current_length} > {self.max_tokens * 0.75}")
self.context.pop(0) # Remove oldest message
current_length = sum(len(msg['content'].split()) for msg in self.context)
class EpisodicMemory:
def __init__(self):
self.memories: List[str] = []
def store_memory(self, memory_item: str):
logging.info(f"Storing episodic memory: {memory_item[:50]}...")
self.memories.append(memory_item)
def retrieve_memory(self, query: str, top_k: int = 1) -> List[str]:
# In a real system, this would be a vector similarity search
logging.info(f"Retrieving from episodic memory for query: {query[:50]}...")
relevant_memories = [m for m in self.memories if query.lower() in m.lower()]
return relevant_memories[:top_k]
class ProceduralMemory:
def __init__(self):
self.tools = {
"get_current_time": self._get_current_time,
"search_knowledge_base": self._search_knowledge_base
}
def _get_current_time(self, **kwargs) -> str:
logging.info("Executing tool: get_current_time")
return f"The current time is {time.strftime('%H:%M:%S')}"
def _search_knowledge_base(self, query: str, **kwargs) -> str:
logging.info(f"Executing tool: search_knowledge_base with query '{query}'")
# Simulate a knowledge base search
if "weather" in query.lower():
return "The weather forecast is sunny with a high of 25C."
return "No specific information found for your query."
def execute_tool(self, tool_name: str, **kwargs) -> str:
if tool_name in self.tools:
try:
return self.tools[tool_name](**kwargs)
except Exception as e:
logging.error(f"Tool '{tool_name}' execution failed: {e}")
return f"Error executing tool '{tool_name}': {e}"
return f"Tool '{tool_name}' not found."
class AIAgent:
def __init__(self, llm_api_key: str):
self.working_memory = WorkingMemory()
self.episodic_memory = EpisodicMemory()
self.procedural_memory = ProceduralMemory()
self.llm_api_key = llm_api_key
self.system_prompt = "You are a helpful AI assistant. Use tools when necessary."
def _call_llm(self, prompt_messages: List[Dict[str, str]]) -> str:
# Simulate an LLM call. In production, this would use an actual LLM API.
# For demonstration, we'll use simple keyword matching to simulate LLM logic.
logging.info("Simulating LLM call...")
full_prompt = "
".join([f"{msg['role']}: {msg['content']}" for msg in prompt_messages])
if "current time" in full_prompt.lower():
return "TOOL_CALL: get_current_time"
elif "weather forecast" in full_prompt.lower():
return "TOOL_CALL: search_knowledge_base(query='weather forecast')"
elif "past interactions" in full_prompt.lower() or "remember" in full_prompt.lower():
return "EPISODIC_RETRIEVAL: What did we discuss before?"
elif "hello" in full_prompt.lower():
return "Hello there! How can I assist you today?"
else:
return "I'm not sure how to respond to that, but I'm learning."
def run(self, user_input: str) -> str:
self.working_memory.add_message("system", self.system_prompt)
self.working_memory.add_message("user", user_input)
# Step 1: LLM processes current working memory
llm_response = self._call_llm(self.working_memory.get_context())
# Step 2: Agent decides action based on LLM response
if llm_response.startswith("TOOL_CALL:"):
tool_call_str = llm_response.replace("TOOL_CALL: ", "")
tool_name_start = tool_call_str.find("(")
if tool_name_start != -1:
tool_name = tool_call_str[:tool_name_start]
params_str = tool_call_str[tool_name_start+1:-1]
# Basic parsing for demonstration; real parsing would be more robust
params = {}
if params_str: # Check if params_str is not empty
try:
for p in params_str.split(','):
key, val = p.split('=')
params[key.strip()] = val.strip().strip("'\") # Remove quotes
except ValueError:
logging.warning(f"Could not parse tool parameters: {params_str}")
tool_output = self.procedural_memory.execute_tool(tool_name, **params)
else:
tool_name = tool_call_str
tool_output = self.procedural_memory.execute_tool(tool_name)
self.working_memory.add_message("tool_output", tool_output)
# Re-call LLM with tool output in context
llm_final_response = self._call_llm(self.working_memory.get_context())
final_agent_response = llm_final_response # Simplified: LLM directly gives final response after tool
elif llm_response.startswith("EPISODIC_RETRIEVAL:"):
query = llm_response.replace("EPISODIC_RETRIEVAL: ", "")
retrieved_data = self.episodic_memory.retrieve_memory(query)
if retrieved_data:
self.working_memory.add_message("retrieved_memory", f"Retrieved: {'; '.join(retrieved_data)}")
else:
self.working_memory.add_message("retrieved_memory", "No relevant past memories found.")
# Re-call LLM with retrieved data in context
llm_final_response = self._call_llm(self.working_memory.get_context())
final_agent_response = llm_final_response # Simplified
else:
final_agent_response = llm_response
# Step 3: Consolidate relevant parts of interaction into episodic memory
self.episodic_memory.store_memory(f"User: {user_input}, Agent: {final_agent_response}")
return final_agent_response
# --- Main execution ---
if __name__ == "__main__":
# Ensure LLM_API_KEY is set in environment variables
llm_key = os.environ.get("LLM_API_KEY", "mock_api_key")
if llm_key == "mock_api_key":
logging.warning("LLM_API_KEY not set. Using mock key. Real LLM calls will fail.")
agent = AIAgent(llm_key)
print("
--- Interaction 1: Simple greeting ---")
response1 = agent.run("Hello!")
print(f"Agent: {response1}")
print("
--- Interaction 2: Tool use (time) ---")
response2 = agent.run("What is the current time?")
print(f"Agent: {response2}")
print("
--- Interaction 3: Tool use (weather) ---")
response3 = agent.run("What's the weather forecast?")
print(f"Agent: {response3}")
print("
--- Interaction 4: Episodic memory retrieval ---")
# Store some initial memories directly for retrieval test
agent.episodic_memory.store_memory("We discussed project Alpha's deadline last week.")
agent.episodic_memory.store_memory("The user asked about the Q3 budget report.")
response4 = agent.run("Can you remember what we discussed about project Alpha?")
print(f"Agent: {response4}")
print("
--- Interaction 5: Context trimming simulation ---")
# Fill working memory to trigger trimming
for i in range(10):
agent.working_memory.add_message("user", f"This is a long message to fill context {i}. " * 50)
response5 = agent.run("What was the very first thing I asked you?")
print(f"Agent: {response5}")
print(f"Working memory size after trimming: {len(agent.working_memory.get_context())} messages")
--- Interaction 1: Simple greeting --- Agent: Hello there! How can I assist you today? --- Interaction 2: Tool use (time) --- Agent: The current time is [current_time_hh:mm:ss] --- Interaction 3: Tool use (weather) --- Agent: The weather forecast is sunny with a high of 25C. --- Interaction 4: Episodic memory retrieval --- Agent: Retrieved: We discussed project Alpha's deadline last week.; User: Can you remember what we discussed about project Alpha?, Agent: Retrieved: We discussed project Alpha's deadline last week. --- Interaction 5: Context trimming simulation --- Agent: I'm not sure how to respond to that, but I'm learning. Working memory size after trimming: [number of messages, likely less than initial 10+]