Imagine you're building a customer support agent for a SaaS platform. On Monday, a user named Priya opens a ticket about a billing discrepancy. The agent walks her through the resolution. On Friday, Priya is back — same account, new issue, but she mentions "like that billing problem last week." Your agent has no idea what she's talking about. It's starting from a blank slate, treating Priya like a stranger. The support experience degrades, trust erodes, and your agent looks less intelligent than a help desk ticket system from 2003.
This is the memory problem in AI agents, and it's one of the most underappreciated engineering challenges in production systems. Most tutorials show you how to connect an LLM to tools. Far fewer show you how agents remember, what they remember, and how to make that memory retrieval fast, relevant, and cost-effective. The difference between a demo-quality agent and a production-quality one often comes down entirely to memory architecture.
By the end of this lesson, you'll have built a working multi-memory agent that combines in-context short-term memory, vector-based long-term retrieval, and a structured episodic memory store. You'll understand when to use each layer, how they interact, and what breaks in production when you get it wrong.
What you'll learn:
You should be comfortable with:
You'll need: langchain, langchain-openai, chromadb, tiktoken, python-dotenv, sqlalchemy
pip install langchain langchain-openai chromadb tiktoken python-dotenv sqlalchemy
Before writing a single line of code, you need to understand the architectural distinction at the heart of this lesson, because conflating these two concepts causes a class of subtle bugs that are genuinely painful to debug.
Tool memory refers to memory that the agent actively writes to and reads from as part of its reasoning loop. It's procedural. The agent decides to call save_to_memory(...) or load_context(...) as a deliberate action during its chain-of-thought. This is memory as a tool — it appears in the agent's available actions just like a web search tool or a database query tool.
Retrieval memory refers to memory that is automatically fetched and injected into context based on relevance, without the agent explicitly choosing to retrieve it. The retrieval happens in the pipeline before the LLM even sees the prompt. This is memory as infrastructure — it runs beneath the agent's reasoning, shaping what it knows before it starts thinking.
Neither pattern is universally better. They have different failure modes and different strengths:
| Dimension | Tool Memory | Retrieval Memory |
|---|---|---|
| Agent awareness | Agent knows it's accessing memory | Memory is injected transparently |
| Precision | High — agent chooses what to fetch | Variable — depends on embedding similarity |
| Latency overhead | Adds a reasoning step | Adds a retrieval step at pipeline entry |
| Token efficiency | Fetches exactly what's needed | May inject irrelevant context |
| Best for | Structured data, explicit lookups | Semantic similarity, fuzzy recall |
In a well-designed production agent, you'll use both. Let's look at how to layer them.
Short-term memory is the simplest of the three layers: it's the conversation history that lives inside the active context window. But "simple" doesn't mean trivial — managing this buffer poorly is the number one cause of runaway token costs and degraded performance in deployed agents.
The naive implementation just appends every message to a list and passes the whole thing to the LLM. This works for five exchanges. At exchange forty, you've burned through your context window, and at exchange four hundred, you're paying for tokens that describe a problem resolved three hours ago.
You need a buffer strategy. The three common approaches are:
The summarizing buffer is almost always the right choice for production agents. Here's a clean implementation:
import os
from dataclasses import dataclass, field
from typing import Optional
import tiktoken
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, AIMessage, SystemMessage, BaseMessage
@dataclass
class ConversationBuffer:
"""
A token-aware conversation buffer that summarizes older context
when the buffer exceeds a configurable threshold.
"""
max_tokens: int = 2000
model_name: str = "gpt-4o"
messages: list[BaseMessage] = field(default_factory=list)
summary: Optional[str] = None
def __post_init__(self):
self.encoder = tiktoken.encoding_for_model(self.model_name)
self.llm = ChatOpenAI(model=self.model_name, temperature=0)
def _count_tokens(self, text: str) -> int:
return len(self.encoder.encode(text))
def _total_buffer_tokens(self) -> int:
total = 0
if self.summary:
total += self._count_tokens(self.summary)
for msg in self.messages:
total += self._count_tokens(msg.content)
return total
def add_message(self, message: BaseMessage) -> None:
self.messages.append(message)
if self._total_buffer_tokens() > self.max_tokens:
self._compress()
def _compress(self) -> None:
"""Summarize the oldest half of the conversation buffer."""
midpoint = len(self.messages) // 2
messages_to_summarize = self.messages[:midpoint]
self.messages = self.messages[midpoint:]
existing_context = f"Previous summary: {self.summary}\n\n" if self.summary else ""
messages_text = "\n".join(
f"{'User' if isinstance(m, HumanMessage) else 'Assistant'}: {m.content}"
for m in messages_to_summarize
)
summary_prompt = f"""Summarize the following conversation excerpt concisely.
Preserve: key decisions made, problems resolved, user preferences stated,
any account or technical details mentioned.
{existing_context}Conversation to summarize:
{messages_text}
Summary:"""
response = self.llm.invoke([HumanMessage(content=summary_prompt)])
self.summary = response.content
print(f"[Buffer compressed] New summary: {self.summary[:100]}...")
def get_context_messages(self) -> list[BaseMessage]:
"""Return the full context including summary as a system message."""
messages = []
if self.summary:
messages.append(SystemMessage(
content=f"Context from earlier in this conversation:\n{self.summary}"
))
messages.extend(self.messages)
return messages
Why the summarizing buffer beats fixed windows: A fixed window of the last 10 messages loses everything said in message 1 — which might include the user's original problem statement, their account tier, or a constraint they established. The summarizing buffer compresses rather than discards, keeping semantic content while reducing token cost.
Let's verify it works before moving on:
buffer = ConversationBuffer(max_tokens=500) # Low threshold for demo
buffer.add_message(HumanMessage(content="Hi, I'm Priya. I'm on the Enterprise plan and I'm having a billing issue."))
buffer.add_message(AIMessage(content="Hi Priya! I can help with that. Can you describe the billing issue?"))
buffer.add_message(HumanMessage(content="I was charged twice for the same invoice in November."))
buffer.add_message(AIMessage(content="I can see that. A duplicate charge of $2,400 was made on Nov 14. I'll escalate this to billing."))
# ... add more messages until compression triggers
context = buffer.get_context_messages()
for msg in context:
print(f"[{type(msg).__name__}] {msg.content[:80]}")
Long-term memory answers a different question than short-term memory. Short-term memory asks: what happened in this conversation? Long-term memory asks: what do we know about this user, this topic, or this type of problem — across all past interactions?
The implementation uses a vector store to semantically index facts, and a retriever to pull the most relevant ones at the start of each new interaction.
The biggest mistake practitioners make with long-term vector memory is storing raw conversation transcripts. Retrieving chunks of raw dialogue is noisy and wasteful. Instead, store distilled facts — structured, self-contained statements about entities in your system.
from dataclasses import dataclass
from datetime import datetime
import json
@dataclass
class MemoryFact:
"""
A single distilled fact for long-term storage.
Designed to be self-contained when retrieved out of context.
"""
entity_id: str # e.g., "user:priya@company.com"
entity_type: str # "user", "account", "product", "issue"
fact: str # Human-readable, self-contained statement
confidence: float # 0.0 to 1.0 — how certain is this fact?
source_session: str # Which session generated this fact
created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
metadata: dict = field(default_factory=dict)
def to_document_text(self) -> str:
"""Format for embedding. Context-rich but concise."""
return f"[{self.entity_type.upper()}:{self.entity_id}] {self.fact}"
import chromadb
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from typing import List
import uuid
class LongTermMemoryStore:
"""
Vector-backed long-term memory for an AI agent.
Stores distilled facts about users, accounts, and past interactions.
"""
def __init__(self, collection_name: str = "agent_long_term_memory", persist_dir: str = "./memory_db"):
self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
self.persist_dir = persist_dir
self.collection_name = collection_name
self.vectorstore = Chroma(
collection_name=collection_name,
embedding_function=self.embeddings,
persist_directory=persist_dir,
)
def store_fact(self, fact: MemoryFact) -> str:
"""Store a single distilled fact. Returns the document ID."""
doc_text = fact.to_document_text()
doc_metadata = {
"entity_id": fact.entity_id,
"entity_type": fact.entity_type,
"confidence": fact.confidence,
"source_session": fact.source_session,
"created_at": fact.created_at,
}
doc_metadata.update(fact.metadata)
doc_id = str(uuid.uuid4())
self.vectorstore.add_documents(
documents=[Document(page_content=doc_text, metadata=doc_metadata)],
ids=[doc_id]
)
return doc_id
def store_facts_batch(self, facts: List[MemoryFact]) -> List[str]:
"""Batch store multiple facts (more efficient for end-of-session writes)."""
documents = []
ids = []
for fact in facts:
doc_id = str(uuid.uuid4())
documents.append(Document(
page_content=fact.to_document_text(),
metadata={
"entity_id": fact.entity_id,
"entity_type": fact.entity_type,
"confidence": fact.confidence,
"source_session": fact.source_session,
"created_at": fact.created_at,
**fact.metadata,
}
))
ids.append(doc_id)
self.vectorstore.add_documents(documents=documents, ids=ids)
return ids
def retrieve_relevant_facts(
self,
query: str,
entity_id: Optional[str] = None,
top_k: int = 5,
min_confidence: float = 0.5
) -> List[Document]:
"""
Retrieve facts semantically similar to query.
Optionally filter to a specific entity (e.g., a specific user).
"""
filter_dict = {"confidence": {"$gte": min_confidence}}
if entity_id:
filter_dict["entity_id"] = entity_id
results = self.vectorstore.similarity_search(
query=query,
k=top_k,
filter=filter_dict
)
return results
def format_for_context(self, documents: List[Document]) -> str:
"""Format retrieved facts as a clean context block for the LLM."""
if not documents:
return ""
lines = ["Relevant context from memory:"]
for doc in documents:
confidence = doc.metadata.get("confidence", 1.0)
created = doc.metadata.get("created_at", "unknown")[:10]
lines.append(f"- [{created}] {doc.page_content} (confidence: {confidence:.0%})")
return "\n".join(lines)
Long-term memory is only as good as the extraction process that feeds it. After each session, you need to distill the conversation into storable facts. This is where most implementations cut corners — and then wonder why their retrieval is noisy.
from langchain_openai import ChatOpenAI
import json
class MemoryExtractor:
"""
Extracts structured facts from conversation history
and writes them to long-term memory.
"""
def __init__(self, llm: ChatOpenAI, memory_store: LongTermMemoryStore):
self.llm = llm
self.memory_store = memory_store
def extract_and_store(
self,
conversation_summary: str,
full_messages: list,
session_id: str,
entity_id: str
) -> List[str]:
"""
Extract facts from a completed session and store them in long-term memory.
Returns list of stored document IDs.
"""
messages_text = "\n".join(
f"{'User' if isinstance(m, HumanMessage) else 'Assistant'}: {m.content}"
for m in full_messages[-20:] # Last 20 messages + summary is enough context
)
extraction_prompt = f"""You are a memory extraction system. Analyze this conversation
and extract factual statements worth remembering for future interactions.
Entity ID: {entity_id}
Session ID: {session_id}
Conversation summary: {conversation_summary}
Recent messages:
{messages_text}
Extract facts in this JSON format. Only include facts that would be useful
in a future unrelated conversation with this user:
{{
"facts": [
{{
"fact": "Self-contained factual statement",
"entity_type": "user|account|issue|preference|technical",
"confidence": 0.0-1.0,
"rationale": "Why this is worth remembering"
}}
]
}}
Focus on: preferences stated, problems resolved, account details, technical constraints,
promises made, and user expertise level. Skip pleasantries and procedural steps.
Return valid JSON only."""
response = self.llm.invoke([HumanMessage(content=extraction_prompt)])
try:
data = json.loads(response.content)
facts_data = data.get("facts", [])
except json.JSONDecodeError:
print(f"[MemoryExtractor] Failed to parse LLM response as JSON")
return []
facts = [
MemoryFact(
entity_id=entity_id,
entity_type=f_data["entity_type"],
fact=f_data["fact"],
confidence=f_data["confidence"],
source_session=session_id,
)
for f_data in facts_data
if f_data.get("confidence", 0) >= 0.5 # Filter low-confidence extractions
]
if facts:
stored_ids = self.memory_store.store_facts_batch(facts)
print(f"[MemoryExtractor] Stored {len(stored_ids)} facts from session {session_id}")
return stored_ids
return []
The confidence threshold is load-bearing. Storing everything with confidence 1.0 floods your vector store with noise. When you retrieve the top-5 facts for a new session, you want signal, not volume. The LLM-assigned confidence — when prompted carefully — is a reasonable proxy for fact reliability.
Episodic memory answers yet a different question: what happened in past sessions as discrete events? It's not "what do we know about Priya" (that's long-term memory) — it's "what happened during Priya's session on November 20th, and how was it resolved?"
Think of it as your agent's case history. It's particularly valuable for support agents, tutoring agents, and any domain where patterns across sessions matter.
Episodic memory benefits from being structured, not just embedded. We'll store it in SQLite via SQLAlchemy so we can query by date, status, user, and resolution type — things semantic search handles poorly.
from sqlalchemy import create_engine, Column, String, Float, DateTime, Text, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import uuid
Base = declarative_base()
class Episode(Base):
"""
A single completed interaction session stored as a structured episode.
"""
__tablename__ = "episodes"
id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
session_id = Column(String, nullable=False, index=True)
entity_id = Column(String, nullable=False, index=True)
# Temporal
started_at = Column(DateTime, nullable=False)
ended_at = Column(DateTime, nullable=False)
duration_seconds = Column(Float)
# Content
summary = Column(Text, nullable=False)
issue_category = Column(String) # e.g., "billing", "technical", "onboarding"
resolution_status = Column(String) # "resolved", "escalated", "pending", "unresolved"
sentiment_score = Column(Float) # -1.0 to 1.0
# Structured outcomes
actions_taken = Column(JSON) # List of actions the agent performed
follow_up_required = Column(JSON) # Any follow-up tasks
key_entities = Column(JSON) # Account IDs, ticket numbers, etc.
created_at = Column(DateTime, default=datetime.utcnow)
class EpisodicMemoryStore:
"""
SQLite-backed episodic memory store for structured session history.
"""
def __init__(self, db_url: str = "sqlite:///episodic_memory.db"):
self.engine = create_engine(db_url, echo=False)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine)
def store_episode(self, episode_data: dict) -> str:
"""Store a completed episode. Returns the episode ID."""
episode = Episode(**episode_data)
with self.Session() as session:
session.add(episode)
session.commit()
return episode.id
def get_recent_episodes(
self,
entity_id: str,
limit: int = 5,
issue_category: Optional[str] = None
) -> List[Episode]:
"""Fetch recent episodes for a user, optionally filtered by category."""
with self.Session() as session:
query = session.query(Episode).filter(
Episode.entity_id == entity_id
)
if issue_category:
query = query.filter(Episode.issue_category == issue_category)
return query.order_by(Episode.ended_at.desc()).limit(limit).all()
def get_unresolved_episodes(self, entity_id: str) -> List[Episode]:
"""Fetch episodes that need follow-up — critical for continuity."""
with self.Session() as session:
return session.query(Episode).filter(
Episode.entity_id == entity_id,
Episode.resolution_status.in_(["pending", "escalated"])
).order_by(Episode.ended_at.desc()).all()
def format_episodes_for_context(self, episodes: List[Episode]) -> str:
"""Format past episodes as structured context for the LLM."""
if not episodes:
return ""
lines = ["Past interaction history:"]
for ep in episodes:
date = ep.ended_at.strftime("%Y-%m-%d")
status_emoji = {"resolved": "✓", "escalated": "↑", "pending": "⏳", "unresolved": "✗"}.get(ep.resolution_status, "?")
lines.append(
f"\n[{date}] {status_emoji} {ep.issue_category or 'General'} | Status: {ep.resolution_status}"
)
lines.append(f" Summary: {ep.summary}")
if ep.follow_up_required:
lines.append(f" Follow-up needed: {', '.join(ep.follow_up_required)}")
return "\n".join(lines)
Like long-term facts, episodes need to be generated — typically by the agent at the end of each session, or by a background process that runs on session close.
class EpisodeGenerator:
"""
Generates structured Episode records from completed conversations.
"""
def __init__(self, llm: ChatOpenAI, episodic_store: EpisodicMemoryStore):
self.llm = llm
self.episodic_store = episodic_store
def generate_and_store(
self,
session_id: str,
entity_id: str,
buffer: ConversationBuffer,
started_at: datetime,
ended_at: datetime
) -> str:
"""Generate an episode from a completed conversation buffer."""
messages_text = "\n".join(
f"{'User' if isinstance(m, HumanMessage) else 'Agent'}: {m.content}"
for m in buffer.messages
)
classification_prompt = f"""Analyze this completed support conversation and return structured JSON.
Session: {session_id}
Buffer summary: {buffer.summary or 'N/A'}
Messages:
{messages_text}
Return JSON matching this exact structure:
{{
"summary": "2-3 sentence factual summary of what happened",
"issue_category": "billing|technical|onboarding|account|general",
"resolution_status": "resolved|escalated|pending|unresolved",
"sentiment_score": -1.0,
"actions_taken": ["list", "of", "concrete", "actions"],
"follow_up_required": ["list of follow-up tasks, empty if none"],
"key_entities": {{"ticket_id": "...", "invoice_id": "...", "etc": "..."}}
}}
Return valid JSON only."""
response = self.llm.invoke([HumanMessage(content=classification_prompt)])
try:
ep_data = json.loads(response.content)
except json.JSONDecodeError:
# Fallback: store minimal episode rather than losing the session
ep_data = {
"summary": buffer.summary or "Session completed — extraction failed.",
"issue_category": "general",
"resolution_status": "unresolved",
"sentiment_score": 0.0,
"actions_taken": [],
"follow_up_required": [],
"key_entities": {}
}
episode_record = {
"session_id": session_id,
"entity_id": entity_id,
"started_at": started_at,
"ended_at": ended_at,
"duration_seconds": (ended_at - started_at).total_seconds(),
**ep_data
}
episode_id = self.episodic_store.store_episode(episode_record)
print(f"[EpisodeGenerator] Stored episode {episode_id} for entity {entity_id}")
return episode_id
Now we build the actual agent that uses all three memory layers together. The key design decision here is when each layer fires:
from langchain_openai import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage, AIMessage
from datetime import datetime
import uuid
class MultiMemoryAgent:
"""
A production-grade agent that integrates short-term, long-term,
and episodic memory to provide contextually aware responses.
"""
SYSTEM_PROMPT = """You are a knowledgeable support agent for ArcSaaS, an enterprise SaaS platform.
You have access to the user's interaction history and known facts about their account.
Always acknowledge relevant past context naturally — don't call attention to the fact that
you're using memory. Just use it the way a good human support agent would.
If you learn new facts about the user during this conversation (preferences, account details,
constraints), note them in your reasoning. They'll be stored after the session ends."""
def __init__(
self,
llm: ChatOpenAI,
short_term: ConversationBuffer,
long_term: LongTermMemoryStore,
episodic: EpisodicMemoryStore,
memory_extractor: MemoryExtractor,
episode_generator: EpisodeGenerator,
):
self.llm = llm
self.short_term = short_term
self.long_term = long_term
self.episodic = episodic
self.extractor = memory_extractor
self.episode_generator = episode_generator
self.session_id = str(uuid.uuid4())
self.entity_id: Optional[str] = None
self.started_at: Optional[datetime] = None
self._session_context_injected = False
def start_session(self, entity_id: str) -> str:
"""
Initialize a session for a known entity.
Retrieves and injects all relevant past context.
"""
self.entity_id = entity_id
self.started_at = datetime.utcnow()
# 1. Pull unresolved episodes — highest priority
unresolved = self.episodic.get_unresolved_episodes(entity_id)
unresolved_context = self.episodic.format_episodes_for_context(unresolved)
# 2. Pull recent resolved episodes for broader context
recent = self.episodic.get_recent_episodes(entity_id, limit=3)
recent_context = self.episodic.format_episodes_for_context(recent)
# 3. Pull relevant long-term facts (use entity_id as broad context query)
long_term_docs = self.long_term.retrieve_relevant_facts(
query=f"user profile preferences account details {entity_id}",
entity_id=entity_id,
top_k=6
)
long_term_context = self.long_term.format_for_context(long_term_docs)
# 4. Assemble session context block
context_parts = []
if unresolved_context:
context_parts.append(f"⚠️ UNRESOLVED ISSUES FROM PRIOR SESSIONS:\n{unresolved_context}")
if recent_context:
context_parts.append(f"RECENT SESSION HISTORY:\n{recent_context}")
if long_term_context:
context_parts.append(long_term_context)
if context_parts:
full_context = "\n\n".join(context_parts)
# Inject as a system message at the start of the buffer
self.short_term.messages.append(
SystemMessage(content=f"SESSION CONTEXT FOR {entity_id}:\n\n{full_context}")
)
self._session_context_injected = True
return self.session_id
def chat(self, user_message: str) -> str:
"""Process a user message and return the agent's response."""
# Mid-session long-term retrieval: fetch facts relevant to THIS message
if self.entity_id and self._session_context_injected:
mid_session_docs = self.long_term.retrieve_relevant_facts(
query=user_message,
entity_id=self.entity_id,
top_k=3,
min_confidence=0.7 # Higher threshold for mid-session — be selective
)
if mid_session_docs:
relevant_context = self.long_term.format_for_context(mid_session_docs)
# Inject as a system hint before the user message
self.short_term.add_message(
SystemMessage(content=f"[Relevant memory for current query]\n{relevant_context}")
)
# Add the user's message to the buffer
self.short_term.add_message(HumanMessage(content=user_message))
# Build the full prompt: system + buffer context
messages = [SystemMessage(content=self.SYSTEM_PROMPT)]
messages.extend(self.short_term.get_context_messages())
# Call the LLM
response = self.llm.invoke(messages)
agent_reply = response.content
# Add response to buffer
self.short_term.add_message(AIMessage(content=agent_reply))
return agent_reply
def end_session(self) -> dict:
"""
Close the session: extract facts for long-term memory
and generate an episodic record.
"""
ended_at = datetime.utcnow()
# Generate episode record
episode_id = self.episode_generator.generate_and_store(
session_id=self.session_id,
entity_id=self.entity_id,
buffer=self.short_term,
started_at=self.started_at,
ended_at=ended_at
)
# Extract long-term facts
stored_fact_ids = self.extractor.extract_and_store(
conversation_summary=self.short_term.summary or "",
full_messages=self.short_term.messages,
session_id=self.session_id,
entity_id=self.entity_id
)
return {
"session_id": self.session_id,
"episode_id": episode_id,
"facts_stored": len(stored_fact_ids),
"duration_seconds": (ended_at - self.started_at).total_seconds()
}
Let's run the complete system end-to-end. This simulates two sessions with the same user — the second session should demonstrate memory continuity.
import os
from dotenv import load_dotenv
load_dotenv()
def build_agent() -> MultiMemoryAgent:
llm = ChatOpenAI(model="gpt-4o", temperature=0.2)
short_term = ConversationBuffer(max_tokens=3000)
long_term_store = LongTermMemoryStore(persist_dir="./memory_db")
episodic_store = EpisodicMemoryStore(db_url="sqlite:///episodic_memory.db")
extractor = MemoryExtractor(llm=llm, memory_store=long_term_store)
generator = EpisodeGenerator(llm=llm, episodic_store=episodic_store)
return MultiMemoryAgent(
llm=llm,
short_term=short_term,
long_term=long_term_store,
episodic=episodic_store,
memory_extractor=extractor,
episode_generator=generator,
)
# ── SESSION 1: Priya's first interaction ─────────────────────────────────────
print("=" * 60)
print("SESSION 1 — First contact with Priya")
print("=" * 60)
agent_session1 = build_agent()
agent_session1.start_session(entity_id="user:priya@techcorp.com")
exchanges = [
"Hi, I'm Priya from TechCorp. We're on the Enterprise plan and I just noticed we were charged twice for the November invoice.",
"The invoice number is INV-2024-1147. The charge is $4,200 and it hit our card on Nov 14 and again on Nov 16.",
"That's great. Also — we prefer to be billed quarterly, not monthly. Can you update that preference?",
"Perfect. One more thing — we're planning to add 50 more seats in Q1. Can you put us in touch with our account manager?",
]
for user_msg in exchanges:
print(f"\nUser: {user_msg}")
response = agent_session1.chat(user_msg)
print(f"Agent: {response}")
session1_summary = agent_session1.end_session()
print(f"\n[Session 1 closed] {session1_summary}")
# ── SESSION 2: Priya returns a week later ────────────────────────────────────
print("\n" + "=" * 60)
print("SESSION 2 — Priya returns (should demonstrate memory)")
print("=" * 60)
agent_session2 = build_agent()
agent_session2.start_session(entity_id="user:priya@techcorp.com")
# Notice: Priya doesn't re-introduce herself. The agent should know who she is.
exchanges_2 = [
"Hi, it's Priya again. Did that duplicate charge get resolved?",
"Good. I'm also here about adding those Q1 seats we mentioned — we're ready to move forward with 50 seats.",
"Can you remind me what our current per-seat pricing is on Enterprise?",
]
for user_msg in exchanges_2:
print(f"\nUser: {user_msg}")
response = agent_session2.chat(user_msg)
print(f"Agent: {response}")
session2_summary = agent_session2.end_session()
print(f"\n[Session 2 closed] {session2_summary}")
When you run this, watch the agent in session 2. It should:
If it doesn't, the retrieval or extraction step has a problem — which leads us to troubleshooting.
Symptom: Retrieval returns irrelevant facts even for clearly relevant queries.
Cause: You're embedding raw conversation text instead of the distilled fact field. Conversation text is long and contains many topics — the embedding becomes an averaged blur.
Fix: Embed only the clean, self-contained fact statement. Your to_document_text() method should produce something like: [USER:priya@techcorp.com] Prefers quarterly billing over monthly billing. — not a paragraph of dialogue.
Symptom: Priya's session returns facts about another user with a similar billing issue.
Cause: You're doing pure semantic retrieval without filtering by entity_id.
Fix: Always pass the entity_id filter to similarity_search. In ChromaDB's where clause syntax: {"entity_id": entity_id}. The filter runs at the metadata level before the similarity ranking, so it doesn't hurt recall for the right entity.
Symptom: Episode records exist but long-term facts are empty, so retrieval is sparse.
Cause: You called store_episode() but forgot to call extract_and_store(), or you called them in the wrong order in end_session().
Fix: In end_session(), generate the episode first (it uses the buffer), then extract facts (it also uses the buffer). Both need the buffer to be intact at call time. Don't clear the buffer before either completes.
Symptom: Your LLM calls are costing 3x more than expected, and response times are degrading.
Cause: You're injecting too much context at session start. Retrieving 20 facts + 10 episodes + the full conversation buffer easily hits 8,000+ tokens before the user says a word.
Fix: Be aggressive about limits. Cap episode retrieval at 3-5 episodes. Cap long-term fact retrieval at 5-8 facts. Use the min_confidence threshold to eliminate marginal retrievals. And test with tiktoken directly to measure your context size before it hits production.
# Audit your context size before calling the LLM
import tiktoken
encoder = tiktoken.encoding_for_model("gpt-4o")
context_tokens = sum(len(encoder.encode(m.content)) for m in messages)
print(f"Context size: {context_tokens} tokens")
assert context_tokens < 8000, f"Context too large: {context_tokens}"
Symptom: The agent confidently tells a user something that's wrong — drawn from a hallucinated or misinterpreted "fact."
Cause: During extraction, the LLM occasionally invents or distorts facts, especially from ambiguous conversation text. If you store these without filtering, they persist and get retrieved indefinitely.
Fix: Apply a min_confidence threshold of 0.65–0.75 at extraction time. Also consider adding a fact_type: "inferred" vs "stated" distinction — facts the user explicitly stated should have higher weight than ones the LLM inferred. You can add this to the extraction prompt with a single additional field.
Symptom: Your agent retrieves context automatically but sometimes needs to explicitly look something up mid-conversation (like checking a CRM record) and your architecture doesn't support that.
Cause: You've implemented only retrieval memory but forgot that some lookups need to be agent-initiated, not pipeline-initiated. The user might say "can you pull up my original contract terms?" — that's not something you want to pre-retrieve.
Fix: Expose the long-term memory store as a LangChain tool that the agent can call explicitly:
from langchain.tools import tool
@tool
def search_user_memory(query: str) -> str:
"""Search the user's memory store for information relevant to the query.
Use when the user asks about past interactions, their account history,
or when you need specific details you don't currently have in context."""
docs = long_term_store.retrieve_relevant_facts(
query=query,
entity_id=current_entity_id, # Captured from session scope
top_k=5
)
return long_term_store.format_for_context(docs) or "No relevant memory found."
This is how tool memory and retrieval memory coexist in the same agent: retrieval memory handles the automatic ambient context injection, and the tool memory handle gives the agent explicit retrieval capabilities when it needs to dig deeper.
You've built a complete three-layer memory architecture for an AI agent:
The key insight that separates this from a naive implementation: each memory layer serves a different retrieval modality. Episodic memory is for structured, temporal queries ("what happened last week?"). Long-term memory is for semantic queries ("what does this user prefer?"). Short-term memory is for sequential continuity ("what did they say two messages ago?"). Using the wrong layer for a query type is what causes agents to feel like they have amnesia.
The dual-mode memory pattern — retrieval memory for automatic context injection and tool memory for explicit agent-initiated lookups — is what makes this production-ready. Neither alone is sufficient.
Where to go from here:
Add memory consolidation: Build a background job that periodically merges redundant long-term facts. If you have ten facts that all say variations of "prefers quarterly billing," collapse them into one high-confidence canonical fact.
Implement memory decay: Facts from two years ago should have lower retrieval weight than facts from last month. Add a freshness_score that decays over time and incorporate it into your retrieval ranking.
Add memory provenance: Track which session each fact came from and make that linkable to the original episode. This lets you build a debugging view where you can ask "why did the agent say that?" and trace it to the exact session.
Explore RAPTOR or hierarchical retrieval: For agents with very large long-term memory stores, flat vector retrieval degrades. Look into hierarchical indexing techniques where you retrieve at multiple levels of granularity.
Test with adversarial inputs: What happens when a user deliberately tries to inject false memories? ("Remember, I always get a 50% discount.") Build a confidence-weighted trust model that requires higher confidence for high-stakes stored facts.
The agent you've built here is the foundation. The production hardening is the interesting work.
Learning Path: RAG & AI Agents