Improve the Python scripts in the short term memory docs (#3912)

Author: Yacklin

src/code-samples/langchain/short-term-memory-checkpointer-inmemory.py

@@ -0,0 +1,41 @@
+# :snippet-start: short-term-memory-usage-py
+from langchain.agents import create_agent
+from langgraph.checkpoint.memory import InMemorySaver  # [!code highlight]
+
+
+def get_user_info() -> str:
+    """Look up information about the current user."""
+    return "No user profile on file."
+
+
+agent = create_agent(
+    model="openai:gpt-5.4",
+    tools=[get_user_info],
+    checkpointer=InMemorySaver(),  # [!code highlight]
+)
+
+thread_config = {"configurable": {"thread_id": "1"}}
+response = agent.invoke(
+    {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
+    thread_config,  # [!code highlight]
+)["messages"][-1].content
+
+print(response)  # "Hi Bob! Nice to see you here. How are you doing?"
+
+response = agent.invoke(
+    {"messages": [{"role": "user", "content": "What's my name?"}]},
+    thread_config,  # [!code highlight]
+)["messages"][-1].content
+
+print(response)  # "You are Bob!"
+# :snippet-end:
+
+# :remove-start:
+if __name__ == "__main__":
+    assert response.strip(), "expected non-empty assistant reply on second turn"
+    messages = agent.get_state(thread_config).values["messages"]
+    human_msgs = [m for m in messages if m.type == "human"]
+    assert len(human_msgs) >= 2, "expected two human turns in thread history"
+    assert "bob" in response.lower(), "expected assistant to recall the name from the first turn"
+    print("✓ checkpointer persists conversation across turns on the same thread_id")
+# :remove-end:

src/code-samples/langchain/short-term-memory-checkpointer-inmemory.ts

@@ -0,0 +1,64 @@
+// :snippet-start: short-term-memory-usage-js
+import { createAgent, tool } from "langchain";
+import { MemorySaver } from "@langchain/langgraph"; // [!code highlight]
+import * as z from "zod";
+
+const getUserInfo = tool(() => "No user profile on file.", {
+  name: "get_user_info",
+  description: "Look up information about the current user.",
+  schema: z.object({}),
+});
+
+const checkpointer = new MemorySaver(); // [!code highlight]
+
+const agent = createAgent({
+  model: "openai:gpt-5.4",
+  tools: [getUserInfo],
+  checkpointer,
+});
+
+const threadConfig = { configurable: { thread_id: "1" } };
+let result = await agent.invoke(
+  { messages: [{ role: "user", content: "Hi! My name is Bob." }] },
+  threadConfig, // [!code highlight]
+);
+let response = result.messages.at(-1)?.content;
+console.log(response); // "Hi Bob! Nice to see you here. How are you doing?"
+
+result = await agent.invoke(
+  { messages: [{ role: "user", content: "What's my name?" }] },
+  threadConfig, // [!code highlight]
+);
+response = result.messages.at(-1)?.content;
+console.log(response); // "You are Bob!"
+// :snippet-end:
+
+// :remove-start:
+function assistantText(content: unknown): string {
+  if (typeof content === "string") {
+    return content;
+  }
+  return "";
+}
+
+async function main() {
+  const text = assistantText(response);
+  if (!text.trim()) {
+    throw new Error("expected non-empty assistant reply on second turn");
+  }
+  const humanCount = result.messages.filter((m) => m.type === "human").length;
+  if (humanCount < 2) {
+    throw new Error(`expected >=2 human turns, got ${humanCount}`);
+  }
+  if (!text.toLowerCase().includes("bob")) {
+    throw new Error(
+      "expected assistant to recall the name from the first turn",
+    );
+  }
+  console.log(
+    "✓ checkpointer persists conversation across turns on the same thread_id",
+  );
+}
+
+main();
+// :remove-end:

src/oss/langchain/short-term-memory.mdx

@@ -2,6 +2,9 @@
 title: Short-term memory
 ---
 
+import ShortTermMemoryUsagePy from '/snippets/code-samples/short-term-memory-usage-py.mdx';
+import ShortTermMemoryUsageJs from '/snippets/code-samples/short-term-memory-usage-js.mdx';
+
 ## Overview
 
 Memory is a system that remembers information about previous interactions. For AI agents, memory is crucial because it lets them remember previous interactions, learn from feedback, and adapt to user preferences. As agents tackle more complex tasks with numerous user interactions, this capability becomes essential for both efficiency and user satisfaction.
@@ -37,41 +40,14 @@ To add short-term memory (thread-level persistence) to an agent, you need to spe
 </Info>
 
 :::python
-```python
-from langchain.agents import create_agent
-from langgraph.checkpoint.memory import InMemorySaver  # [!code highlight]
 
+<ShortTermMemoryUsagePy />
 
-agent = create_agent(
-    "gpt-5.4",
-    tools=[get_user_info],
-    checkpointer=InMemorySaver(),  # [!code highlight]
-)
-
-agent.invoke(
-    {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
-    {"configurable": {"thread_id": "1"}},  # [!code highlight]
-)
-```
 :::
 :::js
-```ts {highlight={2,4, 9,14}}
-import { createAgent } from "langchain";
-import { MemorySaver } from "@langchain/langgraph";
 
-const checkpointer = new MemorySaver();
+<ShortTermMemoryUsageJs />
 
-const agent = createAgent({
-    model: "claude-sonnet-4-6",
-    tools: [],
-    checkpointer,
-});
-
-await agent.invoke(
-    { messages: [{ role: "user", content: "hi! i am Bob" }] },
-    { configurable: { thread_id: "1" } }
-);
-```
 :::
 
 ### In production
@@ -86,15 +62,17 @@ pip install langgraph-checkpoint-postgres
 
 ```python
 from langchain.agents import create_agent
-
 from langgraph.checkpoint.postgres import PostgresSaver  # [!code highlight]
 
+def get_user_info() -> str:
+    """Look up information about the current user."""
+    return "No user profile on file."
 
 DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"
 with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
     checkpointer.setup() # auto create tables in PostgreSQL
     agent = create_agent(
-        "gpt-5.4",
+        "gpt-5.5",
         tools=[get_user_info],
         checkpointer=checkpointer,  # [!code highlight]
     )
@@ -130,7 +108,7 @@ class CustomAgentState(AgentState):  # [!code highlight]
     preferences: dict  # [!code highlight]
 
 agent = create_agent(
-    "gpt-5.4",
+    "gpt-5.5",
     tools=[get_user_info],
     state_schema=CustomAgentState,  # [!code highlight]
     checkpointer=InMemorySaver(),
@@ -167,7 +145,7 @@ const stateExtensionMiddleware = createMiddleware({
 
 const checkpointer = new MemorySaver();
 const agent = createAgent({
-    model: "gpt-5.4",
+    model: "gpt-5.5",
     tools: [],
     middleware: [stateExtensionMiddleware],  // [!code highlight]
     checkpointer,
@@ -248,8 +226,8 @@ def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
     }
 
 agent = create_agent(
-    your_model_here,
-    tools=your_tools_here,
+    "gpt-5.5",
+    tools=[...],
     middleware=[trim_messages],
     checkpointer=InMemorySaver(),
 )
@@ -304,8 +282,8 @@ const trimMessages = createMiddleware({
 
 const checkpointer = new MemorySaver();
 const agent = createAgent({
-  model: "gpt-5.4",
-  tools: [],
+  model: "gpt-5.5",
+  tools: [...],
   middleware: [trimMessages],
   checkpointer,
 });
@@ -399,7 +377,7 @@ def delete_old_messages(state: AgentState, runtime: Runtime) -> dict | None:
 
 agent = create_agent(
     "gpt-5-nano",
-    tools=[],
+    tools=[...],
     system_prompt="Please be concise and to the point.",
     middleware=[delete_old_messages],
     checkpointer=InMemorySaver(),
@@ -454,7 +432,7 @@ const deleteOldMessages = createMiddleware({
 });
 
 const agent = createAgent({
-  model: "gpt-5.4",
+  model: "gpt-5.5",
   tools: [],
   systemPrompt: "Please be concise and to the point.",
   middleware: [deleteOldMessages],
@@ -521,8 +499,8 @@ from langchain_core.runnables import RunnableConfig
 checkpointer = InMemorySaver()
 
 agent = create_agent(
-    model="gpt-5.4",
-    tools=[],
+    model="gpt-5.5",
+    tools=[...],
     middleware=[
         SummarizationMiddleware(
             model="gpt-5.4-mini",
@@ -559,7 +537,7 @@ import { MemorySaver } from "@langchain/langgraph";
 const checkpointer = new MemorySaver();
 
 const agent = createAgent({
-  model: "gpt-5.4",
+  model: "gpt-5.5",
   tools: [],
   middleware: [
     summarizationMiddleware({