Add tests

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,

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):

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__")])

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__")])