[langgraph]: Known limitations with Pydantic BaseModel as State

[sundaraa-deshaw]

Type of issue

issue / bug

Language

Python

Description

The documentation for using Pydantic model for Graph State states that:

Run-time validation only occurs on inputs into nodes, not on the outputs.

However, it seems that the validation occurs only for the input to the first node but not to subsequent nodes. e.g.

from pydantic import BaseModel

class PydanticState(BaseModel):
    messages: list[AnyMessage] = []
    state_1: str = ""
    state_2: str = ""
    source_documents: list[Document] = []

def node_1(state: PydanticState):
    return {
        "state_1": 1,
        "messages": "hi"
    }

def node_2(state: PydanticState):
    print(state)
    print(PydanticState.model_validate(state))
    return {}

pydantic_graph = (
    StateGraph(PydanticState)
    .add_node(node_1)
    .add_node(node_2)
    .add_edge(START, "node_1")
    .add_edge("node_1", "node_2")
    .add_edge("node_2", END)
    .compile()
)

pydantic_graph.invoke({})

produces:

messages='hi' state_1=1 state_2='' source_documents=[]
messages='hi' state_1=1 state_2='' source_documents=[]
{'messages': 'hi', 'state_1': 1, 'state_2': 'from node_2'}

Expected behavior is for Pydantic to throw the validation for messages and state_1