diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1b5ef0191..2dcd2e95e 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -760,23 +760,22 @@ class CompiledStateGraph(CompiledGraph): updates.extend(_get_updates(i) or ()) return updates elif get_type_hints(type(input)): - # if input is a Pydantic model, only update values - # for the keys that have been explicitly set by the users - # (this is needed to avoid sending updates for fields with None defaults) - output_keys_ = output_keys # Pydantic v2 - if hasattr(input, "model_fields_set"): - output_keys_ = [ - k for k in output_keys if k in input.model_fields_set - ] + if hasattr(input, "model_fields"): + defaults = {k: v.default for k, v in input.model_fields.items()} # Pydantic v1 - elif hasattr(input, "__fields_set__"): - output_keys_ = [k for k in output_keys if k in input.__fields_set__] + elif hasattr(input, "__fields__"): + defaults = {k: v.default for k, v in input.__fields__.items()} + else: + defaults = {} + # if input is a Pydantic model, only update values + # that are different from the default values return [ - (k, getattr(input, k)) - for k in output_keys_ - if getattr(input, k, MISSING) is not MISSING + (k, value) + for k in output_keys + if (value := getattr(input, k, MISSING)) is not MISSING + and value != defaults.get(k) ] else: msg = create_error_message( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index dfd3ed55e..938ee63f3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6658,6 +6658,40 @@ def test_pydantic_none_state_update() -> None: assert graph.invoke({"foo": ""}) == {"foo": None} +def test_pydantic_state_mutation() -> None: + from pydantic import BaseModel, Field + + class Inner(BaseModel): + a: int = 0 + + class State(BaseModel): + inner: Inner = Inner() + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return state + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + # test w/ default_factory + class State(BaseModel): + inner: Inner = Field(default_factory=Inner) + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return state + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + def test_get_stream_writer() -> None: class State(TypedDict): foo: str