langgraph: handle pydantic updates consistently in Command (#4255)

Fixes https://github.com/langchain-ai/langgraph/issues/3950

libs/langgraph/langgraph/graph/state.py

@@ -77,7 +77,7 @@ from langgraph.pregel.write import (
 )
 from langgraph.store.base import BaseStore
 from langgraph.types import All, Checkpointer, Command, RetryPolicy
-from langgraph.utils.fields import get_field_default
+from langgraph.utils.fields import get_field_default, get_update_as_tuples
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
 
@@ -761,32 +761,7 @@ class CompiledStateGraph(CompiledGraph):
                         updates.extend(_get_updates(i) or ())
                 return updates
             elif (t := type(input)) and get_type_hints(t):
-                # Pydantic v2
-                if isinstance(input, BaseModelV1):
-                    keep: Optional[set[str]] = input.__fields_set__
-                    defaults = {k: v.default for k, v in t.__fields__.items()}
-                elif isinstance(input, BaseModel):
-                    keep = input.model_fields_set
-                    defaults = {k: v.default for k, v in input.model_fields.items()}
-                # Pydantic v1
-                else:
-                    keep = None
-                    defaults = {}
-
-                # NOTE: This behavior for Pydantic is somewhat inelegant,
-                # but we keep around for backwards compatibility
-                # if input is a Pydantic model, only update values
-                # that are different from the default values or in the keep set
-                return [
-                    (k, value)
-                    for k in output_keys
-                    if (value := getattr(input, k, MISSING)) is not MISSING
-                    and (
-                        value is not None
-                        or defaults.get(k, MISSING) is not None
-                        or (keep is not None and k in keep)
-                    )
-                ]
+                return get_update_as_tuples(input, output_keys)
             else:
                 msg = create_error_message(
                     message=f"Expected dict, got {input}",

libs/langgraph/langgraph/types.py

@@ -23,6 +23,7 @@ from langchain_core.runnables import Runnable, RunnableConfig
 from typing_extensions import Self
 
 from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
+from langgraph.utils.fields import get_update_as_tuples
 
 if TYPE_CHECKING:
     from langgraph.pregel.protocol import PregelProtocol
@@ -318,7 +319,7 @@ class Command(Generic[N], ToolOutputMixin):
         ):
             return self.update
         elif hints := get_type_hints(type(self.update)):
-            return [(k, getattr(self.update, k)) for k in hints]
+            return get_update_as_tuples(self.update, tuple(hints.keys()))
         elif self.update is not None:
             return [("__root__", self.update)]
         else:

libs/langgraph/langgraph/utils/fields.py

@@ -1,8 +1,14 @@
 import dataclasses
-from typing import Any, Generator, Optional, Type, Union, get_type_hints
+from typing import Any, Generator, Optional, Sequence, Type, Union, get_type_hints
 
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import Annotated, NotRequired, ReadOnly, Required, get_origin
 
+# NOTE: this is redefined here separately from langgraph.constants
+# to avoid a circular import
+MISSING = object()
+
 
 def _is_optional_type(type_: Any) -> bool:
     """Check if a type is Optional."""
@@ -147,3 +153,33 @@ def get_enhanced_type_hints(
             pass
 
         yield name, typ, default, description
+
+
+def get_update_as_tuples(input: Any, keys: Sequence[str]) -> list[tuple[str, Any]]:
+    """Get Pydantic state update as a list of (key, value) tuples."""
+    # Pydantic v1
+    if isinstance(input, BaseModelV1):
+        keep: Optional[set[str]] = input.__fields_set__
+        defaults = {k: v.default for k, v in input.__fields__.items()}
+    # Pydantic v2
+    elif isinstance(input, BaseModel):
+        keep = input.model_fields_set
+        defaults = {k: v.default for k, v in input.model_fields.items()}
+    else:
+        keep = None
+        defaults = {}
+
+    # NOTE: This behavior for Pydantic is somewhat inelegant,
+    # but we keep around for backwards compatibility
+    # if input is a Pydantic model, only update values
+    # that are different from the default values or in the keep set
+    return [
+        (k, value)
+        for k in keys
+        if (value := getattr(input, k, MISSING)) is not MISSING
+        and (
+            value is not None
+            or defaults.get(k, MISSING) is not None
+            or (keep is not None and k in keep)
+        )
+    ]

libs/langgraph/tests/test_pregel.py

@@ -7246,6 +7246,39 @@ def test_pydantic_none_state_update() -> None:
     assert graph.invoke({"foo": ""}) == {"foo": None}
 
 
+def test_pydantic_state_update_command() -> None:
+    from pydantic import BaseModel
+
+    class State(BaseModel):
+        foo: Optional[str]
+
+    def node_a(state: State) -> State:
+        return Command(update=State(foo=None))
+
+    graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()
+    assert graph.invoke({"foo": ""}) == {"foo": None}
+
+    class State(BaseModel):
+        foo: Optional[str] = None
+        bar: Optional[str] = None
+
+    def node_a(state: State):
+        return State(foo="foo")
+
+    def node_b(state: State):
+        return Command(update=State(bar="bar"))
+
+    builder = StateGraph(State)
+    builder.add_node(node_a)
+    builder.add_node(node_b)
+    builder.add_edge(START, "node_a")
+    builder.add_edge("node_a", "node_b")
+    builder.add_edge("node_b", END)
+    graph = builder.compile()
+
+    assert graph.invoke(State()) == {"foo": "foo", "bar": "bar"}
+
+
 def test_pydantic_state_mutation() -> None:
     from pydantic import BaseModel, Field
 
@@ -7280,6 +7313,40 @@ def test_pydantic_state_mutation() -> None:
     assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)}
 
 
+def test_pydantic_state_mutation_command() -> 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 Command(update=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 Command(update=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