How to share state between threads¶
By default, state in a graph is scoped to that thread. LangGraph also allows you to specify a "scope" for a given key/value pair that exists between threads. This can be useful for storing information that is shared between threads. For instance, you may want to store information about a user's preferences expressed in one thread, and then use that information in another thread.
In this notebook we will go through an example of how to construct and use such a graph.
Setup¶
First, let's install the required packages and set our API keys
npm install @langchain/openai @langchain/langgraph @langchain/core zod uuid
Then set your enviroment variables for OpenAI:
process.env.OPENAI_API_KEY = "your-openai-api-key";
Create graph¶
In this example we will create a graph that will let us store information about a user's preferences. We will do so by defining a state key that will be scoped to a user_id
, and allowing the model to populate this field as it deems fit (by providing the model with a tool to save information about the user).
Typing shared state keys
Shared state channels (keys) MUST be objects (see info
channel in the AgentState example below)
import { z } from "zod";
import {
START,
END,
Annotation,
StateGraph,
MemoryStore,
SharedValue,
MemorySaver,
} from "@langchain/langgraph";
import { ChatOpenAI } from "@langchain/openai";
import {
type AIMessage,
type BaseMessage,
ToolMessage,
} from "@langchain/core/messages";
import { RunnableConfig } from "@langchain/core/runnables";
import { v4 as uuidv4 } from "uuid";
const infoSchema = z.object({
fact: z.string().describe("The fact about the user"),
topic: z.string().describe("The topic of the fact"),
});
const AgentAnnotation = Annotation.Root({
messages: Annotation<BaseMessage[]>({
reducer: (a, b) => a.concat(b),
default: () => [],
}),
// IMPORTANT
// This is how you define a shared state value.
// The string passed to `.on` is the key that will
// be used to store the value in the shared state.
info: SharedValue.on("user_id"),
});
const prompt = `You are helpful assistant.
Here is what you know about the user:
<info>
{info}
</info>
Help out the user. If the user tells you any information about themselves, save the information using the \`Info\` tool.
This means if the user provides any sort of fact about themselves, be it an opinion they have, a fact about themselves, etc. SAVE IT!
`;
const model = new ChatOpenAI({
model: "gpt-4o",
temperature: 0,
}).bindTools([
{
name: "Info",
description: "Save the information provided by the user",
schema: infoSchema,
},
]);
const callModel = async (
state: typeof AgentAnnotation.State
): Promise<Partial<typeof AgentAnnotation.State>> => {
const facts = Object.values(state.info).map((d) => d.fact);
const info = facts.join("\n");
const systemMsg = prompt.replace("{info}", info);
const response = await model.invoke([
{ role: "system", content: systemMsg },
...state.messages,
]);
return { messages: [response] };
};
const route = (state: typeof AgentAnnotation.State): string => {
const lastMessage = state.messages[state.messages.length - 1];
if (!("tool_calls" in lastMessage)) {
throw new Error("Expected an AI message with tool calls.");
}
return (lastMessage as AIMessage).tool_calls?.length ? "update_memory" : END;
};
const updateMemory = (
state: typeof AgentAnnotation.State
): Partial<typeof AgentAnnotation.State> => {
const toolResponseMessages: ToolMessage[] = [];
const memories: Record<string, z.infer<typeof infoSchema>> = {};
const lastMessage = state.messages[state.messages.length - 1];
if (!("tool_calls" in lastMessage)) {
throw new Error("Expected an AI message with tool calls.");
}
const castLastMessage = lastMessage as AIMessage;
castLastMessage.tool_calls?.forEach((tc) => {
toolResponseMessages.push(
new ToolMessage({
content: "Saved!",
tool_call_id: tc.id as string,
})
);
memories[uuidv4()] = {
fact: tc.args.fact,
topic: tc.args.topic,
};
});
return { messages: toolResponseMessages, info: memories };
};
const memory = new MemorySaver();
// IMPORTANT
// In order to use shared values, you must initialize a store like this:
const kv = new MemoryStore();
const graph = new StateGraph(AgentAnnotation)
.addNode("call_model", callModel)
.addNode("update_memory", updateMemory)
.addEdge("update_memory", END)
.addEdge(START, "call_model")
.addConditionalEdges("call_model", route);
const compiledGraph = graph.compile({
checkpointer: memory,
// Then, pass it to `.compile` like this:
store: kv,
});
Run graph on one thread¶
We can now run the graph on one thread and give it some information
const config = {
configurable: {
thread_id: "1",
// Notice we're specifying `user_id` here, which matches the key name we passed to `SharedValue.on()`
// Without this, our graph wouldn't be able to access the shared state value.
user_id: "1"
},
streamMode: "updates" as const
};
// First let's just say hi to the AI
for await (const update of await compiledGraph.stream({
messages: [{ role: "user", content: "hi" }],
}, config)) {
console.log(update);
}
// Let's continue the conversation (by passing the same config) and tell the AI we like pepperoni pizza
for await (const update of await compiledGraph.stream({
messages: [{ role: "user", content: "i like pepperoni pizza" }],
}, config)) {
console.log(update);
}
// Let's continue the conversation even further (by passing the same config) and tell the AI we live in SF
for await (const update of await compiledGraph.stream({
messages: [{ role: "user", content: "i also just moved to SF" }],
}, config)) {
console.log(update);
}
{ call_model: { messages: [ AIMessage { "id": "chatcmpl-A8VfAQGfZZ4WBbG4zrbwzc34k1Ji3", "content": "Hello! How can I assist you today?", "additional_kwargs": {}, "response_metadata": { "tokenUsage": { "completionTokens": 10, "promptTokens": 137, "totalTokens": 147 }, "finish_reason": "stop", "system_fingerprint": "fp_a5d11b2ef2" }, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": { "input_tokens": 137, "output_tokens": 10, "total_tokens": 147 } } ] } } { call_model: { messages: [ AIMessage { "id": "chatcmpl-A8VfAh46qQIvHejVESnvqYwnAIor6", "content": "", "additional_kwargs": { "tool_calls": [ { "id": "call_IbF0aL78Xep9Xpz3UHLn0POR", "type": "function", "function": "[Object]" } ] }, "response_metadata": { "tokenUsage": { "completionTokens": 23, "promptTokens": 159, "totalTokens": 182 }, "finish_reason": "tool_calls", "system_fingerprint": "fp_25624ae3a5" }, "tool_calls": [ { "name": "Info", "args": { "fact": "The user likes pepperoni pizza", "topic": "Food Preferences" }, "type": "tool_call", "id": "call_IbF0aL78Xep9Xpz3UHLn0POR" } ], "invalid_tool_calls": [], "usage_metadata": { "input_tokens": 159, "output_tokens": 23, "total_tokens": 182 } } ] } } { update_memory: { messages: [ ToolMessage { "content": "Saved!", "additional_kwargs": {}, "response_metadata": {}, "tool_call_id": "call_IbF0aL78Xep9Xpz3UHLn0POR" } ], info: { '2bf3fed4-028a-4c86-93fe-717cde64e8e7': [Object] } } } { call_model: { messages: [ AIMessage { "id": "chatcmpl-A8VfBbqFPyFfN65f4Tk3k0kdtXHtJ", "content": "", "additional_kwargs": { "tool_calls": [ { "id": "call_ZkpyxLljOOfhp4oAMyO2R1Lh", "type": "function", "function": "[Object]" } ] }, "response_metadata": { "tokenUsage": { "completionTokens": 23, "promptTokens": 208, "totalTokens": 231 }, "finish_reason": "tool_calls", "system_fingerprint": "fp_a5d11b2ef2" }, "tool_calls": [ { "name": "Info", "args": { "fact": "The user just moved to San Francisco", "topic": "Location" }, "type": "tool_call", "id": "call_ZkpyxLljOOfhp4oAMyO2R1Lh" } ], "invalid_tool_calls": [], "usage_metadata": { "input_tokens": 208, "output_tokens": 23, "total_tokens": 231 } } ] } } { update_memory: { messages: [ ToolMessage { "content": "Saved!", "additional_kwargs": {}, "response_metadata": {}, "tool_call_id": "call_ZkpyxLljOOfhp4oAMyO2R1Lh" } ], info: { 'ee0085c3-0396-4be8-b0cd-869194aacd5b': [Object] } } }
Run graph on a different thread¶
We can now run the graph on a different thread and see that it remembers facts about the user (specifically that the user likes pepperoni pizza and lives in SF):
const config2 = {
configurable: {
// Notice we have a new thread ID, but the same user ID.
// This allows us to access the shared state value.
thread_id: "2",
user_id: "1"
},
streamMode: "updates" as const
};
for await (const update of await compiledGraph.stream({
messages: [{ role: "user", content: "where and what should i eat for dinner? Can you list some restaurants?" }],
}, config2)) {
console.log(update);
}
{ call_model: { messages: [ AIMessage { "id": "chatcmpl-A8VgEKTv4D2VzGdGwq6FVv09HczZf", "content": "Sure! Since you just moved to San Francisco, here are some popular restaurants you might enjoy:\n\n1. **Tony's Pizza Napoletana**\n - **Cuisine:** Italian, Pizza\n - **Location:** 1570 Stockton St, San Francisco, CA 94133\n - **Why you might like it:** They have a great selection of pizzas, including pepperoni!\n\n2. **House of Prime Rib**\n - **Cuisine:** American, Steakhouse\n - **Location:** 1906 Van Ness Ave, San Francisco, CA 94109\n - **Why you might like it:** If you're in the mood for a hearty meal, their prime rib is highly recommended.\n\n3. **Tartine Bakery**\n - **Cuisine:** Bakery, Cafe\n - **Location:** 600 Guerrero St, San Francisco, CA 94110\n - **Why you might like it:** Perfect for a lighter meal or dessert, their pastries are famous.\n\n4. **La Taqueria**\n - **Cuisine:** Mexican\n - **Location:** 2889 Mission St, San Francisco, CA 94110\n - **Why you might like it:** Known for their delicious tacos and burritos.\n\n5. **Swan Oyster Depot**\n - **Cuisine:** Seafood\n - **Location:** 1517 Polk St, San Francisco, CA 94109\n - **Why you might like it:** Great spot for fresh seafood.\n\nWould you like more information on any of these, or do you have a specific type of cuisine in mind?", "additional_kwargs": {}, "response_metadata": { "tokenUsage": { "completionTokens": 324, "promptTokens": 166, "totalTokens": 490 }, "finish_reason": "stop", "system_fingerprint": "fp_25624ae3a5" }, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": { "input_tokens": 166, "output_tokens": 324, "total_tokens": 490 } } ] } }
Perfect! The AI recommended restaurants in SF, and included a pizza restaurant at the top of it's list.
Notice that the messages
in this new thread do NOT contain the messages from the previous thread since we didn't store them as shared values across the user_id
. However, the info
we saved in the previous thread was saved since we passed in the same user_id
in this new thread.
Let's now run the graph for another user to verify that the preferences of the first user are self contained:
// Once again, we're specifying a new `user_id` value here.
// Like the previous examples, this means the graph will not
// be able to access the memory saved from the previous run.
const config3 = {
configurable: {
thread_id: "3",
user_id: "2"
},
streamMode: "updates" as const
}
for await (const update of await compiledGraph.stream({
messages: [{ role: "user", content: "where and what should i eat for dinner? Can you list some restaurants?" }],
}, config3)) {
console.log(update);
}
{ call_model: { messages: [ AIMessage { "id": "chatcmpl-A8VgwX6MUbLwBdistYNC0LP1t6y7S", "content": "Sure, I can help with that! To give you the best recommendations, could you please tell me your location or the city you're in? Additionally, do you have any preferences or dietary restrictions?", "additional_kwargs": {}, "response_metadata": { "tokenUsage": { "completionTokens": 40, "promptTokens": 151, "totalTokens": 191 }, "finish_reason": "stop", "system_fingerprint": "fp_a5d11b2ef2" }, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": { "input_tokens": 151, "output_tokens": 40, "total_tokens": 191 } } ] } }
Perfect! The graph has forgotten all of the previous preferences and has to ask the user for it's location and dietary preferences.