From 661476e88d163e224a079f683a438a19daf8889d Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 12 Feb 2025 17:50:25 -0800 Subject: [PATCH] langgraph: fix None handling for pydantic state updates --- libs/langgraph/langgraph/graph/state.py | 4 ++-- libs/langgraph/tests/test_pregel.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 8c93c0c1d..ac79cc497 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -33,7 +33,7 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue -from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN +from langgraph.constants import EMPTY_SEQ, MISSING, NS_END, NS_SEP, SELF, TAG_HIDDEN from langgraph.errors import ( ErrorCode, InvalidUpdateError, @@ -694,7 +694,7 @@ class CompiledStateGraph(CompiledGraph): return [ (k, getattr(input, k)) for k in output_keys - if getattr(input, k, None) is not None + if getattr(input, k, MISSING) is not MISSING ] else: msg = create_error_message( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ccd4fae94..6c3603499 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6482,3 +6482,16 @@ def test_node_destinations() -> None: Edge(source="child", target="node_b", data="foo", conditional=True), Edge(source="child", target="node_c", data="bar", conditional=True), ] == graph.edges + + +def test_pydantic_none_state_update() -> None: + from pydantic import BaseModel + + class State(BaseModel): + foo: str | None + + def node_a(state: State) -> State: + return State(foo=None) + + graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() + assert graph.invoke({"foo": ""}) == {"foo": None}