Using Pydantic model as state
Previously, we might have used TypedDict the most, but Pydantic model is also a great choice. Let’s see how this type can serve as the graph state.
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from pydantic import BaseModel
# Overall state of the graph (this is the shared public state between nodes)
class OverallState(BaseModel):
a: str
def node(state: OverallState):
return {"a": "goodbye"}
# Build state graph
builder = StateGraph(OverallState)
builder.add_node(node) # node_1 is the first node
builder.add_edge(START, "node") # Start the graph with node_1
builder.add_edge("node", END) # End the graph after node_1
graph = builder.compile()
# Test the graph with valid input
graph.invoke({"a": "hello"})
When we inherit from BaseModel in OverallState, and define the property a as a str type, if our input is correct, it will output:
try:
graph.invoke({"a": 123}) # Should be a string
except Exception as e:
print("An exception was raised because `a` is an integer rather than a string.")
print(e)
An exception was raised because `a` is an integer rather than a string.
1 validation error for OverallState
a Input should be a valid string [type=string_type, input_value=123, input_type=int] For further information visit https://errors.pydantic.dev/2.10/v/string_type
As we can see, it acts as a form of input validation.
Runtime validation also applies to multi-node graphs. In the example below, bad_node will update a to an integer.
Since runtime validation occurs on the input, a validation error will occur when calling ok_node (rather than when bad_node returns a state update inconsistent with the model).
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from pydantic import BaseModel
# Overall state of the graph (this is the shared public state between nodes)
class OverallState(BaseModel):
a: str
def bad_node(state: OverallState):
return {"a": 123 # Invalid}
def ok_node(state: OverallState):
return {"a": "goodbye"}
# Build state graph
builder = StateGraph(OverallState)
builder.add_node(bad_node)
builder.add_node(ok_node)
builder.add_edge(START, "bad_node")
builder.add_edge("bad_node", "ok_node")
builder.add_edge("ok_node", END)
graph = builder.compile()
# Test the graph with valid input
try:
graph.invoke({"a": "hello"})
except Exception as e:
print("An exception was raised because bad_node set `a` to an integer.")
print(e)
Note that it will raise an error during conversion, rather than discovering the type mismatch during the update.
