hands-on labin your browser

Add Long-Term Memory to an AI Agent: LangGraph + Milvus

Build a sales assistant that remembers: short-term state in a LangGraph checkpointer, long-term facts in Milvus, and reflection loops that extract knowledge.

time
35 min
checked steps
6
level
Intermediate
setup
None

$29.99a month with Pro

Pro opens all 6 steps of this lab and every other lab. Cancel any time.

no account needed · Pro opens all 6 steps

See Pro plans
Lab cockpit35 min · 6 stepsSession running
1 / 6 steps passingShort-Term Memory with Checkpointer · step 2 of 6
agent_checkpointed.py▶ Run✓ Check
# Step 2: Short-term memory with MemorySaver + thread_idfrom langchain_nvidia_ai_endpoints import ChatNVIDIAfrom langchain.agents import create_agent llm = ChatNVIDIA(    base_url="http://nim-proxy.labs.svc:8080/v1",    model="meta/llama-3.3-70b-instruct",    temperature=0,)    
TerminalOutput
step 2 of 6, replayed: Short-Term Memory with Checkpointer. The code, the output and the check are the lab's own.

The job

You're building an AI sales intelligence assistant for Nimbus Sales, a B2B SaaS with 50+ account executives. Every rep has dozens of active deals, talks to multiple buyers per week, and loses context between calls. Their biggest complaint: "I tell the AI about Acme on Monday, then Friday it has no idea who Acme is."

Your job: give the assistant two kinds of memory that production agents need.

By the end of this lab, you'll have:

ComponentWhat it doesPersistence
agent_no_memory.pyStateless baseline, shows the problemNone
agent_checkpointed.pyRemembers within a conversation threadIn-memory, per-session
memory_tools.pysave_memory + search_memory (Milvus-backed)Across sessions, permanent
agent_with_memory.pyCombines short + long-term memoryBoth
reflect.pyAuto-extracts salient facts from recent turnsPopulates long-term memory
sales_assistant.pyComplete system: checkpointer + Milvus + reflectionProduction-ready pattern

The test conversation that should work by end of lab:

TurnUser saysExpected behavior
1 (Mon)"I had a call with Acme Corp. VP Sarah is interested but worried about integration time."Responds helpfully, silently saves facts to memory
2 (Mon, same thread)"What did Sarah care about?"Recalls from thread state, "integration time"
3 (Fri, new thread)"Remind me about Acme"Retrieves from Milvus, "Sarah is VP, worried about integration time"
4 (Fri, same)"What was the objection?"Continues the new thread with long-term context

6 steps, each checked when you finish it

A check runs your work at the end of every step. Hints and the full solution are there if you get stuck.

  1. 1

    The Problem: A Stateless Agent

    The default LangGraph agent is stateless.

    read this step below ↓

  2. 2

    Short-Term Memory with Checkpointer

    The fix for step 1's problem is a checkpointer, a component that saves the graph state after each step, so the next invoke can resume from where you left off.

  3. 3

    Long-Term Memory with Milvus

    The checkpointer solves within-conversation memory.

  4. 4

    Memory-Aware Agent

    Now we wire the memory tools from step 3 into an agent.

  5. 5

    Reflection: Auto-Extracting Facts

    Relying on the agent to call save_memory at the right moment is unreliable.

  6. 6

    Putting It Together: The Sales Assistant

    You now have every piece of a production-grade memory architecture.

Step 1 as it appears in the lab

The lab’s own text. The hint and the solution stay inside the lab.

Step 1: The Problem: A Stateless Agent

Why memory matters

The default LangGraph agent is stateless. Each agent.invoke() call starts fresh — the agent has no knowledge of what was said before. For a sales assistant, this means:

Turn 1: "I just called Acme Corp — Sarah is the VP, she's worried about integration."
  → Agent: "Got it, I can help with integration concerns."

Turn 2: "What did Sarah care about?"
  → Agent: "I'm sorry, I don't have any information about Sarah. Can you share more context?"

Every conversation starts from zero. Unusable in production.

What we'll build across this lab

Production agents need two kinds of memory:

TypeScopeUse caseOur implementation
Short-termWithin one conversationRemember what we just discussedLangGraph MemorySaver + thread_id
Long-termAcross all conversationsRemember facts permanentlyMilvus vector store

