Surprising results when using pydantic for state schema

[eyurtsev]

Privileged issue

• I am a LangChain maintainer, or was asked directly by a LangChain maintainer to create an issue here.

Issue Content

Run time validation with pydantic schema does not work with output schema.

from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

from pydantic import BaseModel

# The overall state of the graph (this is the public state shared across nodes)
class OverallState(BaseModel):
    a: float

def node_1(state: OverallState) -> OverallState:
    return {
        "a": "foo"
    }

# Build the state graph
builder = StateGraph(OverallState, output=OverallState)
builder.add_node(node_1)  # node_1 is the first node
builder.add_edge(START, "node_1")  # Start the graph with node_1
builder.add_edge("node_1", END)  # End the graph after node_1
graph = builder.compile()

graph.invoke(
    {
        "a": 2.3
    }
)

Result

{'a': 'foo'}

This result does not match the output schema!

Expected

Expected a validation error

[gbaian10]

Hi, I would like to ask if you discourage performing operations on the state (the BaseModel object itself) within the action function and returning the entire BaseModel object?

After I studying the code, the biggest drawback of doing so seems to be that when you have many fields and only need to update one of them, the current underlying logic still reads all fields of the BaseModel and updates all values (even though only one value has actually changed).

from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, ConfigDict


class OverallState(BaseModel):
    model_config = ConfigDict(validate_assignment=False)

    a: float


def node_1(state: OverallState) -> OverallState:
    # state.a = "foo"  # If validate_assignment is True, it will raise an error here
    state.a = "foo"
    return state


def node_2(state: dict) -> OverallState:
    print(state)  # It will raise an error here
    return state


builder = StateGraph(OverallState, output=OverallState)
builder.add_node(node_1)
builder.add_node(node_2)
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
builder.add_edge("node_2", end_key=END)
graph = builder.compile()

print(graph.invoke(OverallState(a=2.3)))

Additionally, the BaseModel has a validate_assignment setting which is set to False by default.

Therefore, even if the user wants to apply a "monkey patch" (since Pydantic allows reassignment to a different type by default), an error will still occur at the next node due to the reconstruction of the object.

[gbaian10]

I once mistakenly thought that BaseModel was the same object being passed in the graph, but it seems that’s not the case.

Because it involves database writing and reading, unless it's in memory, the object might always have to be recreated.