From 805059d3f1035c225308ca38ce3ecd118408b480 Mon Sep 17 00:00:00 2001
From: Lander McFall Kerbey <lander.kerbey@mongodb.com>
Date: Thu, 13 Aug 2026 12:21:06 -0400
Subject: [PATCH] (DOCSP-62238) Updating notebook for Semantic Memory

---
 .../langchain-memory-semantic-cache.ipynb     | 68 +++++++++++--------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/ai-integrations/langchain-memory-semantic-cache.ipynb b/ai-integrations/langchain-memory-semantic-cache.ipynb
index 7c76090..21a8254 100644
--- a/ai-integrations/langchain-memory-semantic-cache.ipynb
+++ b/ai-integrations/langchain-memory-semantic-cache.ipynb
@@ -31,7 +31,7 @@
    },
    "outputs": [],
    "source": [
-    "pip install --quiet --upgrade langchain langchain-community langchain-core langchain-mongodb langchain-voyageai langchain-openai langchain-text-splitters pypdf"
+    "pip install --quiet --upgrade langchain langchain-community langchain-core langchain-mongodb langchain-voyageai langchain-openai langchain-text-splitters langgraph langgraph-checkpoint-mongodb pypdf"
    ]
   },
   {
@@ -141,18 +141,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory\n",
-    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
+    "from langgraph.checkpoint.mongodb import MongoDBSaver\n",
     "from langchain_core.prompts import MessagesPlaceholder\n",
-    "         \n",
-    "# Define a function that gets the chat message history \n",
-    "def get_session_history(session_id: str) -> MongoDBChatMessageHistory:\n",
-    "    return MongoDBChatMessageHistory(\n",
-    "        connection_string=MONGODB_URI,\n",
-    "        session_id=session_id,\n",
-    "        database_name=\"langchain_db\",\n",
-    "        collection_name=\"rag_with_memory\"\n",
-    "    )"
+    "from pymongo import MongoClient\n",
+    "\n",
+    "# Connect to your MongoDB cluster\n",
+    "mongo_client = MongoClient(MONGODB_URI)\n",
+    "\n",
+    "# Create the checkpointer to persist conversation state\n",
+    "checkpointer = MongoDBSaver(mongo_client, db_name=\"langchain_db\")"
    ]
   },
   {
@@ -229,6 +226,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from langgraph.graph import StateGraph, START, END, MessagesState\n",
+    "from langchain_core.messages import AIMessage\n",
+    "\n",
     "# Build the RAG chain\n",
     "rag_chain = (\n",
     "    retriever_chain\n",
@@ -237,13 +237,23 @@
     "    | parse_output\n",
     ")\n",
     "\n",
-    "# Wrap the chain with message history\n",
-    "rag_with_memory = RunnableWithMessageHistory(\n",
-    "    rag_chain,\n",
-    "    get_session_history,\n",
-    "    input_messages_key=\"question\",\n",
-    "    history_messages_key=\"history\",\n",
-    ")"
+    "# Define the node that runs the RAG chain\n",
+    "def call_rag_chain(state: MessagesState):\n",
+    "    question = state[\"messages\"][-1].content\n",
+    "    history = state[\"messages\"][:-1]\n",
+    "    answer = rag_chain.invoke(\n",
+    "        {\"question\": question, \"history\": history}\n",
+    "    )\n",
+    "    return {\"messages\": [AIMessage(content=answer)]}\n",
+    "\n",
+    "# Build the graph\n",
+    "workflow = StateGraph(MessagesState)\n",
+    "workflow.add_node(\"rag\", call_rag_chain)\n",
+    "workflow.add_edge(START, \"rag\")\n",
+    "workflow.add_edge(\"rag\", END)\n",
+    "\n",
+    "# Compile the graph with the checkpointer\n",
+    "rag_with_memory = workflow.compile(checkpointer=checkpointer)"
    ]
   },
   {
@@ -255,10 +265,10 @@
    "source": [
     "# First question\n",
     "response_1 = rag_with_memory.invoke(\n",
-    "    {\"question\": \"What was MongoDB's latest acquisition?\"},\n",
-    "    {\"configurable\": {\"session_id\": \"user_1\"}}\n",
+    "    {\"messages\": [(\"human\", \"What was MongoDB's latest acquisition?\")]},\n",
+    "    {\"configurable\": {\"thread_id\": \"user_1\"}}\n",
     ")\n",
-    "print(response_1)"
+    "print(response_1[\"messages\"][-1].content)"
    ]
   },
   {
@@ -270,10 +280,10 @@
    "source": [
     "# Follow-up question that references the previous question\n",
     "response_2 = rag_with_memory.invoke(\n",
-    "    {\"question\": \"Why did they do it?\"},\n",
-    "    {\"configurable\": {\"session_id\": \"user_1\"}}\n",
+    "    {\"messages\": [(\"human\", \"Why did they do it?\")]},\n",
+    "    {\"configurable\": {\"thread_id\": \"user_1\"}}\n",
     ")\n",
-    "print(response_2)"
+    "print(response_2[\"messages\"][-1].content)"
    ]
   },
   {
@@ -320,8 +330,8 @@
     "\n",
     "# First query (not cached)\n",
     "rag_with_memory.invoke(\n",
-    "  {\"question\": \"What was MongoDB's latest acquisition?\"},\n",
-    "  {\"configurable\": {\"session_id\": \"user_2\"}}\n",
+    "  {\"messages\": [(\"human\", \"What was MongoDB's latest acquisition?\")]},\n",
+    "  {\"configurable\": {\"thread_id\": \"user_2\"}}\n",
     ")"
    ]
   },
@@ -336,8 +346,8 @@
     "\n",
     "# Second query (cached)\n",
     "rag_with_memory.invoke(\n",
-    "  {\"question\": \"What company did MongoDB acquire recently?\"},\n",
-    "  {\"configurable\": {\"session_id\": \"user_2\"}}\n",
+    "  {\"messages\": [(\"human\", \"What company did MongoDB acquire recently?\")]},\n",
+    "  {\"configurable\": {\"thread_id\": \"user_2\"}}\n",
     ")"
    ]
   }