For step 1, we'll build the broken baseline — an agent with no memory — to see the problem firsthand. Then each subsequent step adds a layer.

How create_agent works without a checkpointer

from langchain.agents import create_agent

agent = create_agent(model=llm, tools=[])

# Each call is independent — no shared state
agent.invoke({"messages": [{"role": "user", "content": "My name is Alice"}]})
agent.invoke({"messages": [{"role": "user", "content": "What is my name?"}]})
# Agent: "I don't know your name."

The fix in later steps will be a checkpointer argument that persists state keyed by thread_id.

Your task

Create agent_no_memory.py:

  1. Create a ChatNVIDIA LLM pointed at the NIM proxy (use meta/llama-3.3-70b-instruct)
  2. Build a create_agent without a checkpointer — just model=llm, tools=[]
  3. Run a 2-turn conversation demonstrating the problem:
    • Turn 1: Tell the agent about a client (name, role, concern)
    • Turn 2: Ask the agent to recall what the client cared about
  4. Print both responses so you can see the agent fails to remember

This sets up the baseline we'll improve in step 2.

NCP-AAI exam note: Cognition, Planning & Memory is 10% of the exam. Key concept: agents are stateless by default — all context must be explicitly passed in the messages list. Understanding when you need to add memory vs. pass context via tools is a core design decision.

agent_no_memory.py, the file you edit32 lines
# Step 1: The Stateless Baseline — demonstrates the problem
from langchain_nvidia_ai_endpoints import ChatNVIDIA
from langchain.agents import create_agent

# TODO: Create ChatNVIDIA LLM
#   base_url = "http://nim-proxy.labs.svc:8080/v1"
#   model    = "meta/llama-3.3-70b-instruct"
#   temperature = 0
llm = ______

# TODO: Build an agent with NO checkpointer — just model and empty tools list
agent = ______

# Turn 1 — tell the agent about a client
print("--- Turn 1 ---")
r1 = agent.invoke({"messages": [{
    "role": "user",
    "content": "I just got off a call with Acme Corp. Their VP, Sarah, is interested but worried about integration time."
}]})
print(f"User: I just got off a call with Acme Corp...")
print(f"Agent: {r1['messages'][-1].content[:200]}")

# Turn 2 — ask it to recall (it will fail)
print("\n--- Turn 2 ---")
r2 = agent.invoke({"messages": [{
    "role": "user",
    "content": "What was Sarah's main concern?"
}]})
print(f"User: What was Sarah's main concern?")
print(f"Agent: {r2['messages'][-1].content[:200]}")

print("\n(The agent has no memory of turn 1 — that's the problem we'll fix next.)")

what you get

  • A hosted environment in your browser
  • A check at the end of every step
  • Hints and the full solution
  • Nothing to install or configure

before you start

  • Completed Lab 1 (ReAct Agent with NIM) or equivalent
  • Completed Lab 2 (RAG Pipeline with NIM), we reuse Milvus for long-term memory
  • Basic LangGraph knowledge (create_agent, StateGraph, thread_id)

Exam domains covered

Cognition, Planning, and MemoryAgent Architecture and DesignAgent Development

Part of the AI agents and MCP track

Lab 12 of 26. Tool loops, ReAct, MCP servers and clients, multi-agent supervisors and agent evaluation.

Questions about this lab

Why use two different memory stores instead of one?

Because the access patterns are different. Short-term memory is exact and ordered, 'what were the last 12 messages in this thread?', and a checkpointer backed by a dict or Redis handles it in O(1) with zero embedding cost. Long-term memory is approximate and semantic, 'everything we know about Acme, retrieved by meaning', and needs a vector store with an embedding model. Stuffing everything into one store loses you either exactness (vector stores are lossy) or scalability (you cannot prompt-stuff a whole user history into every call). The production pattern is both, with a clear boundary.

What is thread_id in LangGraph and how does it relate to users?

thread_id is the key the checkpointer uses to load and save state for one conversation. Same thread_id = same conversation history; new thread_id = fresh context. It is not the same as a user ID, a single user typically has many threads (one per conversation, ticket, session) and long-term memory is keyed by user across all of them. The lab wires both: the checkpointer scopes by thread_id, save_memory stores facts in Milvus under the user's id so Friday's brand-new thread can still retrieve Monday's notes.

