Introduction to Agent Memory Architectures
Memory is one of the most critical components of any intelligent agent system. Without memory, an AI agent is essentially stateless — it processes each request in isolation, unable to recall previous interactions, learn from past mistakes, or maintain context across conversations. Pydantic AI, a framework built on top of Pydantic for building production-grade AI agents, provides robust primitives for implementing memory architectures that range from simple conversation buffers to sophisticated retrieval-augmented memory stores.
In this tutorial, we will explore how to design and implement agent memory architectures using Pydantic AI. We will cover the conceptual foundations, walk through practical implementations, and discuss best practices that will help you build agents capable of maintaining meaningful, long-term relationships with users.
Why Memory Matters in AI Agents
Consider a customer support agent that helps users troubleshoot software issues. Without memory, every time a user returns with a follow-up question, the agent starts from scratch — asking for account details, re-establishing the context of the problem, and re-running diagnostics. This is frustrating for users and wasteful in terms of token consumption. With a well-designed memory architecture, the agent can:
- Recall previous conversations and user preferences
- Maintain context across multi-turn interactions within a single session
- Learn facts about users that persist between sessions
- Avoid repeating questions or re-doing work already completed
- Provide personalized responses based on historical interactions
- Reduce latency and token costs by caching relevant information
Memory transforms an agent from a stateless function into a persistent, context-aware assistant. The challenge lies in deciding what to remember, how long to remember it, and how to retrieve relevant memories efficiently when needed.
Types of Memory Architectures
Before diving into code, it is important to understand the different categories of memory that agents typically employ. Drawing inspiration from cognitive science, most modern agent frameworks distinguish between several types of memory:
Short-Term Memory (Working Memory)
Short-term memory holds the immediate context of the current conversation. This includes the message history, intermediate tool call results, and any temporary state needed to complete the current task. Short-term memory is typically bounded — either by a maximum number of messages or by a token limit — and is discarded when the conversation ends.
Long-Term Memory
Long-term memory persists across sessions and conversations. It can store user preferences, factual knowledge, learned patterns, and historical interaction summaries. Long-term memory is usually backed by a database or vector store and requires explicit retrieval mechanisms to surface relevant information.
Episodic Memory
Episodic memory records specific past events or interactions — essentially a log of what happened, when, and in what context. This is useful for agents that need to recall specific past conversations or actions, such as a personal assistant that remembers what you discussed last week.
Semantic Memory
Semantic memory stores general knowledge and facts, independent of when they were learned. For an agent, this might include a knowledge base of product information, user profiles, or domain-specific facts that inform decision-making.
Procedural Memory
Procedural memory represents learned skills and procedures — how to do things. In the context of AI agents, this often manifests as learned tool-use patterns, optimized prompts, or cached plans for recurring task types.
Getting Started with Pydantic AI
Pydantic AI is a framework that brings type safety and structured data validation to AI agent development. It integrates naturally with the Pydantic ecosystem and supports multiple LLM providers. Let us start by installing the necessary dependencies.
pip install pydantic-ai pydantic-ai-slim openai
For memory backends that involve vector storage, you will also want to install additional packages:
pip install chromadb sentence-transformers
Let us begin with a basic Pydantic AI agent to establish our foundation:
from pydantic_ai import Agent
from pydantic import BaseModel, Field
class Response(BaseModel):
answer: str = Field(description="The agent's response")
confidence: float = Field(description="Confidence score 0-1")
agent = Agent(
model="openai:gpt-4o",
result_type=Response,
system_prompt="You are a helpful assistant."
)
result = agent.run_sync("What is the capital of France?")
print(result.data)
This basic agent has no memory. Each call to agent.run_sync() is independent. Let us now build memory into it step by step.
Implementing Short-Term Memory with Message History
The simplest form of memory is maintaining conversation history across multiple turns within a single session. Pydantic AI supports this through its dependency injection system and message history management.
Basic Conversation History
Pydantic AI allows you to pass previous messages back into subsequent calls. The RunContext and message history APIs make this straightforward:
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
agent = Agent(
model="openai:gpt-4o",
system_prompt="You are a helpful conversational assistant."
)
# First turn
result1 = agent.run_sync("My name is Alice and I live in Berlin.")
print(f"Turn 1: {result1.data}")
# Second turn — pass the message history from the first turn
result2 = agent.run_sync(
"What is my name and where do I live?",
message_history=result1.all_messages()
)
print(f"Turn 2: {result2.data}")
The key here is result.all_messages(), which returns the full list of messages from the previous run. By passing this as message_history to the next run, the agent retains the context of the conversation.
Managing Conversation Sessions
In a real application, you will want to manage multiple conversation sessions, each with its own history. A common pattern is to use a session manager that stores message histories keyed by a session identifier:
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
from typing import Dict, List
import uuid
class SessionManager:
def __init__(self):
self._sessions: Dict[str, List[ModelMessage]] = {}
def create_session(self) -> str:
session_id = str(uuid.uuid4())
self._sessions[session_id] = []
return session_id
def get_history(self, session_id: str) -> List[ModelMessage]:
return self._sessions.get(session_id, [])
def update_history(self, session_id: str, messages: List[ModelMessage]):
self._sessions[session_id] = messages
def clear_session(self, session_id: str):
self._sessions.pop(session_id, None)
agent = Agent(
model="openai:gpt-4o",
system_prompt="You are a helpful assistant with memory of our conversation."
)
session_manager = SessionManager()
session_id = session_manager.create_session()
def chat(user_input: str, session_id: str) -> str:
history = session_manager.get_history(session_id)
result = agent.run_sync(user_input, message_history=history)
session_manager.update_history(session_id, result.all_messages())
return result.data
# Simulate a multi-turn conversation
print(chat("Hi, I'm Bob and I love hiking.", session_id))
print(chat("What are my hobbies?", session_id))
print(chat("Suggest a weekend activity for me.", session_id))
This pattern works well for single-server applications. For distributed deployments, you would replace the in-memory dictionary with a persistent store like Redis or a database.
Bounded Conversation History
As conversations grow longer, the message history can exceed the model's context window or become prohibitively expensive. A common strategy is to bound the history to the most recent N messages or to implement a sliding window with summarization:
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
class BoundedHistory:
def __init__(self, max_messages: int = 20):
self.max_messages = max_messages
self._messages: List[ModelMessage] = []
def add_messages(self, messages: List[ModelMessage]):
self._messages.extend(messages)
# Keep only the most recent messages
if len(self._messages) > self.max_messages:
self._messages = self._messages[-self.max_messages:]
def get_messages(self) -> List[ModelMessage]:
return list(self._messages)
agent = Agent(model="openai:gpt-4o")
history = BoundedHistory(max_messages=10)
result = agent.run_sync("Tell me about Python.", message_history=history.get_messages())
history.add_messages(result.all_messages())
Implementing Long-Term Memory
Short-term memory handles within-session context, but what about information that should persist across sessions? This is where long-term memory comes in. The most common approach is to store memories in a vector database and retrieve relevant ones using semantic search at the start of each conversation turn.
Vector-Based Memory Store
Let us build a long-term memory system using ChromaDB as the vector store. The idea is to extract important facts from conversations, embed them, and store them for later retrieval:
import chromadb
from chromadb.utils import embedding_functions
from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import List, Optional
import json
import uuid
class Memory(BaseModel):
content: str = Field(description="The memory content")
user_id: str = Field(description="The user this memory belongs to")
metadata: dict = Field(default_factory=dict, description="Additional metadata")
class VectorMemoryStore:
def __init__(self, collection_name: str = "agent_memories"):
self.client = chromadb.PersistentClient(path="./chroma_db")
self.embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2"
)
self.collection = self.client.get_or_create_collection(
name=collection_name,
embedding_function=self.embed_fn
)
def store(self, memory: Memory):
memory_id = str(uuid.uuid4())
self.collection.add(
ids=[memory_id],
documents=[memory.content],
metadatas=[{
"user_id": memory.user_id,
**memory.metadata
}]
)
return memory_id
def retrieve(self, query: str, user_id: str, n_results: int = 5) -> List[str]:
results = self.collection.query(
query_texts=[query],
n_results=n_results,
where={"user_id": user_id}
)
return results["documents"][0] if results["documents"] else []
def delete_user_memories(self, user_id: str):
self.collection.delete(where={"user_id": user_id})
memory_store = VectorMemoryStore()
Memory Extraction Agent
Not every message in a conversation contains information worth remembering. We can use a secondary agent whose sole job is to extract salient facts from conversations and store them in the long-term memory:
from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import List
class ExtractedMemories(BaseModel):
memories: List[str] = Field(
description="A list of distinct, factual memories worth remembering long-term"
)
memory_extractor = Agent(
model="openai:gpt-4o",
result_type=ExtractedMemories,
system_prompt=(
"You are a memory extraction system. Analyze the conversation and extract "
"key facts worth remembering about the user. Focus on: personal preferences, "
"important life events, recurring topics, explicit requests to remember something, "
"and factual information shared by the user. Ignore small talk and transient details. "
"Each memory should be a self-contained, factual statement."
)
)
def extract_and_store_memories(conversation_text: str, user_id: str):
result = memory_extractor.run_sync(
f"Extract memories from this conversation:\n\n{conversation_text}"
)
for memory_text in result.data.memories:
memory_store.store(Memory(
content=memory_text,
user_id=user_id,
metadata={"source": "conversation"}
))
return result.data.memories
Integrating Long-Term Memory into the Main Agent
Now let us tie everything together. The main agent will retrieve relevant memories before responding and store new memories after each conversation turn:
from pydantic_ai import Agent, RunContext
from pydantic_ai.messages import ModelMessage
from typing import List
main_agent = Agent(
model="openai:gpt-4o",
system_prompt=(
"You are a helpful, personalized assistant. You have access to long-term "
"memories about the user. Use these memories to provide personalized responses. "
"If the user shares new important information, acknowledge it."
)
)
@main_agent.system_prompt
async def inject_memories(ctx: RunContext[str]) -> str:
"""Inject relevant long-term memories into the system prompt."""
user_id = ctx.deps
# We need a query — use the last user message if available
last_message = ""
if ctx.messages:
for msg in reversed(ctx.messages):
if hasattr(msg, "content") and isinstance(msg.content, str):
last_message = msg.content
break
if not last_message:
return "No long-term memories available yet."
memories = memory_store.retrieve(last_message, user_id, n_results=5)
if not memories:
return "No relevant long-term memories found."
memory_text = "\n".join(f"- {m}" for m in memories)
return f"Long-term memories about this user:\n{memory_text}"
def chat_with_memory(user_input: str, user_id: str, history: List[ModelMessage]) -> tuple:
"""Run a conversation turn with full memory integration."""
# Run the main agent
result = main_agent.run_sync(
user_input,
deps=user_id,
message_history=history
)
# Extract and store new memories from this turn
conversation_snippet = f"User: {user_input}\nAssistant: {result.data}"
new_memories = extract_and_store_memories(conversation_snippet, user_id)
if new_memories:
print(f" [Stored {len(new_memories)} new memories]")
return result.data, result.all_messages()
# Example usage across multiple sessions
user_id = "user_123"
print("=== Session 1 ===")
history = []
response, history = chat_with_memory(
"Hi! I'm Sarah, I work as a data scientist at a startup in Austin.",
user_id, history
)
print(f"Agent: {response}")
response, history = chat_with_memory(
"I'm really into rock climbing and I have a dog named Max.",
user_id, history
)
print(f"Agent: {response}")
# Simulate a new session — history is empty but long-term memory persists
print("\n=== Session 2 (new session) ===")
new_session_history = []
response, new_session_history = chat_with_memory(
"Can you suggest a weekend activity for me?",
user_id, new_session_history
)
print(f"Agent: {response}")
# The agent should reference rock climbing and/or Max from long-term memory
Building a Custom Memory Backend
While vector stores are excellent for semantic retrieval, some use cases call for structured memory that can be queried precisely. Let us build a hybrid memory system that combines structured storage (for facts with known schemas) with vector storage (for free-form memories).
Structured Fact Memory
from pydantic import BaseModel, Field
from typing import Dict, List, Optional
from datetime import datetime
import json
class UserFact(BaseModel):
user_id: str
fact_type: str = Field(description="Category: preference, profile, relationship, goal")
key: str = Field(description="The fact key, e.g. 'favorite_food'")
value: str = Field(description="The fact value, e.g. 'sushi'")
confidence: float = Field(default=1.0, description="Confidence in this fact")
updated_at: datetime = Field(default_factory=datetime.now)
class StructuredMemoryStore:
def __init__(self):
self._facts: Dict[str, Dict[str, UserFact]] = {}
def upsert(self, fact: UserFact):
if fact.user_id not in self._facts:
self._facts[fact.user_id] = {}
self._facts[fact.user_id][fact.key] = fact
def get(self, user_id: str, key: str) -> Optional[UserFact]:
return self._facts.get(user_id, {}).get(key)
def get_all(self, user_id: str) -> List[UserFact]:
return list(self._facts.get(user_id, {}).values())
def get_by_type(self, user_id: str, fact_type: str) -> List[UserFact]:
return [
f for f in self.get_all(user_id)
if f.fact_type == fact_type
]
def delete(self, user_id: str, key: str):
if user_id in self._facts:
self._facts[user_id].pop(key, None)
structured_store = StructuredMemoryStore()
Fact Extraction Agent
from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import List
class ExtractedFact(BaseModel):
fact_type: str = Field(description="Category: preference, profile, relationship, goal")
key: str = Field(description="A snake_case key for this fact")
value: str = Field(description="The fact value")
class ExtractedFacts(BaseModel):
facts: List[ExtractedFact] = Field(default_factory=list)
fact_extractor = Agent(
model="openai:gpt-4o",
result_type=ExtractedFacts,
system_prompt=(
"You extract structured facts from user messages. "
"Categorize each fact as: preference (likes/dislikes), "
"profile (personal info like name, job, location), "
"relationship (people and pets in their life), "
"or goal (things they want to achieve). "
"Only extract clear, factual statements. Return an empty list if nothing notable."
)
)
def extract_and_store_facts(user_input: str, user_id: str):
result = fact_extractor.run_sync(user_input)
for fact in result.data.facts:
structured_store.upsert(UserFact(
user_id=user_id,
fact_type=fact.fact_type,
key=fact.key,
value=fact.value
))
return result.data.facts
Hybrid Memory Manager
Now let us combine structured and vector memory into a unified memory manager:
from typing import List, Optional
from dataclasses import dataclass
@dataclass
class MemoryContext:
"""Aggregated memory context for a single conversation turn."""
structured_facts: List[UserFact]
semantic_memories: List[str]
def to_prompt_string(self) -> str:
parts = []
if self.structured_facts:
facts_text = "\n".join(
f" - {f.key} ({f.fact_type}): {f.value}"
for f in self.structured_facts
)
parts.append(f"Known facts about the user:\n{facts_text}")
if self.semantic_memories:
mem_text = "\n".join(f" - {m}" for m in self.semantic_memories)
parts.append(f"Relevant memories:\n{mem_text}")
return "\n\n".join(parts) if parts else "No prior memory about this user."
class HybridMemoryManager:
def __init__(
self,
structured: StructuredMemoryStore,
vector: VectorMemoryStore
):
self.structured = structured
self.vector = vector
def retrieve_context(self, query: str, user_id: str) -> MemoryContext:
# Get all structured facts (typically small enough to include all)
structured_facts = self.structured.get_all(user_id)
# Get semantically relevant memories
semantic_memories = self.vector.retrieve(query, user_id, n_results=5)
return MemoryContext(
structured_facts=structured_facts,
semantic_memories=semantic_memories
)
def store_turn(self, user_input: str, agent_response: str, user_id: str):
# Extract and store structured facts
extract_and_store_facts(user_input, user_id)
# Extract and store semantic memories
conversation = f"User: {user_input}\nAssistant: {agent_response}"
extract_and_store_memories(conversation, user_id)
memory_manager = HybridMemoryManager(structured_store, memory_store)
Using the Hybrid Memory Manager with an Agent
from pydantic_ai import Agent, RunContext
hybrid_agent = Agent(
model="openai:gpt-4o",
deps_type=str, # user_id passed as deps
system_prompt=(
"You are a personalized assistant with access to the user's memory. "
"Use the provided facts and memories to give personalized, context-aware responses."
)
)
@hybrid_agent.system_prompt
async def inject_hybrid_memory(ctx: RunContext[str]) -> str:
user_id = ctx.deps
query = ""
for msg in reversed(ctx.messages):
if hasattr(msg, "content") and isinstance(msg.content, str):
query = msg.content
break
if not query:
query = "general"
context = memory_manager.retrieve_context(query, user_id)
return context.to_prompt_string()
def hybrid_chat(user_input: str, user_id: str, history=None):
if history is None:
history = []
result = hybrid_agent.run_sync(
user_input,
deps=user_id,
message_history=history
)
# Store new memories from this turn
memory_manager.store_turn(user_input, result.data, user_id)
return result.data, result.all_messages()
# Demonstration
user_id = "user_456"
print("=== Turn 1 ===")
resp, hist = hybrid_chat(
"I'm Tom, a software engineer in Seattle. I prefer Python over JavaScript.",
user_id
)
print(f"Agent: {resp}")
print("\n=== Turn 2 ===")
resp, hist = hybrid_chat(
"I'm trying to learn Rust this year. Any tips?",
user_id, hist
)
print(f"Agent: {resp}")
print("\n=== Turn 3 (new session) ===")
resp, hist = hybrid_chat(
"What programming languages should I focus on?",
user_id
)
print(f"Agent: {resp}")
# Agent should reference Python preference, Seattle, and Rust learning goal
Implementing Episodic Memory with Summarization
Episodic memory stores records of past conversations. Rather than keeping full transcripts (which are expensive), a common approach is to summarize each conversation episode and store the summary. Let us implement this:
from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import List, Dict
from datetime import datetime
class EpisodeSummary(BaseModel):
summary: str = Field(description="A concise summary of the conversation")
key_topics: List[str] = Field(description="Main topics discussed")
user_sentiment: str = Field(description="Overall user sentiment: positive, neutral, negative")
action_items: List[str] = Field(default_factory=list, description="Any follow-up actions")
summarizer_agent = Agent(
model="openai:gpt-4o",
result_type=EpisodeSummary,
system_prompt=(
"Summarize the following conversation episode. Capture the key topics, "
"the user's overall sentiment, and any action items or follow-ups that were mentioned."
)
)
class EpisodicMemoryStore:
def __init__(self):
self._episodes: Dict[str, List[dict]] = {}
def store_episode(self, user_id: str, episode: EpisodeSummary, timestamp: datetime):
if user_id not in self._episodes:
self._episodes[user_id] = []
self._episodes[user_id].append({
"summary": episode.summary,
"key_topics": episode.key_topics,
"user_sentiment": episode.user_sentiment,
"action_items": episode.action_items,
"timestamp": timestamp.isoformat()
})
def get_episodes(self, user_id: str, limit: int = 10) -> List[dict]:
episodes = self._episodes.get(user_id, [])
return episodes[-limit:]
def search_episodes(self, user_id: str, topic: str) -> List[dict]:
episodes = self.get_episodes(user_id)
return [
ep for ep in episodes
if topic.lower() in ep["summary"].lower()
or any(topic.lower() in t.lower() for t in ep["key_topics"])
]
episodic_store = EpisodicMemoryStore()
def summarize_and_store_conversation(
messages: List,
user_id: str
):
"""Summarize a completed conversation and store it as an episode."""
# Format messages into text
conversation_text = "\n".join(
f"{msg.kind if hasattr(msg, 'kind') else 'Unknown'}: {msg.content}"
for msg in messages
if hasattr(msg, "content")
)
result = summarizer_agent.run_sync(
f"Summarize this conversation:\n\n{conversation_text}"
)
episodic_store.store_episode(user_id, result.data, datetime.now())
return result.data
Memory Consolidation and Forgetting
A critical but often overlooked aspect of memory architecture is forgetting. Without a strategy for pruning, consolidating, and forgetting memories, your memory store will grow unboundedly, leading to increased retrieval latency, higher storage costs, and potentially contradictory or stale information.
Memory Consolidation
Consolidation is the process of merging related memories, resolving contradictions, and maintaining a clean, coherent memory store. Here is a consolidation agent that periodically reviews and merges memories:
from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import List
class ConsolidatedMemory(BaseModel):
content: str = Field(description="The merged memory content")
should_replace: List[str] = Field(
description="IDs of memories that should be replaced by this consolidated version"
)
class ConsolidationResult(BaseModel):
consolidated: List[ConsolidatedMemory] = Field(default_factory=list)
to_delete: List[str] = Field(default_factory=list, description="IDs of memories to delete")
consolidation_agent = Agent(
model="openai:gpt-4o",
result_type=ConsolidationResult,
system_prompt=(
"You are a memory consolidation system. Given a list of memories about a user, "
"identify: (1) memories that can be merged into a single, more complete memory, "
"and (2) memories that are outdated, contradictory, or no longer relevant and should "
"be deleted. Be conservative — only consolidate clearly related memories and only "
"delete clearly outdated ones."
)
)
def consolidate_user_memories(user_id: str, memory_store: VectorMemoryStore):
"""Consolidate all memories for a given user."""
# Retrieve a large batch of memories
all_memories = memory_store.retrieve(
query="user information preferences facts",
user_id=user_id,
n_results=50
)
if len(all_memories) < 5:
return # Not enough memories to warrant consolidation
memories_text = "\n".join(
f"{i}. {m}" for i, m in enumerate(all_memories)
)
result = consolidation_agent.run_sync(
f"Consolidate these memories:\n\n{memories_text}"
)
# In a real implementation, you would map the text back to IDs
# and perform the actual replacement/deletion in the store
print(f"Consolidated {len(result.data.consolidated)} memory groups")
print(f"Flagged {len(result.data.to_delete)} memories for deletion")
return result.data
Time-Based Forgetting
Some memories should naturally expire. For example, a user's current location or job may change, while their name or core preferences are more stable. You can implement time-based decay:
from datetime import datetime, timedelta
from typing import Dict, List
class MemoryWithDecay:
def __init__(self, half_life_days: int = 90):
self.half_life = timedelta(days=half_life_days)
def should_expire(
self,
created_at: datetime,
fact_type: str,
current_time: datetime = None
) -> bool:
current_time = current_time or datetime.now()
age = current_time - created_at
# Different fact types have different retention policies
retention_policies = {
"profile": timedelta(days=365 * 5), # Very long retention
"preference": timedelta(days=365), # Long retention
"goal": timedelta(days=180), # Medium retention
"transient": timedelta(days=7), # Short retention
}
max_age = retention_policies.get(fact_type, timedelta(days=90))
return age > max_age
def get_relevance_score(
self,
created_at: datetime,
last_accessed: datetime,
current_time: datetime = None
) -> float:
"""Calculate a relevance score based on recency and access frequency."""
current_time = current_time or datetime.now()
age = (current_time - created_at).total_seconds() / 86400 # in days
half_life_days = self.half_life.days
# Exponential decay: relevance halves every half_life_days
return 0.5 ** (age / half_life_days)
decay_manager = MemoryWithDecay(half_life_days=90)
Best Practices for Agent Memory Architectures
1. Separate Memory Concerns
Do not try to cram everything into a single memory store. Different types of information have different storage and retrieval requirements. Structured facts need exact lookup; conversational memories need semantic search; episodic memories need temporal ordering. Use separate stores for each and aggregate them at retrieval time.
2. Be Selective About What You Store
Not every message contains memorable information. Use an extraction agent to filter out noise before storing. Storing irrelevant memories pollutes your retrieval results and increases costs. A good rule of thumb: if a human would not bother remembering it, your agent probably should not either.
3. Implement Memory Consolidation
Over time, memories accumulate, overlap, and sometimes contradict each other. Schedule periodic consolidation runs that merge related memories, resolve contradictions, and prune stale entries. This keeps your memory store clean and your retrieval results relevant.
4. Respect User Privacy
Memory systems store potentially sensitive user information. Always provide users with the ability to view, edit, and delete their stored memories. Implement clear data retention policies and ensure compliance with regulations like GDPR. Never store information the user has explicitly asked you to forget.
5. Handle Memory Retrieval Failures Gracefully
Your vector store might be down, or retrieval might return irrelevant results. Your agent should degrade gracefully — operating without memory rather than crashing. Always wrap memory retrieval in error handling:
def safe_retrieve(query: str, user_id: str, store: VectorMemoryStore) -> List[str]:
try:
return store.retrieve(query, user_id, n_results=5)
except Exception as e:
print(f"Memory retrieval failed: {e}")
return []
6. Use Appropriate Embedding Models
The quality of your semantic memory retrieval depends heavily on the embedding model. For most use cases, lightweight models like all-MiniLM-L6-v2 provide a good balance of speed and quality. For domain-specific applications, consider fine-tuning embeddings on your specific data.
7. Monitor Memory Store Health
Track metrics like memory count per user, average retrieval latency, retrieval relevance scores, and storage costs. Set up alerts for anomalous growth patterns. A sudden spike in stored memories might indicate a bug in your extraction logic.
8. Test Memory Behavior Explicitly
Memory introduces statefulness, which makes testing more complex. Write integration tests that verify memories are correctly stored, retrieved, and surfaced in agent responses. Test edge cases like empty memory stores, corrupted data, and concurrent access:
import pytest
def test_memory_storage_and_retrieval():
store = VectorMemoryStore(collection_name="test_memories")
user_id = "test_user"
store.store(Memory(
content="User loves Italian food",
user_id=user_id
))
results = store.retrieve("What food does the user like?", user_id)
assert len(results) > 0
assert any("Italian" in r for r in results)
def test_empty_memory_retrieval():
store = VectorMemoryStore(collection_name="test_memories_empty")
results = store.retrieve("anything", "nonexistent_user")
assert results == []
def test_memory_isolation_between_users():
store = VectorMemoryStore(collection_name="test_isolation")
store.store(Memory(content="User A likes tennis", user_id="user_a"))
store.store(Memory(content="User B likes chess", user_id="user_b"))
results_a = store.retrieve("sports", "user_a")
results_b = store.retrieve("sports", "user_b")
assert all("tennis" in r or "User A" in r for r in results_a)
assert all("chess" in r or "User B" in r for r in results_b)
9. Consider Token Budgets
Memories injected into the system prompt consume tokens. Be mindful of how many memories you retrieve and how you format them. A practical approach is to allocate a fixed token budget for memory context and truncate or prioritize within that budget:
def format_memories_within_budget(
memories: List[str],
max_tokens: int = 500,
chars_per_token: int = 4
) -> str:
max_chars = max_tokens * chars_per_token
formatted = []
total_chars = 0
for memory in memories:
entry = f"- {memory}"
if total_chars + len(entry) > max_chars:
break
formatted.append(entry)
total_chars += len(entry)
return "\n".join(formatted) if formatted else "No memories available."
10. Version Your Memory Schema
As your application evolves, the structure of your memories may change. Include a schema version in your memory metadata so you can migrate old memories when you change your storage format. This prevents silent breakage when you update your memory extraction logic.
Putting It All Together: A Complete Memory-Enabled Agent
Let us assemble everything we have built into a complete, production-ready memory architecture:
from pydantic_ai import Agent, RunContext
from pydantic import BaseModel, Field
from pydantic_ai.messages import ModelMessage
from typing import List, Dict, Optional
from datetime import datetime
import uuid
# --- Memory Models ---
class Memory(BaseModel):
content: str
user_id: str
metadata: dict = Field(default_factory=dict)
class UserFact(BaseModel):
user_id: str
fact_type: str
key: str
value: str
confidence: float = 1.0
updated_at: datetime = Field(default_factory=datetime.now)
# --- Memory Stores ---
class InMemory