langgraph_memorysaver_checkpoint_thread_state_persistence.py

1from typing import Annotated
2from typing_extensions import TypedDict
3
4from langgraph.graph import StateGraph, START, END
5from langgraph.graph.message import add_messages
6from langgraph.checkpoint.memory import MemorySaver
7
8# 1. Define the state of the graph
9class State(TypedDict):
10    # add_messages allows us to append new messages to the existing list
11    messages: Annotated[list, add_messages]
12
13# 2. Define a simple node
14def chatbot(state: State):
15    return {"messages": [("assistant", "Hello! How can I help you today?")]}
16
17# 3. Build the graph
18builder = StateGraph(State)
19builder.add_node("chatbot", chatbot)
20builder.add_edge(START, "chatbot")
21builder.add_edge("chatbot", END)
22
23# 4. Initialize the checkpointer
24memory = MemorySaver()
25
26# 5. Compile the graph with the checkpointer
27graph = builder.compile(checkpointer=memory)
28
29# 6. Use the graph with a config containing a thread_id
30config = {"configurable": {"thread_id": "1"}}
31
32# First interaction
33input_message = {"messages": [("user", "Hi there!")]}
34for event in graph.stream(input_message, config):
35    for value in event.values():
36        print("Assistant:", value["messages"][-1].content)
37
38# The state is now persisted. You can resume the conversation using the same thread_id.
39print("\n--- Resuming Conversation ---")
40input_message_2 = {"messages": [("user", "Remember my name is Alice.")]}
41for event in graph.stream(input_message_2, config):
42    for value in event.values():
43        print("Assistant:", value["messages"][-1].content)