diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a37f1d951..e9b92ac37 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1221,7 +1221,7 @@ class Pregel(PregelProtocol): # no values as END, just clear all tasks if values is None and as_node == END: if len(updates) > 1: - raise ValueError( + raise InvalidUpdateError( "Cannot apply multiple updates when clearing state" ) @@ -1284,7 +1284,7 @@ class Pregel(PregelProtocol): # no values, empty checkpoint if values is None and as_node is None: if len(updates) > 1: - raise ValueError( + raise InvalidUpdateError( "Cannot create empty checkpoint with multiple updates" ) @@ -1308,7 +1308,9 @@ class Pregel(PregelProtocol): # no values, copy checkpoint if values is None and as_node == "__copy__": if len(updates) > 1: - raise ValueError("Cannot copy checkpoint with multiple updates") + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint @@ -1429,8 +1431,8 @@ class Pregel(PregelProtocol): run_tasks.append(task) run_task_ids.append(task_id) - for task in run_tasks: run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task run.invoke( values, @@ -1560,7 +1562,7 @@ class Pregel(PregelProtocol): # no values, just clear all tasks if values is None and as_node == END: if len(updates) > 1: - raise ValueError( + raise InvalidUpdateError( "Cannot apply multiple updates when clearing state" ) @@ -1623,7 +1625,7 @@ class Pregel(PregelProtocol): # no values, empty checkpoint if values is None and as_node is None: if len(updates) > 1: - raise ValueError( + raise InvalidUpdateError( "Cannot create empty checkpoint with multiple updates" ) @@ -1647,7 +1649,9 @@ class Pregel(PregelProtocol): # no values, copy checkpoint if values is None and as_node == "__copy__": if len(updates) > 1: - raise ValueError("Cannot copy checkpoint with multiple updates") + raise InvalidUpdateError( + "Cannot copy checkpoint with multiple updates" + ) next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint @@ -1765,8 +1769,8 @@ class Pregel(PregelProtocol): run_tasks.append(task) run_task_ids.append(task_id) - for task in run_tasks: run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task await run.ainvoke( values, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 0d530ddb5..4cca670d6 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -135,7 +135,7 @@ class Interrupt: class BulkUpdate(NamedTuple): values: dict[str, Any] | None - as_node: Optional[str] + as_node: Optional[str] = None class PregelTask(NamedTuple): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d8d51a8a0..8bbeb3a3b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7613,3 +7613,70 @@ def test_parallel_interrupts_double( assert invokes == 5 assert len(events) == 5 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_bulk_state_updates( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + graph.bulk_update_state(config, [({"foo": "bar"}, "node_a")]) + + # Then bulk update with both nodes + graph.bulk_update_state( + config, + [ + ({"foo": "updated"}, "node_a"), + ({"baz": "new"}, "node_b"), + ], + ) + + state = graph.get_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = list(checkpointer.list({"configurable": {"thread_id": "1"}})) + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state( + config, + [({"foo": "error"}, None), ({"bar": "error"}, None)], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + graph.bulk_update_state(config, []) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state(config, [(None, "__end__"), (None, "__copy__")]) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 60cbd2547..f6f38300a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7838,3 +7838,70 @@ async def test_handles_multiple_interrupts_from_tasks() -> None: assert len(result) == 2 assert result[0] == "Added James!" assert result[1] == "Added Will!" + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_bulk_state_updates( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + await graph.abulk_update_state(config, [({"foo": "bar"}, "node_a")]) + + # Then bulk update with both nodes + await graph.abulk_update_state( + config, + [ + ({"foo": "updated"}, "node_a"), + ({"baz": "new"}, "node_b"), + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = list(checkpointer.list({"configurable": {"thread_id": "1"}})) + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state( + config, + [({"foo": "error"}, None), ({"bar": "error"}, None)], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + await graph.abulk_update_state(config, []) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state(config, [(None, "__end__"), (None, "__copy__")])