Why is reflection better than asking the LLM to call save_memory itself?

Because LLMs under conversational pressure forget to use side-channel tools. The agent is focused on answering the user's question; calling save_memory is a non-load-bearing step that the model skips about as often as not. Reflection decouples the save from the chat: after each turn a separate LLM pass reads the most recent few exchanges with a prompt like 'extract durable facts about clients, deals, or preferences mentioned above' and writes them to Milvus without the main agent ever having to notice. It's slower and more tokens but dramatically more reliable.

How do I prevent long-term memory from filling with garbage?

Three levers. First, the reflection prompt should explicitly distinguish durable facts ('Sarah is VP at Acme') from ephemeral context ('Sarah is out of office Thursday'), only the former should land in long-term memory. Second, tag every memory with a category and a confidence score, and surface both at retrieval so the agent can weight them. Third, add a periodic compaction pass that dedupes similar memories and expires low-confidence entries. The lab's save_memory takes a category argument specifically so you can filter by type at query time.

Why Milvus instead of pgvector or Chroma?

Milvus is the default in the NVIDIA NeMo Retriever reference architecture and runs at scale with approximate-nearest-neighbour indexes (IVF, HNSW, GPU-accelerated CAGRA) that pgvector does not ship natively. For a small personal assistant, pgvector or Chroma are completely fine and the API shape is very similar; the pattern you learn here, embed with NIM, upsert, cosine search, ports directly. The lab uses Milvus Lite (embedded) so the concepts transfer to a full Milvus cluster without changing the application code.

What exactly does the reflection step write?

A list of {fact, category} objects extracted from the last few turns by a structured-output LLM call. Typical categories the lab prompts for are client_context ('Sarah is VP at Acme'), deal_state ('Acme is evaluating the Pro tier'), objection ('Sarah is worried about integration time'), and preference ('Rep prefers SOC 2 one-pagers for security-sensitive buyers'). Each fact becomes a Milvus row with the category as metadata, which lets search_memory filter by type ('only pull objections for this account before the call') in addition to semantic match.

What you'll build in this agent memory lab

Agent memory is the feature that turns a chatbot into an assistant, and every production LLM app hits the wall the moment users complain 'I told it about Acme on Monday, Friday it has no idea who Acme is.' This lab assembles the two-tier memory architecture every serious agent team ships: short-term state via a LangGraph MemorySaver checkpointer keyed by thread_id, long-term semantic memory backed by Milvus, and a reflection loop that auto-extracts durable facts so the agent doesn't have to remember to save them. You finish with a sales-intelligence assistant that passes the Monday-call / Friday-recall test across fresh threads, running on NVIDIA NIM endpoints we provision.

The substance is the split between short-term (thread-scoped, exact message history, cheap) and long-term (user-scoped, semantic, permanent, cross-session). You start by proving the stateless baseline fails, then add a checkpointer keyed by thread_id so conversation state survives multiple invoke() calls. Long-term memory is save_memory(fact, category) writing embedded records into Milvus and search_memory(query, top_k) doing cosine similarity over NVIDIA NIM text-embedding vectors. The reflection loop, a separate LLM pass that reads recent turns and extracts anchor facts, is what makes it robust: relying on the agent to call save_memory mid-conversation is brittle because the model under user-facing pressure forgets side-channel tool calls about as often as not. This is the exact pattern real products ship.

Prerequisites: Python, the react-agent-nim and rag-pipeline-nim labs (you reuse Milvus), and comfort with create_agent, StateGraph, and thread_id. The hosted environment ships with LangGraph, the LangChain NIM integration, and pymilvus (Milvus Lite) preinstalled, running against our managed NIM proxy serving meta/llama-3.3-70b-instruct and the NVIDIA text-embedding model, no keys, no GPU provisioning. About 35 minutes of focused work. You leave with a stateless baseline that demonstrates the failure mode, a thread-scoped agent that recalls within a conversation, a Milvus-backed long-term store, a reflection pass that extracts durable facts, and a composed assistant that holds the whole thing together.