Skip to content

Commit 8bfd74e

Browse files
author
Rtudo
committed
Add advance-state-management
1 parent 2f0e8e5 commit 8bfd74e

4 files changed

Lines changed: 467 additions & 82 deletions

File tree

‎docs/ai/langgraph/advance-state-management.md‎

Lines changed: 381 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,381 @@
1+
# Advanced State Management: Beyond Just Messages
2+
3+
So far, our agent's state has been a simple list of messages. LangGraph's `add_messages` helper has cleverly managed the conversation history for us. But what happens when we need to track more than just the chat? What if we need to store structured data, track approvals, or manage complex application states?
4+
5+
This is where LangGraph's state management truly shines. The state `TypedDict` can hold anything you want, giving you a powerful way to control the agent's flow. In this tutorial, we'll explore two advanced techniques:
6+
7+
1. **Customizing the state** to manage an approval workflow.
8+
2. **Time travel**, allowing us to rewind the agent's state to correct its course.
9+
10+
## Customizing and Updating State: An Approval Workflow
11+
12+
Imagine an agent designed to help with software deployments. It needs to fetch deployment details, present them for human review, and only proceed if a human approver gives the green light. A simple list of messages won't be enough to track the approval status.
13+
14+
We need a richer state object.
15+
16+
### 1. Defining a Custom State
17+
18+
Let's create a new `State` dictionary that includes not only our messages but also fields for the deployment information and the approval status.
19+
20+
```python
21+
from typing import TypedDict, Annotated, Optional
22+
from langchain_core.messages import BaseMessage
23+
24+
# We can define a structure for our deployment info
25+
class DeploymentInfo(TypedDict):
26+
build_number: int
27+
changelog: str
28+
deployed_by: Optional[str]
29+
deployed_at: Optional[str]
30+
31+
class State(TypedDict):
32+
messages: Annotated[list, add_messages]
33+
# A new field to hold our structured data
34+
deployment_info: Optional[DeploymentInfo]
35+
# A simple flag to track the approval status
36+
approved: bool
37+
```
38+
Our state now has a dedicated place for `deployment_info` and a boolean `approved` flag, which we'll initialize to `False`.
39+
40+
### 2. The Approval Flow
41+
42+
Our new workflow will look like this:
43+
44+
1. Fetch deployment information and add it to the state.
45+
2. Interrupt and present this information to a human.
46+
3. The human either approves or denies the deployment.
47+
4. Based on the approval, the graph will branch to either a "deploy" node or the end.
48+
49+
50+
```mermaid
51+
graph LR
52+
A[Start] --> B(Fetch Deployment Info);
53+
B --> C((PAUSE: Human Review));
54+
C --> D{Approved?};
55+
D -- Yes --> E[Deploy Service Node];
56+
D -- No --> F[End];
57+
E --> F;
58+
```
59+
60+
### 3. Updating the State from Outside the Graph
61+
62+
The key to this workflow is updating the state *after* the graph has been interrupted. When the human reviewer gives their approval, our application code will directly modify the graph's state using `graph.update_state()`.
63+
64+
Here’s how we can implement this:
65+
66+
1. **New Nodes**: We'll define a `fetch_deployment_info` node that simulates getting data and a `deploy_service` node that runs after approval.
67+
2. **Conditional Edge**: We'll create a new function `check_approval_status` that inspects `state['approved']` to decide the next step.
68+
3. **Application Logic**: The main script will get user input and call `graph.update_state()` to change the `approved` flag from `False` to `True`.
69+
70+
Here's a look at the application logic that handles the human interaction:
71+
72+
```python
73+
# The graph will interrupt before the 'check_approval_status' edge
74+
# ... graph execution starts ...
75+
paused_state = graph.get_state(config)
76+
info = paused_state.values['deployment_info']
77+
print("--- Deployment Review Required ---")
78+
print(f"Build Number: {info['build_number']}")
79+
print(f"Changelog: {info['changelog']}")
80+
81+
approval = input("Approve deployment? (yes/no): ").lower()
82+
83+
if approval == "yes":
84+
# If approved, UPDATE THE STATE
85+
update = {"approved": True}
86+
graph.update_state(config, update)
87+
# Resume the graph
88+
graph.invoke(None, config)
89+
else:
90+
print("Deployment aborted.")
91+
```
92+
This pattern is incredibly powerful. It allows the graph to handle the automated parts of a workflow while giving the user precise control over key decision points.
93+
94+
---
95+
96+
## Customizing and Resetting State: The Time Travel Approach
97+
98+
99+
### Scenario: The random number generation
100+
101+
To make the concept of time travel crystal clear, we'll use a simple, deterministic task. Imagine an agent designed to follow a multi-step plan:
102+
103+
1. The user asks the agent to generate **three random numbers**, one at a time.
104+
2. The agent calls a `generate_random_number` tool three separate times, adding each result to the conversation state.
105+
3. After generating all three numbers, the agent presents the final list to the user.
106+
4. The graph **pauses** for human review.
107+
5. The user can either tell the agent to **rewind** to a previous step (e.g., "start over from the 2nd number") or proceed with the generated numbers.
108+
6. If a rewind is requested, the application code will find the correct checkpoint in the graph's history, roll back the state, and let the agent continue from that exact point.
109+
110+
### Visualizing the Time Travel Workflow
111+
112+
<!-- ??? "Time Travel Flow" -->
113+
``` mermaid
114+
flowchart LR
115+
A[Start] --> B{"Agent Logic"}
116+
B -->|Needs Tool 1/3| C["Call Tool"]
117+
B -->|Needs Tool 2/3| C
118+
B -->|Needs Tool 3/3| C
119+
C -->|Tool Result|B
120+
C -->|Tool Result|B
121+
C --> D["Present for Review"]
122+
```
123+
124+
125+
``` mermaid
126+
flowchart LR
127+
B{"Agent Logic"}
128+
I["Resume from Checkpoint"]
129+
G["End"]
130+
H["Find Checkpoint N in History"]
131+
E(("Pause: Human Review")) --> F{"Decision"}
132+
F --> G
133+
F --> H
134+
H --> I
135+
I --> B
136+
```
137+
138+
139+
### The Key Ingredients for Time Travel
140+
141+
This advanced pattern relies on a few core LangGraph features working together.
142+
143+
!!! info "A Persistent Checkpointer is Essential"
144+
Time travel is only possible because a **checkpointer** saves the state of the graph after every single step. For this to work, you *must* compile your graph with a checkpointer like `MemorySaver` or, for production apps, `SqliteSaver`. It's this saved history that we will navigate.
145+
146+
Here are the key components of our implementation:
147+
148+
<div class="grid cards" markdown>
149+
150+
- **Multi-Step Task**
151+
Our prompt explicitly asks the agent to generate three numbers one at a time. This forces the agent to call its tool multiple times, creating a distinct history of checkpoints that we can later choose from.
152+
153+
- **`graph.get_state_history()`**
154+
This is our time machine. After the graph pauses, we can call this method to retrieve a complete, ordered list of every state the graph has been in for the current conversation thread.
155+
156+
- **Application-Side Logic**
157+
The graph itself doesn't contain the "rewind" logic. The intelligence to inspect the history, select a checkpoint, and resume the graph lives in our Python application. This separates the agent's automated workflow from the user's manual control.
158+
159+
</div>
160+
161+
### Building the Time-Traveling Agent
162+
163+
164+
#### 1. The Tool and State
165+
166+
First, we define our simple tool and the graph's state. The state itself remains a simple list of messages; the complexity lies in how we manage its history.
167+
168+
```python
169+
import random
170+
from langchain_core.tools import tool
171+
from typing import TypedDict, Annotated
172+
from langgraph.graph.message import add_messages
173+
174+
# A simple tool that generates a number
175+
@tool
176+
def generate_random_number(min_val: int = 1, max_val: int = 100) -> int:
177+
"""Generates a random number within a specified range."""
178+
print(f"--- TOOL CALL: Generating a random number ---")
179+
return random.randint(min_val, max_val)
180+
181+
# The state remains simple
182+
class State(TypedDict):
183+
messages: Annotated[list, add_messages]
184+
```
185+
186+
#### 2. The Time Travel Cockpit
187+
188+
The core of our solution is the application logic that runs after the graph pauses/interrupts for the human review:
189+
190+
1. Getting the current state of the paused graph.
191+
2. Presenting the agent's work to the user.
192+
3. Handling user feedback (a number to rewind to state before that number was generated).
193+
4. **Executing the time travel:**
194+
* Fetching the entire history of checkpoints.
195+
* Finding the specific checkpoint to rewind to by counting the `ToolMessage` instances.
196+
* Creating a new configuration and resuming the graph's execution from that past state.
197+
198+
Let's look at the logic for finding the right checkpoint:
199+
200+
```python
201+
# 'choice' is the step the user wants to rewind to (e.g., 2)
202+
current_history = list(graph.get_state_history(config))
203+
204+
# We iterate through the checkpoints to find the one
205+
# that occurred just *before* the chosen tool call.
206+
target_checkpoint = None
207+
for state in reversed(current_history):
208+
tool_messages = [msg for msg in state.values['messages'] if isinstance(msg, ToolMessage)]
209+
print(f"Checking checkpoint: {len(tool_messages)} tool messages")
210+
211+
if len(tool_messages) == choice - 1:
212+
target_checkpoint = state
213+
print(f"Found target checkpoint with {len(tool_messages)} tool messages")
214+
break
215+
216+
if target_checkpoint:
217+
print(f"Rewinding to checkpoint with {len([msg for msg in target_checkpoint.values['messages'] if isinstance(msg, ToolMessage)])} tool messages")
218+
print(f"Target checkpoint config: {target_checkpoint.config}")
219+
print("=== RESUMING FROM CHECKPOINT ===")
220+
221+
for event in graph.stream(None, target_checkpoint.config):
222+
print(f"Event: {event}")
223+
224+
final_state = graph.get_state(config).values
225+
```
226+
227+
This ability to programmatically search the agent's memory and restart its "thought process" from a specific point is what makes this technique so powerful.
228+
229+
### Full Code Example
230+
231+
??? "Full Code"
232+
```python title="agent_with_state_management.py" hl_lines="71 85 97-100 106-109"
233+
import os
234+
import random
235+
from typing import Annotated, TypedDict
236+
237+
from dotenv import load_dotenv
238+
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
239+
from langchain_core.tools import tool
240+
from langchain_google_genai import ChatGoogleGenerativeAI
241+
from langgraph.graph import StateGraph, END, START
242+
from langgraph.graph.message import add_messages
243+
from langgraph.prebuilt import ToolNode, tools_condition
244+
from langgraph.checkpoint.memory import MemorySaver
245+
246+
# --- 1. Define our tool ---
247+
@tool
248+
def generate_random_number(min_val: int = 1, max_val: int = 100) -> int:
249+
"""Generates a random number within a specified range."""
250+
print(f"--- TOOL CALL: Generating a random number ---")
251+
return random.randint(min_val, max_val)
252+
253+
# --- 2. Define a State ---
254+
class State(TypedDict):
255+
messages: Annotated[list, add_messages]
256+
257+
# --- Setup LLM and Tools ---
258+
load_dotenv()
259+
llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
260+
tools = [generate_random_number]
261+
llm_with_tools = llm.bind_tools(tools)
262+
263+
# --- Define Graph Nodes ---
264+
def chatbot(state: State):
265+
print("--- CHATBOT: Deciding next action ---")
266+
return {"messages": [llm_with_tools.invoke(state["messages"])]}
267+
268+
tool_node = ToolNode(tools)
269+
270+
# --- Assemble the Graph ---
271+
graph_builder = StateGraph(State)
272+
graph_builder.add_node("chatbot", chatbot)
273+
graph_builder.add_node("tools", tool_node)
274+
graph_builder.add_node("human_review", lambda state: state)
275+
276+
graph_builder.add_edge(START, "chatbot")
277+
graph_builder.add_conditional_edges("chatbot", tools_condition, {"tools": "tools", END: "human_review"})
278+
graph_builder.add_edge("tools", "chatbot")
279+
graph_builder.add_edge("human_review", END)
280+
281+
# --- Compile and Run ---
282+
memory = MemorySaver()
283+
graph = graph_builder.compile(
284+
checkpointer=memory,
285+
interrupt_before=["human_review"],
286+
)
287+
288+
# Use a consistent config throughout
289+
config = {"configurable": {"thread_id": "rn-game"}}
290+
user_input = "Generate exactly one random number at a time. I want 3 random numbers total. In your output, include the generate numbers"
291+
initial_state = {"messages": [HumanMessage(content=user_input)]}
292+
293+
# Initial run
294+
print("=== INITIAL RUN ===")
295+
for event in graph.stream(initial_state, config):
296+
print(f"Event: {event}")
297+
298+
final_state = None
299+
300+
# --- Asking User to rewind to desired state in graph by choosing a number n (1/2/3), by resetting the graph to state before the nth number generation---
301+
while True:
302+
# Get current state - this is always the latest state
303+
current_state = graph.get_state(config)
304+
305+
print("\n--- HUMAN REVIEW ---")
306+
print("Agent's message:")
307+
print(current_state.values["messages"][-1].content)
308+
309+
feedback = input("Type a number to rewind to stage before nth random number generation or press Enter to proceed as it is: ").lower()
310+
311+
try:
312+
choice = int(feedback)
313+
if 1 <= choice <= 3:
314+
print(f"Rewinding to re-generate from step {choice}...")
315+
316+
# Get the CURRENT state history
317+
current_history = list(graph.get_state_history(config))
318+
319+
print(f"Current history has {len(current_history)} checkpoints")
320+
321+
# Find the checkpoint that has exactly (choice - 1) tool messages
322+
target_checkpoint = None
323+
324+
# Process in chronological order to find the right checkpoint
325+
for state in reversed(current_history):
326+
tool_messages = [msg for msg in state.values['messages'] if isinstance(msg, ToolMessage)]
327+
print(f"Checking checkpoint: {len(tool_messages)} tool messages")
328+
329+
if len(tool_messages) == choice - 1:
330+
target_checkpoint = state
331+
print(f"Found target checkpoint with {len(tool_messages)} tool messages")
332+
break
333+
334+
if target_checkpoint:
335+
print(f"Rewinding to checkpoint with {len([msg for msg in target_checkpoint.values['messages'] if isinstance(msg, ToolMessage)])} tool messages")
336+
print(f"Target checkpoint config: {target_checkpoint.config}")
337+
338+
print("=== RESUMING FROM CHECKPOINT ===")
339+
340+
for event in graph.stream(None, target_checkpoint.config):
341+
print(f"Event: {event}")
342+
343+
344+
final_state = graph.get_state(config).values
345+
346+
print("Checkpoint resume completed.")
347+
break
348+
else:
349+
print("Could not find the specified checkpoint.")
350+
print("Available checkpoints in current history:")
351+
for i, state in enumerate(reversed(current_history)):
352+
tool_messages = [msg for msg in state.values['messages'] if isinstance(msg, ToolMessage)]
353+
print(f" Checkpoint {i}: {len(tool_messages)} tool messages")
354+
else:
355+
print("Invalid number. Please enter 1, 2, or 3.")
356+
except Exception as e:
357+
final_state = current_state.values
358+
print("No reset performed. Either a number was not provided or reset was not needed.")
359+
break
360+
361+
362+
if final_state:
363+
debug_state = final_state
364+
else:
365+
debug_state = current_state
366+
367+
## State Inspection
368+
print("\n" + "="*50)
369+
print("--- FINAL, CLEAN GRAPH HISTORY ---")
370+
for i, message in enumerate(final_state['messages']):
371+
msg_type = message.__class__.__name__
372+
print(f"Step {i}: {msg_type}")
373+
if isinstance(message, AIMessage) and message.tool_calls:
374+
print(f" - Details: Requesting tool '{message.tool_calls[0]['name']}'")
375+
elif isinstance(message, ToolMessage):
376+
print(f" - Details: Tool returned '{message.content}'")
377+
else:
378+
print(f" - Details: '{message.content}'")
379+
```
380+
If we check the final state history, it won't have any trace of any reset and will render as first time random numbers being output. This demonstration shows how we can ask agent to travel back in time.
381+
![final-state](./images/state-management-output.png)

0 commit comments

Comments
 (0)