Building a Memory-Augmented Agent
Duration: 15 min
Building a Memory-Augmented Agent
Duration: 15 min
Overview
This capstone module brings together everything from the course: working memory, short-term buffers, long-term vector storage, episodic memory, knowledge graphs, retrieval strategies, and compression. You will build a complete memory-augmented agent that remembers across conversations, learns from experience, and retrieves contextually relevant information to improve its responses.
Key Concepts & Foundations
- What: A memory-augmented agent integrates multiple memory systems into a unified architecture
- Why: Real agents need all memory types working together - no single system handles every use case
- How: Layer memory systems with a unified retrieval interface and memory lifecycle management
Full Memory Architecture
User Message
|
v
[Working Memory] <-- current context window
|
+--> [Short-Term Buffer] <-- recent conversation
|
+--> [Long-Term Store] <-- vector search for relevant past memories
|
+--> [Knowledge Graph] <-- structured facts about the user
|
+--> [Episodic Memory] <-- similar past experiences
|
v
[Memory Fusion] --> Combined context for LLM
|
v
[LLM Response]
|
v
[Memory Update] --> Store new facts, update graph, record episode
Design Principles
1. Read before write - Always retrieve existing memories before storing new ones (avoids duplicates) 2. Layered retrieval - Fast stores first (working memory), slow stores last (vector search) 3. Budget-aware - Never exceed the context window; compress or drop low-priority memories 4. Graceful degradation - If memory retrieval fails, the agent should still function (just without context)
Hands-On Implementation
from dataclasses import dataclass, field
from datetime import datetime
from collections import defaultdict
import hashlib
import math
@dataclass
class UnifiedMemory:
content: str
source: str # "working", "short_term", "long_term", "knowledge", "episode"
relevance: float = 0.0
importance: float = 0.5
class MemoryAugmentedAgent:
"""A complete agent with integrated multi-layer memory."""
def __init__(self, agent_name: str = "Assistant", context_budget: int = 4000):
self.name = agent_name
self.context_budget = context_budget
# Working memory (current conversation)
self.working_memory: list[dict] = []
# Short-term buffer (session history)
self.short_term: list[dict] = []
self.short_term_limit = 20
# Long-term memory (persistent facts)
self.long_term: list[dict] = []
# Knowledge graph (structured facts)
self.knowledge: dict[str, list[tuple[str, str]]] = defaultdict(list)
# Episodic memory (past experiences)
self.episodes: list[dict] = []
def receive_message(self, user_message: str) -> str:
"""Process a user message with full memory augmentation."""
# 1. Retrieve relevant memories
context = self._build_context(user_message)
# 2. Generate response (simulated - replace with real LLM call)
response = self._generate_response(user_message, context)
# 3. Update memories
self._update_memories(user_message, response)
return response
def _build_context(self, query: str) -> list[UnifiedMemory]:
"""Retrieve and rank memories from all stores."""
memories = []
# Working memory (always included)
for msg in self.working_memory[-3:]:
memories.append(UnifiedMemory(
content=f"[recent] {msg['role']}: {msg['content']}",
source="working", relevance=1.0, importance=0.7
))
# Long-term retrieval (keyword matching for demo)
query_words = set(query.lower().split())
for mem in self.long_term:
overlap = len(query_words & set(mem["content"].lower().split()))
if overlap > 0:
memories.append(UnifiedMemory(
content=f"[remembered] {mem['content']}",
source="long_term",
relevance=overlap / len(query_words),
importance=mem.get("importance", 0.5)
))
# Knowledge graph facts
for entity in query_words:
for predicate, obj in self.knowledge.get(entity, []):
memories.append(UnifiedMemory(
content=f"[fact] {entity} {predicate} {obj}",
source="knowledge", relevance=0.8, importance=0.7
))
# Episodic memory (similar past tasks)
for ep in self.episodes:
ep_words = set(ep["trigger"].lower().split())
overlap = len(query_words & ep_words)
if overlap > 0:
memories.append(UnifiedMemory(
content=f"[experience] {ep['summary']}",
source="episode", relevance=overlap / max(len(query_words), 1),
importance=0.6
))
# Rank by combined score and fit within budget
memories.sort(key=lambda m: m.relevance 0.6 + m.importance 0.4, reverse=True)
return self._fit_to_budget(memories)
def _fit_to_budget(self, memories: list[UnifiedMemory]) -> list[UnifiedMemory]:
"""Select memories that fit within the context budget."""
selected = []
chars_used = 0
for mem in memories:
if chars_used + len(mem.content) > self.context_budget:
break
selected.append(mem)
chars_used += len(mem.content)
return selected
def _generate_response(self, message: str, context: list[UnifiedMemory]) -> str:
"""Simulate LLM response (replace with actual API call)."""
context_str = "\n".join(m.content for m in context)
# In production: call OpenAI/Bedrock/etc with context + message
return f"[Response using {len(context)} memories about: {message[:50]}]"
def _update_memories(self, user_message: str, response: str):
"""Update all memory stores after an interaction."""
# Update working memory
self.working_memory.append({"role": "user", "content": user_message})
self.working_memory.append({"role": "assistant", "content": response})
# Overflow to short-term
while len(self.working_memory) > 6:
evicted = self.working_memory.pop(0)
self.short_term.append(evicted)
# Extract and store facts (simple heuristic)
if "my name is" in user_message.lower():
name = user_message.split("is")[-1].strip().rstrip(".")
self.knowledge["user"].append(("name_is", name))
self.long_term.append({"content": f"User's name is {name}", "importance": 0.9})
if "i prefer" in user_message.lower() or "i like" in user_message.lower():
self.long_term.append({"content": user_message, "importance": 0.7})
def add_knowledge(self, entity: str, predicate: str, obj: str):
"""Manually add a fact to the knowledge graph."""
self.knowledge[entity].append((predicate, obj))
def add_episode(self, trigger: str, summary: str, outcome: str):
"""Record a past experience."""
self.episodes.append({
"trigger": trigger, "summary": summary, "outcome": outcome
})
def get_memory_stats(self) -> dict:
return {
"working_memory": len(self.working_memory),
"short_term": len(self.short_term),
"long_term": len(self.long_term),
"knowledge_entities": len(self.knowledge),
"episodes": len(self.episodes),
}
Demo: Full memory-augmented agent session
agent = MemoryAugmentedAgent("CodeMentor")Pre-load some memories from past sessions
agent.add_knowledge("python", "is_a", "programming language")
agent.add_knowledge("user", "prefers", "type hints")
agent.long_term.append({"content": "User is building an ML pipeline with PyTorch", "importance": 0.8})
agent.add_episode(
trigger="deploy model to production",
summary="Used canary deployment after staging tests passed. Successful.",
outcome="success"
)Simulate a conversation
print("=== Agent Conversation ===\n")
r1 = agent.receive_message("Hi, my name is Sarah")
print(f"User: Hi, my name is Sarah\nAgent: {r1}\n")r2 = agent.receive_message("I prefer concise Python examples")
print(f"User: I prefer concise Python examples\nAgent: {r2}\n")
r3 = agent.receive_message("Help me deploy my model to production")
print(f"User: Help me deploy my model to production\nAgent: {r3}\n")
print("=== Memory Stats ===")
print(agent.get_memory_stats())
Advanced Techniques
1. Memory Lifecycle Management - Automate the full cycle: create, retrieve, update, compress, archive, delete 2. Multi-Agent Shared Memory - Let multiple specialized agents read/write a shared memory layer 3. Memory Observability - Log what memories are retrieved per turn; debug why the agent forgot something 4. A/B Testing Memory Strategies - Compare retrieval approaches by measuring response quality 5. User Memory Controls - Let users view, edit, and delete what the agent remembers about them (privacy)
Quiz
Q1: What is the correct order of retrieval in a layered memory system?
- A) Long-term first, then working memory
- B) Working memory first, then short-term, then long-term (fast to slow) ✓
- C) Random order does not matter
- D) Episodic memory first always
Q2: Why should an agent "read before write" when updating memory?
- A) To be polite
- B) To check if the information already exists and avoid storing duplicates ✓
- C) To slow down the response
- D) Because databases require it
Q3: What does graceful degradation mean for agent memory?
- A) The agent crashes politely
- B) If memory retrieval fails, the agent still functions without context rather than erroring ✓
- C) The agent slowly forgets everything
- D) Memory quality decreases over time