langgraph_memory_saver_checkpointer_state_persistence_quickstart.py

1from typing import Annotated
2from typing_extensions import TypedDict
3from langgraph.graph import StateGraph, START, END
4from langgraph.graph.message import add_messages
5from langgraph.checkpoint.memory import MemorySaver
6
7# 1. Define the state of our graph
8class State(TypedDict):
9    # add_messages allows us to append new messages to the history
10    messages: Annotated[list, add_messages]
11
12# 2. Define a simple node
13def chatbot(state: State):
14    return {"messages": [("assistant", "Hello! How can I help you today?")]}
15
16# 3. Initialize the graph with a checkpointer
17workflow = StateGraph(State)
18workflow.add_node("chatbot", chatbot)
19workflow.add_edge(START, "chatbot")
20workflow.add_edge("chatbot", END)
21
22# MemorySaver is the default in-memory checkpointer provided by langgraph-checkpoint
23memory = MemorySaver()
24
25# Compile the graph with the checkpointer
26app = workflow.compile(checkpointer=memory)
27
28# 4. Use the graph with a thread_id to persist state
29config = {"configurable": {"thread_id": "1"}}
30
31# First interaction
32input_message = {"messages": [("user", "Hi, I'm Alice")]}
33for event in app.stream(input_message, config):
34    for value in event.values():
35        print("Assistant:", value["messages"][-1].content)
36
37# The state is now saved under thread_id "1". 
38# Subsequent calls with the same thread_id will load the previous state.
39snapshot = app.get_state(config)
40print(f"Current State Messages Count: {len(snapshot.values['messages'])}")