From 857f3e4a38921ac8ea3add026c67c0bc042742e3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 13 Mar 2025 17:31:38 -0700 Subject: [PATCH] Update handling of updates/inputs passed in as pydantic models - remove usage of require_at_least_one_of, we shouldn't be enforcing presence of keys in inputs/updates, an empty dict is a valid input/update - ensure that values that were explcitly set/assigned in pydantic model are saved even if equal to default value --- libs/langgraph/langgraph/graph/state.py | 20 +++-- libs/langgraph/langgraph/pregel/debug.py | 96 ++++++++++++++--------- libs/langgraph/langgraph/pregel/read.py | 1 - libs/langgraph/langgraph/pregel/write.py | 15 +--- libs/langgraph/tests/test_pregel.py | 95 +++++++++++++--------- libs/langgraph/tests/test_pregel_async.py | 33 -------- 6 files changed, 131 insertions(+), 129 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f4143ca1f..140f97ef5 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -762,23 +762,28 @@ class CompiledStateGraph(CompiledGraph): else: updates.extend(_get_updates(i) or ()) return updates - elif get_type_hints(type(input)): + elif (t := type(input)) and get_type_hints(t): # Pydantic v2 - if hasattr(input, "model_fields"): - defaults = {k: v.default for k, v in input.model_fields.items()} + if isinstance(input, BaseModel): + keep: Optional[set[str]] = input.model_fields_set + defaults = {k: v.default for k, v in t.model_fields.items()} # Pydantic v1 - elif hasattr(input, "__fields__"): - defaults = {k: v.default for k, v in input.__fields__.items()} + elif isinstance(input, BaseModelV1): + keep = input.__fields_set__ + defaults = {k: v.default for k, v in t.__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 + # 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 != defaults.get(k) + and (value != defaults.get(k) or (keep is not None and k in keep)) ] else: msg = create_error_message( @@ -804,7 +809,6 @@ class CompiledStateGraph(CompiledGraph): ChannelWrite( write_entries, tags=[TAG_HIDDEN], - require_at_least_one_of=output_keys, ), ], ) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 8429fd538..54c260682 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -25,8 +25,10 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_NS, ERROR, INTERRUPT, + MISSING, NS_END, NS_SEP, + RETURN, TAG_HIDDEN, ) from langgraph.pregel.io import read_channels @@ -132,7 +134,9 @@ def map_debug_task_results( "id": task.id, "name": task.name, "error": next((w[1] for w in writes if w[0] == ERROR), None), - "result": [w for w in writes if w[0] in stream_channels_list], + "result": [ + w for w in writes if w[0] in stream_channels_list or w[0] == RETURN + ], "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], }, } @@ -264,49 +268,63 @@ def tasks_w_writes( ) -> tuple[PregelTask, ...]: """Apply writes / subgraph states to tasks to be returned in a StateSnapshot.""" pending_writes = pending_writes or [] - return tuple( - PregelTask( - task.id, - task.name, - task.path, - next( - ( - exc - for tid, n, exc in pending_writes - if tid == task.id and n == ERROR - ), - None, - ), - tuple( - v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT - ), - states.get(task.id) if states else None, + out: list[PregelTask] = [] + for task in tasks: + rtn = next( ( + val + for tid, chan, val in pending_writes + if tid == task.id and chan == RETURN + ), + MISSING, + ) + out.append( + PregelTask( + task.id, + task.name, + task.path, next( ( - val - for tid, chan, val in pending_writes - if tid == task.id and chan == output_keys + exc + for tid, n, exc in pending_writes + if tid == task.id and n == ERROR ), None, - ) - if isinstance(output_keys, str) - else { - chan: val - for tid, chan, val in pending_writes - if tid == task.id - and ( - chan == output_keys - if isinstance(output_keys, str) - else chan in output_keys + ), + tuple( + v + for tid, n, v in pending_writes + if tid == task.id and n == INTERRUPT + ), + states.get(task.id) if states else None, + ( + rtn + if rtn is not MISSING + else next( + ( + val + for tid, chan, val in pending_writes + if tid == task.id and chan == output_keys + ), + None, ) - } + if isinstance(output_keys, str) + else { + chan: val + for tid, chan, val in pending_writes + if tid == task.id + and ( + chan == output_keys + if isinstance(output_keys, str) + else chan in output_keys + ) + } + ) + if any( + w[0] == task.id and w[1] not in (ERROR, INTERRUPT) + for w in pending_writes + ) + else None, ) - if any( - w[0] == task.id and w[1] not in (ERROR, INTERRUPT) - for w in pending_writes - ) - else None, ) - for task in tasks - ) + return tuple(out) diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index e3d5aab3d..728159156 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -201,7 +201,6 @@ class PregelNode(Runnable): writers[-2] = ChannelWrite( writes=writers[-2].writes + writers[-1].writes, tags=writers[-2].tags, - require_at_least_one_of=writers[-2].require_at_least_one_of, ) writers.pop() return writers diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 7469824e4..0744b4434 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -49,21 +49,18 @@ class ChannelWrite(RunnableCallable): writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]] """Sequence of write entries or Send objects to write.""" - require_at_least_one_of: Optional[Sequence[str]] - """If defined, at least one of these channels must be written to.""" def __init__( self, writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], *, tags: Optional[Sequence[str]] = None, - require_at_least_one_of: Optional[Sequence[str]] = None, + require_at_least_one_of: Optional[Sequence[str]] = None, # ignored ): super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) self.writes = cast( list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes ) - self.require_at_least_one_of = require_at_least_one_of def get_name( self, suffix: Optional[str] = None, *, name: Optional[str] = None @@ -96,7 +93,6 @@ class ChannelWrite(RunnableCallable): self.do_write( config, writes, - self.require_at_least_one_of if input is not None else None, ) return input @@ -112,7 +108,6 @@ class ChannelWrite(RunnableCallable): self.do_write( config, writes, - self.require_at_least_one_of if input is not None else None, ) return input @@ -120,7 +115,7 @@ class ChannelWrite(RunnableCallable): def do_write( config: RunnableConfig, writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], - require_at_least_one_of: Optional[Sequence[str]] = None, + require_at_least_one_of: Optional[Sequence[str]] = None, # ignored ) -> None: # validate for w in writes: @@ -151,12 +146,6 @@ class ChannelWrite(RunnableCallable): tuples.append((w.channel, value)) else: raise ValueError(f"Invalid write entry: {w}") - # assert required channels - if require_at_least_one_of is not None: - if not {chan for chan, _ in tuples} & set(require_at_least_one_of): - raise InvalidUpdateError( - f"Must write to at least one of {require_at_least_one_of}" - ) write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND] write(tuples) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 5a4c9bfa8..85b5f1ad1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2563,7 +2563,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( request: pytest.FixtureRequest, checkpointer_name: str, ) -> None: - from pydantic.v1 import BaseModel, ValidationError + from pydantic.v1 import BaseModel, Field, ValidationError checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() @@ -2625,6 +2625,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( answer: Optional[str] = None docs: Optional[list[str]] = None + class UpdateDocs34(BaseModel): + docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"]) + def rewrite_query(data: State) -> State: assert isinstance(data.inner, InnerObject) return {"query": f"query: {data.query}"} @@ -2638,7 +2641,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( def retriever_two(data: State) -> State: time.sleep(0.1) - return {"docs": ["doc3", "doc4"]} + return UpdateDocs34() def qa(data: State) -> State: return {"answer": ",".join(data.docs)} @@ -2734,7 +2737,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( request: pytest.FixtureRequest, checkpointer_name: str, ) -> None: - from pydantic import BaseModel, ConfigDict, ValidationError + from pydantic import BaseModel, ConfigDict, Field, ValidationError checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() @@ -2787,6 +2790,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( answer: Optional[str] = None docs: Optional[list[str]] = None + class UpdateDocs34(BaseModel): + docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"]) + class Input(BaseModel): query: str inner: InnerObject @@ -2808,7 +2814,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( def retriever_two(data: State) -> State: time.sleep(0.1) - return {"docs": ["doc3", "doc4"]} + return UpdateDocs34() def qa(data: State) -> State: return {"answer": ",".join(data.docs)} @@ -5671,37 +5677,6 @@ def test_command_goto_with_static_breakpoints( assert result == {"foo": "abc|node-1|node-2|node-2"} -def test_nested_graph_state_error_handling(): - """Test error handling when updating state in nested graphs.""" - - class State(TypedDict): - count: int - - def child_node(state: State): - return {"count": state["count"] + 1} - - child = StateGraph(State) - child.add_node("child", child_node) - child.add_edge(START, "child") - - parent = StateGraph(State) - parent.add_node("child_graph", child.compile()) - parent.add_edge(START, "child_graph") - - app = parent.compile(checkpointer=MemorySaver()) - - # Test invalid state update on parent - with pytest.raises(InvalidUpdateError): - app.update_state({"configurable": {"thread_id": "1"}}, {"invalid_key": "value"}) - - # Test invalid state update on child - with pytest.raises(InvalidUpdateError): - app.update_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, - {"invalid_key": "value"}, - ) - - def test_parallel_node_execution(): """Test that parallel nodes execute concurrently.""" @@ -7033,3 +7008,53 @@ def test_interrupt_subgraph_reenter_checkpointer_true( } # confirm that we preserve the state values from the previous invocation assert bar_values == [None, "barbaz", "quxbaz"] + + +def test_empty_invoke() -> None: + from pydantic import BaseModel + + def reducer_merge_dicts( + dict1: dict[Any, Any], dict2: dict[Any, Any] + ) -> dict[Any, Any]: + merged = {**dict1, **dict2} + return merged + + class SimpleGraphState(BaseModel): + x1: Annotated[list[str], operator.add] = [] + x2: Annotated[dict[str, Any], reducer_merge_dicts] = {} + + def update_x1_1(state: SimpleGraphState): + print(state) + return {"x1": ["111"]} + + def update_x1_2(state: SimpleGraphState): + print(state) + state.x1.append("222") + return {"x1": ["222"]} + + def update_x2_1(state: SimpleGraphState): + print(state) + return {"x2": {"111": 111}} + + def update_x2_2(state: SimpleGraphState): + print(state) + return {"x2": {"222": 222}} + + graph = StateGraph(SimpleGraphState) + graph.add_node("x1_1_node", update_x1_1) + graph.add_node("x1_2_node", update_x1_2) + graph.add_node("x2_1_node", update_x2_1) + graph.add_node("x2_2_node", update_x2_2) + graph.add_edge("x1_1_node", "x1_2_node") + graph.add_edge("x1_2_node", "x2_1_node") + graph.add_edge("x2_1_node", "x2_2_node") + + graph.add_edge(START, "x1_1_node") + graph.add_edge("x2_2_node", END) + + compiled = graph.compile() + + assert compiled.invoke(SimpleGraphState()).get("x2") == { + "111": 111, + "222": 222, + } diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 105a0c9e6..39d3d34a6 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6654,39 +6654,6 @@ async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> N assert result == {"foo": "abc|node-1|node-2|node-2"} -async def test_nested_graph_state_error_handling(): - """Test error handling when updating state in nested graphs.""" - - class State(TypedDict): - count: int - - def child_node(state: State): - return {"count": state["count"] + 1} - - child = StateGraph(State) - child.add_node("child", child_node) - child.add_edge(START, "child") - - parent = StateGraph(State) - parent.add_node("child_graph", child.compile()) - parent.add_edge(START, "child_graph") - - app = parent.compile(checkpointer=MemorySaver()) - - # Test invalid state update on parent - with pytest.raises(InvalidUpdateError): - await app.aupdate_state( - {"configurable": {"thread_id": "1"}}, {"invalid_key": "value"} - ) - - # Test invalid state update on child - with pytest.raises(InvalidUpdateError): - await app.aupdate_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, - {"invalid_key": "value"}, - ) - - async def test_parallel_node_execution(): """Test that parallel nodes execute concurrently."""