GraphAI provides a state management system that allows data to persist across node executions. This state is essential for maintaining context throughout the execution of a graph.
The Graph class provides several methods for working with state:
Copy
Ask AI
# Get the current statecurrent_state = graph.get_state()# Set a new state (replaces existing state)graph.set_state({"new_state": "value"})# Update the state (merges with existing state)graph.update_state({"additional": "data"})# Reset the state to an empty dictionarygraph.reset_state()
Nodes can access the graph state by including a state parameter in their function signature:
Copy
Ask AI
from graphai import node@nodeasync def stateful_node(input: dict, state: dict): # Access the state history = state.get("history", []) context = state.get("context", "") # Use the state in processing processed_data = process_with_context(input["data"], context) # Return updated information return {"result": processed_data}
The state is passed automatically to any node that includes a state parameter.
The most common way to modify state is to include state changes in the node’s return value:
Copy
Ask AI
@nodeasync def update_history(input: dict, state: dict): # Get current history history = state.get("history", []) # Add new entry to history new_history = history + [input["query"]] # Return with updated history return { "result": process(input["query"]), "history": new_history # This updates the state's history }
When the node returns, any keys in the return dictionary are merged with the current state.
For more complex workflows, you can use graph methods directly:
Copy
Ask AI
@nodeasync def complex_state_update(input: dict, graph): # Process data result = process(input["data"]) # Get current state current_state = graph.get_state() # Make complex updates current_state["history"].append(input["query"]) current_state["metadata"]["last_processed"] = datetime.now().isoformat() # Set the updated state graph.set_state(current_state) return {"result": result}
This approach is less common but provides more flexibility for complex state manipulations.
The state persists for the lifetime of the graph object. If you need to persist state between graph executions:
Copy
Ask AI
# Save state after executionresult = await graph.execute(input_data)saved_state = graph.get_state()# Store saved_state somewhere (e.g., database)store_state(saved_state)# Later, restore staterestored_state = load_state()graph.set_state(restored_state)
State is scoped to the graph instance. If you create multiple graph instances, each will have its own independent state:
Copy
Ask AI
graph1 = Graph(initial_state={"id": "graph1"})graph2 = Graph(initial_state={"id": "graph2"})# These operate on different state objectsgraph1.update_state({"value": 1})graph2.update_state({"value": 2})
It’s important to understand the difference between state and input:
Input: Data passed to the current node execution
State: Persistent data that’s available across multiple node executions
For example:
Copy
Ask AI
@nodeasync def process_with_both(input: dict, state: dict): # Input is specific to this execution query = input["query"] # State persists across executions history = state.get("history", []) # Use both result = process_with_history(query, history) return { "result": result, "history": history + [query] # Update state for future nodes }