Tool Memory vs. Retrieval Memory in AI Agents: Implementing Short-Term, Long-Term, and Episodic Memory Architectures

Introduction

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:

• The conceptual distinction between tool memory and retrieval memory, and why it matters architecturally
• How to implement short-term (in-context) memory using conversation buffer management
• How to build long-term memory with a vector store and semantic retrieval
• How to design and populate an episodic memory system that stores structured summaries of past interactions
• How to wire all three memory types into a single agent that selects the right memory at the right time

---

Prerequisites

You should be comfortable with:

• Python at an intermediate level (async functions, dataclasses, type hints)
• Basic LangChain or LlamaIndex concepts (chains, agents, retrievers) — we'll use LangChain here
• Familiarity with embeddings and vector stores (ChromaDB or similar)
• REST API basics and JSON manipulation

You'll need: langchain, langchain-openai, chromadb, tiktoken, python-dotenv, sqlalchemy

pip install langchain langchain-openai chromadb tiktoken python-dotenv sqlalchemy

---

The Conceptual Split: Tool Memory vs. Retrieval Memory

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.

---

Memory Layer 1: Short-Term Memory (In-Context Buffer)

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 Buffer Management Problem

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:

1. Fixed window: Keep the last N messages. Simple but loses crucial early context.
2. Token-capped window: Keep messages until you hit a token budget. More precise than message count.
3. Summarizing buffer: When the buffer exceeds a threshold, summarize older messages and replace them with the summary.

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.

Testing the Buffer

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]}")

---

Memory Layer 2: Long-Term Memory (Vector Store Retrieval)

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.

Structuring What You Store

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}"

Building the Long-Term Memory Store

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)

Populating Long-Term Memory (The Extraction Step)

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.

---

Memory Layer 3: Episodic Memory (Structured Session History)

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.

The Episodic Memory Schema

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)

Generating Episodic Records

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

---

Wiring It Together: The Multi-Memory Agent

Now we build the actual agent that uses all three memory layers together. The key design decision here is when each layer fires:

• Episodic memory retrieves at session start (before the first message is processed) — you want the agent to know about unresolved past issues immediately.
• Long-term memory retrieves at session start and on demand when the agent recognizes it needs background knowledge.
• Short-term memory is active throughout the session, always present in the context window.

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()
        }

---

Hands-On Exercise: Building the Full Pipeline

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:

1. Know Priya's name and company without being told
2. Know about the duplicate charge from session 1
3. Know that TechCorp prefers quarterly billing
4. Connect the Q1 seats mention to the prior discussion

If it doesn't, the retrieval or extraction step has a problem — which leads us to troubleshooting.

---

Common Mistakes & Troubleshooting

Mistake 1: Embedding the Wrong Text

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.

---

Mistake 2: No Entity Isolation in Retrieval

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.

---

Mistake 3: Storing Episodes Before Extraction

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.

---

Mistake 4: Token Explosion from Over-Injection

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}"

---

Mistake 5: Treating Low-Confidence LLM Extractions as Ground Truth

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.

---

Mistake 6: Missing the Tool Memory Use Case

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.

---

Summary & Next Steps

You've built a complete three-layer memory architecture for an AI agent:

• Short-term memory (ConversationBuffer): token-aware, compressing, handles in-session continuity
• Long-term memory (LongTermMemoryStore): vector-backed, semantic retrieval of distilled facts that persist across sessions
• Episodic memory (EpisodicMemoryStore): structured SQL records of past sessions, queryable by date, status, and category

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:

1. 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.

2. 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.

3. 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.

4. 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.

5. 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.