From 865fba9d50c28026904ef8a4b1ac485772889d37 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 24 Jun 2025 19:10:02 +0200 Subject: [PATCH 01/15] feat(langgraph): task masquerading with update state --- libs/langgraph/langgraph/pregel/__init__.py | 116 ++++-- libs/langgraph/langgraph/types.py | 1 + libs/langgraph/tests/test_pregel.py | 406 ++++++++++++++++++++ 3 files changed, 495 insertions(+), 28 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 05c157d38..4f180566b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance(self._channels, list), ( - "Cannot read additional channels when subscribed to single channels" - ) + assert isinstance( + self._channels, list + ), "Cannot read additional channels when subscribed to single channels" self._channels.extend(channels) return self @@ -1354,7 +1354,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou Args: config: The config to apply the updates to. supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. - Each update is a tuple of the form `(values, as_node)`. + Each update is a tuple of the form `(values, as_node, task_id)` where task_id is optional. Raises: ValueError: If no checkpointer is set or no updates are provided. @@ -1544,28 +1544,78 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou f"Received no input writes for {self.input_channels}" ) - # no values, copy checkpoint - if values is None and as_node == "__copy__": + # copy checkpoint + if as_node == "__copy__": if len(updates) > 1: raise InvalidUpdateError( "Cannot copy checkpoint with multiple updates" ) + if saved is None: + raise InvalidUpdateError("Cannot copy a non-existent checkpoint") + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint next_config = checkpointer.put( - saved.parent_config or saved.config if saved else checkpoint_config, + saved.parent_config + or patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ), next_checkpoint, { "source": "fork", "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, + "parents": saved.metadata.get("parents", {}), }, {}, ) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) + + # we want to both clone a checkpoint and update state in one go. + # reuse the same task ID if possible. + if isinstance(values, list) and len(values) > 0: + # figure out the task IDs for the next update checkpoint + next_tasks = prepare_next_tasks( + next_checkpoint, + saved.pending_writes, + self.nodes, + channels, + managed, + next_config, + step + 2, + step + 4, + for_execution=False, + ) + + tasks_group_by = defaultdict(list) + for task in next_tasks.values(): + tasks_group_by[task.name].append(task.id) + + user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) + for item in values: + if not isinstance(item, dict): + continue + + values = item["values"] + as_node = item["as_node"] + + user_group = user_group_by[as_node] + tasks_group = tasks_group_by[as_node] + + target_idx = len(user_group) + task_id = tasks_group[target_idx] if target_idx < len(tasks_group) else None + + user_group_by[as_node].append( + StateUpdate(values=values, as_node=as_node, task_id=task_id) + ) + + return perform_superstep( + patch_checkpoint_map(next_config, saved.metadata), + [item for lst in user_group_by.values() for item in lst], + ) + + return patch_checkpoint_map(next_config, saved.metadata) + # apply pending writes, if not on specific checkpoint if ( CONFIG_KEY_CHECKPOINT_ID not in config[CONF] @@ -1613,9 +1663,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou apply_writes( checkpoint, channels, tasks, None, self.trigger_to_nodes ) - valid_updates: list[tuple[str, dict[str, Any] | None]] = [] + valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: - values, as_node = updates[0] + values, as_node, task_id = updates[0] # find last node that updated the state, if not provided if as_node is None and len(self.nodes) == 1: as_node = tuple(self.nodes)[0] @@ -1646,9 +1696,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") - valid_updates.append((as_node, values)) + valid_updates.append((as_node, values, task_id)) else: - for values, as_node in updates: + for values, as_node, task_id in updates: if as_node is None: raise InvalidUpdateError( "as_node is required when applying multiple updates" @@ -1656,19 +1706,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") - valid_updates.append((as_node, values)) + valid_updates.append((as_node, values, task_id)) run_tasks: list[PregelTaskWrites] = [] run_task_ids: list[str] = [] - for as_node, values in valid_updates: + for as_node, values, provided_task_id in valid_updates: # create task to run all writers of the chosen node writers = self.nodes[as_node].flat_writers if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + task_id = provided_task_id or str( + uuid5(UUID(checkpoint["id"]), INTERRUPT) + ) run_tasks.append(task) run_task_ids.append(task_id) run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] @@ -1681,6 +1733,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou configurable={ # deque.extend is thread-safe CONFIG_KEY_SEND: writes.extend, + CONFIG_KEY_TASK_ID: task_id, CONFIG_KEY_READ: partial( local_read, _scratchpad( @@ -1750,7 +1803,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou Args: config: The config to apply the updates to. supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. - Each update is a tuple of the form `(values, as_node)`. + Each update is a tuple of the form `(values, as_node, task_id)` where task_id is optional. Raises: ValueError: If no checkpointer is set or no updates are provided. @@ -2006,9 +2059,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou apply_writes( checkpoint, channels, tasks, None, self.trigger_to_nodes ) - valid_updates: list[tuple[str, dict[str, Any] | None]] = [] + valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: - values, as_node = updates[0] + values, as_node, task_id = updates[0] # find last node that updated the state, if not provided if as_node is None and len(self.nodes) == 1: as_node = tuple(self.nodes)[0] @@ -2035,9 +2088,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") - valid_updates.append((as_node, values)) + valid_updates.append((as_node, values, task_id)) else: - for values, as_node in updates: + for values, as_node, task_id in updates: if as_node is None: raise InvalidUpdateError( "as_node is required when applying multiple updates" @@ -2045,19 +2098,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") - valid_updates.append((as_node, values)) + valid_updates.append((as_node, values, task_id)) run_tasks: list[PregelTaskWrites] = [] run_task_ids: list[str] = [] - for as_node, values in valid_updates: + for as_node, values, provided_task_id in valid_updates: # create task to run all writers of the chosen node writers = self.nodes[as_node].flat_writers if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + task_id = provided_task_id or str( + uuid5(UUID(checkpoint["id"]), INTERRUPT) + ) run_tasks.append(task) run_task_ids.append(task_id) run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] @@ -2070,6 +2125,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou configurable={ # deque.extend is thread-safe CONFIG_KEY_SEND: writes.extend, + CONFIG_KEY_TASK_ID: task_id, CONFIG_KEY_READ: partial( local_read, _scratchpad( @@ -2136,24 +2192,28 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou config: RunnableConfig, values: dict[str, Any] | Any | None, as_node: str | None = None, + task_id: str | None = None, ) -> RunnableConfig: """Update the state of the graph with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ - return self.bulk_update_state(config, [[StateUpdate(values, as_node)]]) + return self.bulk_update_state(config, [[StateUpdate(values, as_node, task_id)]]) async def aupdate_state( self, config: RunnableConfig, values: dict[str, Any] | Any, as_node: str | None = None, + task_id: str | None = None, ) -> RunnableConfig: """Asynchronously update the state of the graph with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ - return await self.abulk_update_state(config, [[StateUpdate(values, as_node)]]) + return await self.abulk_update_state( + config, [[StateUpdate(values, as_node, task_id)]] + ) def _defaults( self, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 292deee03..d95c4fd51 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -164,6 +164,7 @@ class Interrupt: class StateUpdate(NamedTuple): values: dict[str, Any] | None as_node: str | None = None + task_id: str | None = None class PregelTask(NamedTuple): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bc32e2480..b6cc11332 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8047,3 +8047,409 @@ def test_timeout_with_parent_command( graph.invoke({"value": "start"}, thread1) assert exc_info.value.args[0].goto == "test_cmd" assert exc_info.value.args[0].update == {"key": "value"} + + +def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None: + """Test forking and updating task results with state history.""" + two_count = 0 + + def checkpoint(snapshot: StateSnapshot) -> tuple[str, dict]: + """Convert a checkpoint to a tuple representation for comparison.""" + return ("checkpoint", {"values": snapshot.values}) + + def task(task_info: PregelTask) -> tuple[str, dict]: + """Convert a task to a tuple representation for comparison.""" + return ("task", {"name": task_info.name, "result": task_info.result}) + + def get_tree(history: list[StateSnapshot]) -> list: + """Build a tree structure from state history for comparison.""" + if not history: + return [] + + # Build a tree structure similar to renderForks + node_map: dict[str, dict] = {} + root_nodes: list[dict] = [] + + # Second pass: establish parent-child relationships + for item in reversed(history): + checkpoint_id = item.config["configurable"]["checkpoint_id"] + parent_checkpoint_id = ( + item.parent_config["configurable"]["checkpoint_id"] + if item.parent_config + else None + ) + node_map[checkpoint_id] = {"item": item, "children": []} + + parent = node_map.get(parent_checkpoint_id) + if parent: + parent["children"].append(node_map[checkpoint_id]) + else: + root_nodes.append(node_map[checkpoint_id]) + + def node_to_tree(node: dict) -> list: + """Convert a node to tree structure.""" + result = [checkpoint(node["item"])] + [ + task(task_info) for task_info in node["item"].tasks + ] + + if len(node["children"]) > 1: + branches = [node_to_tree(child) for child in node["children"]] + return result + [branches] + elif len(node["children"]) == 1: + return result + node_to_tree(node["children"][0]) + else: + return result + + # Process all root nodes + if len(root_nodes) == 1: + return node_to_tree(root_nodes[0]) + elif len(root_nodes) > 1: + # Multiple root nodes - treat as branches + branches = [node_to_tree(node) for node in root_nodes] + return branches + else: + return [] + + # Define the graph with a sequence of nodes + def one(state: dict) -> Command: + return Command(goto=[Send("two", {})], update={"name": "one"}) + + def two(state: dict) -> dict: + nonlocal two_count + two_count += 1 + return {"name": f"two {two_count}"} + + def three(state: dict) -> dict: + return {"name": "three"} + + graph = ( + StateGraph({"name": str}) + .add_node("one", one) + .add_node("two", two) + .add_node("three", three) + .add_edge(START, "one") + .add_edge("one", "two") + .add_edge("two", "three") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + history: list[StateSnapshot] = [] + + # Initial run + graph.invoke({"name": "start"}, config) + history = list(graph.get_state_history(config)) + + assert get_tree(history) == [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + ] + + # Update the start state + graph.invoke( + None, + graph.update_state( + history[4].config, + [StateUpdate(values={"name": "start*"}, as_node="__start__")], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + ], + [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start*"}}), + ("checkpoint", {"values": {"name": "start*"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + ("checkpoint", {"values": {"name": "start* > one"}}), + ("task", {"name": "two", "result": {"name": "two 3"}}), + ("task", {"name": "two", "result": {"name": "two 4"}}), + ("checkpoint", {"values": {"name": "start* > one > two 4 > two 3"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start* > one > two 4 > two 3 > three"}}, + ), + ], + ] + + # Fork from task "one" + # Start from the checkpoint that has the task "one" + assert history[3].values == {"name": "start*"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "one" + + graph.invoke( + None, + graph.update_state( + history[3].config, + [StateUpdate(values={"name": "one*"}, as_node="one")], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + ], + [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start*"}}), + [ + [ + ("checkpoint", {"values": {"name": "start*"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + ("checkpoint", {"values": {"name": "start* > one"}}), + ("task", {"name": "two", "result": {"name": "two 3"}}), + ("task", {"name": "two", "result": {"name": "two 4"}}), + ( + "checkpoint", + {"values": {"name": "start* > one > two 4 > two 3"}}, + ), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start* > one > two 4 > two 3 > three"}}, + ), + ], + [ + ("checkpoint", {"values": {"name": "start*"}}), + ("task", {"name": "one", "result": {"name": "one*"}}), + ("checkpoint", {"values": {"name": "start* > one*"}}), + ("task", {"name": "two", "result": {"name": "two 5"}}), + ("checkpoint", {"values": {"name": "start* > one* > two 5"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start* > one* > two 5 > three"}}, + ), + ], + ], + ], + ] + + two_count = 0 + config = {"configurable": {"thread_id": "2"}} + + # Initialize the thread once again + graph.invoke({"name": "start"}, config) + history = list(graph.get_state_history(config)) + + # Fork from task "two" + # Start from the checkpoint that has the task "two" + assert history[2].values == {"name": "start > one"} + + graph.invoke( + None, + graph.update_state( + history[2].config, + [ + StateUpdate(values={"name": "two 3"}, as_node="two"), + StateUpdate(values={"name": "two 4"}, as_node="two"), + ], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + [ + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 2 > two 1 > three"}}, + ), + ], + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 3"}}), + ("task", {"name": "two", "result": {"name": "two 4"}}), + ("checkpoint", {"values": {"name": "start > one > two 3 > two 4"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4 > three"}}, + ), + ], + ], + ] + + # Fork task three + assert history[1].values == {"name": "start > one > two 3 > two 4"} + assert len(history[1].tasks) == 1 + assert history[1].tasks[0].name == "three" + + graph.invoke( + None, + graph.update_state( + history[1].config, + [StateUpdate(values={"name": "three*"}, as_node="three")], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + [ + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 2 > two 1 > three"}}, + ), + ], + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 3"}}), + ("task", {"name": "two", "result": {"name": "two 4"}}), + [ + [ + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4"}}, + ), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4 > three"}}, + ), + ], + [ + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4"}}, + ), + ("task", {"name": "three", "result": {"name": "three*"}}), + ( + "checkpoint", + { + "values": { + "name": "start > one > two 3 > two 4 > three*" + } + }, + ), + ], + ], + ], + ], + ] + + # Regenerate task three + assert history[3].values == {"name": "start > one > two 3 > two 4"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "three" + + graph.invoke(None, graph.update_state(history[3].config, None, "__copy__")) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + ("checkpoint", {"values": {}}), + ("task", {"name": "__start__", "result": {"name": "start"}}), + ("checkpoint", {"values": {"name": "start"}}), + ("task", {"name": "one", "result": {"name": "one"}}), + [ + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 1"}}), + ("task", {"name": "two", "result": {"name": "two 2"}}), + ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 2 > two 1 > three"}}, + ), + ], + [ + ("checkpoint", {"values": {"name": "start > one"}}), + ("task", {"name": "two", "result": {"name": "two 3"}}), + ("task", {"name": "two", "result": {"name": "two 4"}}), + [ + [ + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4"}}, + ), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4 > three"}}, + ), + ], + [ + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4"}}, + ), + ("task", {"name": "three", "result": {"name": "three*"}}), + ( + "checkpoint", + { + "values": { + "name": "start > one > two 3 > two 4 > three*" + } + }, + ), + ], + [ + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4"}}, + ), + ("task", {"name": "three", "result": {"name": "three"}}), + ( + "checkpoint", + {"values": {"name": "start > one > two 3 > two 4 > three"}}, + ), + ], + ], + ], + ], + ] From 565d52975cc7da56ca586b49371e8f00b9b87323 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:30:04 +0200 Subject: [PATCH 02/15] Fix test --- libs/langgraph/langgraph/pregel/__init__.py | 96 +++++- libs/langgraph/tests/test_pregel.py | 346 ++++++++----------- libs/langgraph/tests/test_pregel_async.py | 348 ++++++++++++++++++++ 3 files changed, 573 insertions(+), 217 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 4f180566b..a6e2773b9 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1421,7 +1421,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self.channels, checkpoint, ) - values, as_node = updates[0] + values, as_node, _ = updates[0] # no values as END, just clear all tasks if values is None and as_node == END: @@ -1588,22 +1588,32 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) tasks_group_by = defaultdict(list) + user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) + for task in next_tasks.values(): tasks_group_by[task.name].append(task.id) - user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) for item in values: - if not isinstance(item, dict): - continue - - values = item["values"] - as_node = item["as_node"] + if isinstance(item, dict): + values = item["values"] + as_node = item["as_node"] + elif isinstance(item, StateUpdate): + values = item.values + as_node = item.as_node + else: + raise InvalidUpdateError( + f"Invalid update item: {item} when copying checkpoint" + ) user_group = user_group_by[as_node] tasks_group = tasks_group_by[as_node] - + target_idx = len(user_group) - task_id = tasks_group[target_idx] if target_idx < len(tasks_group) else None + task_id = ( + tasks_group[target_idx] + if target_idx < len(tasks_group) + else None + ) user_group_by[as_node].append( StateUpdate(values=values, as_node=as_node, task_id=task_id) @@ -1870,7 +1880,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self.channels, checkpoint, ) - values, as_node = updates[0] + values, as_node, _ = updates[0] # no values, just clear all tasks if values is None and as_node == END: if len(updates) > 1: @@ -1992,24 +2002,84 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) # no values, copy checkpoint - if values is None and as_node == "__copy__": + if as_node == "__copy__": if len(updates) > 1: raise InvalidUpdateError( "Cannot copy checkpoint with multiple updates" ) + if saved is None: + raise InvalidUpdateError("Cannot copy a non-existent checkpoint") + next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput( - saved.parent_config or saved.config if saved else checkpoint_config, + saved.parent_config + or patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ), next_checkpoint, { "source": "fork", "step": step + 1, - "parents": saved.metadata.get("parents", {}) if saved else {}, + "parents": saved.metadata.get("parents", {}), }, {}, ) + + # we want to both clone a checkpoint and update state in one go. + # reuse the same task ID if possible. + if isinstance(values, list) and len(values) > 0: + # figure out the task IDs for the next update checkpoint + next_tasks = prepare_next_tasks( + next_checkpoint, + saved.pending_writes, + self.nodes, + channels, + managed, + next_config, + step + 2, + step + 4, + for_execution=False, + ) + + tasks_group_by = defaultdict(list) + user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) + + for task in next_tasks.values(): + tasks_group_by[task.name].append(task.id) + + for item in values: + if isinstance(item, dict): + values = item["values"] + as_node = item["as_node"] + elif isinstance(item, StateUpdate): + values = item.values + as_node = item.as_node + else: + raise InvalidUpdateError( + f"Invalid update item: {item} when copying checkpoint" + ) + + user_group = user_group_by[as_node] + tasks_group = tasks_group_by[as_node] + + target_idx = len(user_group) + task_id = ( + tasks_group[target_idx] + if target_idx < len(tasks_group) + else None + ) + + user_group_by[as_node].append( + StateUpdate(values=values, as_node=as_node, task_id=task_id) + ) + + return await aperform_superstep( + patch_checkpoint_map(next_config, saved.metadata), + [item for lst in user_group_by.values() for item in lst], + ) + return patch_checkpoint_map( next_config, saved.metadata if saved else None ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b6cc11332..dc732da91 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8051,15 +8051,12 @@ def test_timeout_with_parent_command( def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None: """Test forking and updating task results with state history.""" - two_count = 0 - def checkpoint(snapshot: StateSnapshot) -> tuple[str, dict]: - """Convert a checkpoint to a tuple representation for comparison.""" - return ("checkpoint", {"values": snapshot.values}) + def checkpoint(values: dict[str, Any]): + return ("checkpoint", {"values": values}) - def task(task_info: PregelTask) -> tuple[str, dict]: - """Convert a task to a tuple representation for comparison.""" - return ("task", {"name": task_info.name, "result": task_info.result}) + def task(name: str, result: Any): + return ("task", {"name": name, "result": result}) def get_tree(history: list[StateSnapshot]) -> list: """Build a tree structure from state history for comparison.""" @@ -8081,15 +8078,17 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> node_map[checkpoint_id] = {"item": item, "children": []} parent = node_map.get(parent_checkpoint_id) - if parent: - parent["children"].append(node_map[checkpoint_id]) - else: - root_nodes.append(node_map[checkpoint_id]) + (parent["children"] if parent else root_nodes).append( + node_map[checkpoint_id] + ) def node_to_tree(node: dict) -> list: """Convert a node to tree structure.""" - result = [checkpoint(node["item"])] + [ - task(task_info) for task_info in node["item"].tasks + result = [ + checkpoint(node["item"].values), + ] + [ + task(task_info.name, task_info.result) + for task_info in node["item"].tasks ] if len(node["children"]) > 1: @@ -8100,9 +8099,10 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> else: return result - # Process all root nodes if len(root_nodes) == 1: + # Process all root nodes return node_to_tree(root_nodes[0]) + elif len(root_nodes) > 1: # Multiple root nodes - treat as branches branches = [node_to_tree(node) for node in root_nodes] @@ -8110,20 +8110,21 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> else: return [] + class State(TypedDict): + name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b] + # Define the graph with a sequence of nodes - def one(state: dict) -> Command: + def one(state: State) -> Command: return Command(goto=[Send("two", {})], update={"name": "one"}) - def two(state: dict) -> dict: - nonlocal two_count - two_count += 1 - return {"name": f"two {two_count}"} + def two(state: State) -> State: + return {"name": "two"} - def three(state: dict) -> dict: + def three(state: State) -> State: return {"name": "three"} graph = ( - StateGraph({"name": str}) + StateGraph(State) .add_node("one", one) .add_node("two", two) .add_node("three", three) @@ -8141,16 +8142,16 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> history = list(graph.get_state_history(config)) assert get_tree(history) == [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ] # Update the start state @@ -8158,39 +8159,36 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None, graph.update_state( history[4].config, - [StateUpdate(values={"name": "start*"}, as_node="__start__")], - "__copy__", + values=[StateUpdate(values={"name": "start*"}, as_node="__start__")], + as_node="__copy__", ), ) history = list(graph.get_state_history(config)) assert get_tree(history) == [ [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start*"}}), - ("checkpoint", {"values": {"name": "start*"}}), - ("task", {"name": "one", "result": {"name": "one"}}), - ("checkpoint", {"values": {"name": "start* > one"}}), - ("task", {"name": "two", "result": {"name": "two 3"}}), - ("task", {"name": "two", "result": {"name": "two 4"}}), - ("checkpoint", {"values": {"name": "start* > one > two 4 > two 3"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start* > one > two 4 > two 3 > three"}}, - ), + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), ], ] @@ -8212,54 +8210,44 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> history = list(graph.get_state_history(config)) assert get_tree(history) == [ [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1 > three"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start*"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), [ [ - ("checkpoint", {"values": {"name": "start*"}}), - ("task", {"name": "one", "result": {"name": "one"}}), - ("checkpoint", {"values": {"name": "start* > one"}}), - ("task", {"name": "two", "result": {"name": "two 3"}}), - ("task", {"name": "two", "result": {"name": "two 4"}}), - ( - "checkpoint", - {"values": {"name": "start* > one > two 4 > two 3"}}, - ), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start* > one > two 4 > two 3 > three"}}, - ), + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), ], [ - ("checkpoint", {"values": {"name": "start*"}}), - ("task", {"name": "one", "result": {"name": "one*"}}), - ("checkpoint", {"values": {"name": "start* > one*"}}), - ("task", {"name": "two", "result": {"name": "two 5"}}), - ("checkpoint", {"values": {"name": "start* > one* > two 5"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start* > one* > two 5 > three"}}, - ), + checkpoint({"name": "start*"}), + task("one", {"name": "one*"}), + checkpoint({"name": "start* > one*"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one* > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one* > two > three"}), ], ], ], ] - two_count = 0 config = {"configurable": {"thread_id": "2"}} # Initialize the thread once again @@ -8275,8 +8263,8 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> graph.update_state( history[2].config, [ - StateUpdate(values={"name": "two 3"}, as_node="two"), - StateUpdate(values={"name": "two 4"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), ], "__copy__", ), @@ -8284,38 +8272,32 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> history = list(graph.get_state_history(config)) assert get_tree(history) == [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), [ [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 2 > two 1 > three"}}, - ), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 3"}}), - ("task", {"name": "two", "result": {"name": "two 4"}}), - ("checkpoint", {"values": {"name": "start > one > two 3 > two 4"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4 > three"}}, - ), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], ], ] # Fork task three - assert history[1].values == {"name": "start > one > two 3 > two 4"} + assert history[1].values == {"name": "start > one > two > two"} assert len(history[1].tasks) == 1 assert history[1].tasks[0].name == "three" @@ -8330,52 +8312,33 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> history = list(graph.get_state_history(config)) assert get_tree(history) == [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), [ [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 2 > two 1 > three"}}, - ), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 3"}}), - ("task", {"name": "two", "result": {"name": "two 4"}}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), [ [ - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4"}}, - ), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4 > three"}}, - ), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4"}}, - ), - ("task", {"name": "three", "result": {"name": "three*"}}), - ( - "checkpoint", - { - "values": { - "name": "start > one > two 3 > two 4 > three*" - } - }, - ), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), ], ], ], @@ -8383,7 +8346,7 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> ] # Regenerate task three - assert history[3].values == {"name": "start > one > two 3 > two 4"} + assert history[3].values == {"name": "start > one > two > two"} assert len(history[3].tasks) == 1 assert history[3].tasks[0].name == "three" @@ -8391,63 +8354,38 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> history = list(graph.get_state_history(config)) assert get_tree(history) == [ - ("checkpoint", {"values": {}}), - ("task", {"name": "__start__", "result": {"name": "start"}}), - ("checkpoint", {"values": {"name": "start"}}), - ("task", {"name": "one", "result": {"name": "one"}}), + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), [ [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 1"}}), - ("task", {"name": "two", "result": {"name": "two 2"}}), - ("checkpoint", {"values": {"name": "start > one > two 2 > two 1"}}), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 2 > two 1 > three"}}, - ), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ("checkpoint", {"values": {"name": "start > one"}}), - ("task", {"name": "two", "result": {"name": "two 3"}}), - ("task", {"name": "two", "result": {"name": "two 4"}}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), [ [ - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4"}}, - ), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4 > three"}}, - ), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], [ - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4"}}, - ), - ("task", {"name": "three", "result": {"name": "three*"}}), - ( - "checkpoint", - { - "values": { - "name": "start > one > two 3 > two 4 > three*" - } - }, - ), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), ], [ - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4"}}, - ), - ("task", {"name": "three", "result": {"name": "three"}}), - ( - "checkpoint", - {"values": {"name": "start > one > two 3 > two 4 > three"}}, - ), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), ], ], ], diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b088eb7c5..3ab761db9 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8770,3 +8770,351 @@ async def test_timeout_with_parent_command( await graph.ainvoke({"value": "start"}, thread1) assert exc_info.value.args[0].goto == "test_cmd" assert exc_info.value.args[0].update == {"key": "value"} + + +async def test_fork_and_update_task_results( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test forking and updating task results with state history.""" + + def checkpoint(values: dict[str, Any]): + return ("checkpoint", {"values": values}) + + def task(name: str, result: Any): + return ("task", {"name": name, "result": result}) + + def get_tree(history: list[StateSnapshot]) -> list: + """Build a tree structure from state history for comparison.""" + if not history: + return [] + + # Build a tree structure similar to renderForks + node_map: dict[str, dict] = {} + root_nodes: list[dict] = [] + + # Second pass: establish parent-child relationships + for item in reversed(history): + checkpoint_id = item.config["configurable"]["checkpoint_id"] + parent_checkpoint_id = ( + item.parent_config["configurable"]["checkpoint_id"] + if item.parent_config + else None + ) + node_map[checkpoint_id] = {"item": item, "children": []} + + parent = node_map.get(parent_checkpoint_id) + (parent["children"] if parent else root_nodes).append( + node_map[checkpoint_id] + ) + + def node_to_tree(node: dict) -> list: + """Convert a node to tree structure.""" + result = [ + checkpoint(node["item"].values), + ] + [ + task(task_info.name, task_info.result) + for task_info in node["item"].tasks + ] + + if len(node["children"]) > 1: + branches = [node_to_tree(child) for child in node["children"]] + return result + [branches] + elif len(node["children"]) == 1: + return result + node_to_tree(node["children"][0]) + else: + return result + + if len(root_nodes) == 1: + # Process all root nodes + return node_to_tree(root_nodes[0]) + + elif len(root_nodes) > 1: + # Multiple root nodes - treat as branches + branches = [node_to_tree(node) for node in root_nodes] + return branches + else: + return [] + + class State(TypedDict): + name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b] + + # Define the graph with a sequence of nodes + def one(state: State) -> Command: + return Command(goto=[Send("two", {})], update={"name": "one"}) + + def two(state: State) -> State: + return {"name": "two"} + + def three(state: State) -> State: + return {"name": "three"} + + graph = ( + StateGraph(State) + .add_node("one", one) + .add_node("two", two) + .add_node("three", three) + .add_edge(START, "one") + .add_edge("one", "two") + .add_edge("two", "three") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + history: list[StateSnapshot] = [] + + # Initial run + await graph.ainvoke({"name": "start"}, config) + history = [c async for c in graph.aget_state_history(config)] + + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ] + + # Update the start state + await graph.ainvoke( + None, + await graph.aupdate_state( + history[4].config, + values=[StateUpdate(values={"name": "start*"}, as_node="__start__")], + as_node="__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + ] + + # Fork from task "one" + # Start from the checkpoint that has the task "one" + assert history[3].values == {"name": "start*"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "one" + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[3].config, + [StateUpdate(values={"name": "one*"}, as_node="one")], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + [ + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one*"}), + checkpoint({"name": "start* > one*"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one* > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one* > two > three"}), + ], + ], + ], + ] + + config = {"configurable": {"thread_id": "2"}} + + # Initialize the thread once again + await graph.ainvoke({"name": "start"}, config) + history = [c async for c in graph.aget_state_history(config)] + + # Fork from task "two" + # Start from the checkpoint that has the task "two" + assert history[2].values == {"name": "start > one"} + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[2].config, + [ + StateUpdate(values={"name": "two"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), + ], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ] + + # Fork task three + assert history[1].values == {"name": "start > one > two > two"} + assert len(history[1].tasks) == 1 + assert history[1].tasks[0].name == "three" + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[1].config, + [StateUpdate(values={"name": "three*"}, as_node="three")], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + ], + ], + ], + ] + + # Regenerate task three + assert history[3].values == {"name": "start > one > two > two"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "three" + + await graph.ainvoke( + None, await graph.aupdate_state(history[3].config, None, "__copy__") + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ], + ], + ] From 263a583f01fc31619b8a630898877ea31aec0859 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:32:05 +0200 Subject: [PATCH 03/15] Fix lint --- libs/langgraph/langgraph/pregel/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a6e2773b9..7d5646b85 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance( - self._channels, list - ), "Cannot read additional channels when subscribed to single channels" + assert isinstance(self._channels, list), ( + "Cannot read additional channels when subscribed to single channels" + ) self._channels.extend(channels) return self From adb2eee54b7c821f734d6135f7754dbfbb5998c3 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:45:48 +0200 Subject: [PATCH 04/15] Remove dict --- libs/langgraph/langgraph/pregel/__init__.py | 30 ++++++++------------- 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7d5646b85..d8b2e3f8a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance(self._channels, list), ( - "Cannot read additional channels when subscribed to single channels" - ) + assert isinstance( + self._channels, list + ), "Cannot read additional channels when subscribed to single channels" self._channels.extend(channels) return self @@ -1421,7 +1421,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self.channels, checkpoint, ) - values, as_node, _ = updates[0] + values, as_node = updates[0][:2] # no values as END, just clear all tasks if values is None and as_node == END: @@ -1594,17 +1594,13 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou tasks_group_by[task.name].append(task.id) for item in values: - if isinstance(item, dict): - values = item["values"] - as_node = item["as_node"] - elif isinstance(item, StateUpdate): - values = item.values - as_node = item.as_node - else: + if not isinstance(item, Sequence): raise InvalidUpdateError( f"Invalid update item: {item} when copying checkpoint" ) + values, as_node = item[:2] + user_group = user_group_by[as_node] tasks_group = tasks_group_by[as_node] @@ -1880,7 +1876,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self.channels, checkpoint, ) - values, as_node, _ = updates[0] + values, as_node = updates[0][:2] # no values, just clear all tasks if values is None and as_node == END: if len(updates) > 1: @@ -2012,6 +2008,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint next_config = await checkpointer.aput( saved.parent_config @@ -2050,17 +2047,12 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou tasks_group_by[task.name].append(task.id) for item in values: - if isinstance(item, dict): - values = item["values"] - as_node = item["as_node"] - elif isinstance(item, StateUpdate): - values = item.values - as_node = item.as_node - else: + if not isinstance(item, Sequence): raise InvalidUpdateError( f"Invalid update item: {item} when copying checkpoint" ) + values, as_node = item[:2] user_group = user_group_by[as_node] tasks_group = tasks_group_by[as_node] From 8a67a5ac25be39539071349e5dbd469dda845c10 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:48:25 +0200 Subject: [PATCH 05/15] Fix format --- libs/langgraph/langgraph/pregel/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d8b2e3f8a..d6cc3e432 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance( - self._channels, list - ), "Cannot read additional channels when subscribed to single channels" + assert isinstance(self._channels, list), ( + "Cannot read additional channels when subscribed to single channels" + ) self._channels.extend(channels) return self From 239df52e74a0cf3b88db1ada7a9e5b58540cbec5 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:56:53 +0200 Subject: [PATCH 06/15] Avoid creating new root checkpoint to preserve behaviour --- libs/langgraph/langgraph/pregel/__init__.py | 22 +++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d6cc3e432..fd65810c5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1555,12 +1555,17 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) + copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = checkpointer.put( saved.parent_config - or patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + or ( + patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ) + if copy_and_apply_tasks + else saved.config ), next_checkpoint, { @@ -1573,7 +1578,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. - if isinstance(values, list) and len(values) > 0: + if copy_and_apply_tasks: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, @@ -2008,12 +2013,17 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) + copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = await checkpointer.aput( saved.parent_config - or patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + or ( + patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} + ) + if copy_and_apply_tasks + else saved.config ), next_checkpoint, { @@ -2026,7 +2036,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. - if isinstance(values, list) and len(values) > 0: + if copy_and_apply_tasks: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, From 04b9947a417c70f7b9826a2977f30b994fe1ae4d Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 03:02:53 +0200 Subject: [PATCH 07/15] Fix mypy --- libs/langgraph/langgraph/pregel/__init__.py | 28 +++++++++++++-------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index fd65810c5..a05f3c5aa 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1555,7 +1555,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) - copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = checkpointer.put( @@ -1564,7 +1563,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou patch_configurable( saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ) - if copy_and_apply_tasks + if isinstance(values, list) and len(values) > 0 else saved.config ), next_checkpoint, @@ -1578,18 +1577,23 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. - if copy_and_apply_tasks: + if isinstance(values, list) and len(values) > 0: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, - saved.pending_writes, + saved.pending_writes or [], self.nodes, channels, managed, next_config, step + 2, step + 4, - for_execution=False, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + manager=None, ) tasks_group_by = defaultdict(list) @@ -2013,7 +2017,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) - copy_and_apply_tasks = isinstance(values, list) and len(values) > 0 # copy checkpoint next_config = await checkpointer.aput( @@ -2022,7 +2025,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou patch_configurable( saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ) - if copy_and_apply_tasks + if isinstance(values, list) and len(values) > 0 else saved.config ), next_checkpoint, @@ -2036,18 +2039,23 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. - if copy_and_apply_tasks: + if isinstance(values, list) and len(values) > 0: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, - saved.pending_writes, + saved.pending_writes or [], self.nodes, channels, managed, next_config, step + 2, step + 4, - for_execution=False, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + manager=None, ) tasks_group_by = defaultdict(list) From 06a42aee53eb01e4f14c86fa566a2051e0c61008 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 17:29:26 +0200 Subject: [PATCH 08/15] Remove if/else --- libs/langgraph/langgraph/pregel/__init__.py | 22 +++++++-------------- libs/langgraph/tests/test_large_cases.py | 7 ++++--- 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a05f3c5aa..3df2bef35 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance(self._channels, list), ( - "Cannot read additional channels when subscribed to single channels" - ) + assert isinstance( + self._channels, list + ), "Cannot read additional channels when subscribed to single channels" self._channels.extend(channels) return self @@ -1559,12 +1559,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # copy checkpoint next_config = checkpointer.put( saved.parent_config - or ( - patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} - ) - if isinstance(values, list) and len(values) > 0 - else saved.config + or patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ), next_checkpoint, { @@ -2021,12 +2017,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou # copy checkpoint next_config = await checkpointer.aput( saved.parent_config - or ( - patch_configurable( - saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} - ) - if isinstance(values, list) and len(values) > 0 - else saved.config + or patch_configurable( + saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ), next_checkpoint, { diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 9fd2f40f9..ec1383fc2 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -4301,7 +4301,7 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: ) -def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: +def test_clear_tasks_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -4450,7 +4450,8 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: ) # clear the interrupt and next tasks - tool_two.update_state(thread1, None, as_node="__copy__") + tool_two.update_state(thread1, None) + # interrupt is cleared, next task is kept assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -4481,7 +4482,7 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: created_at=AnyStr(), metadata={ "parents": {}, - "source": "fork", + "source": "update", "step": 1, }, parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config), From c5a851cd916a5974325d4283fb57c5f7f5404add Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 26 Jun 2025 00:24:24 +0200 Subject: [PATCH 09/15] Use as_node=END instead --- libs/langgraph/tests/test_large_cases.py | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index ec1383fc2..ab8752897 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -1677,9 +1677,9 @@ def test_state_graph_packets( # Define decision-making logic def should_continue(data: dict) -> str: - assert data["something_extra"] == "hi there", ( - "nodes can pass extra data to their cond edges, which isn't saved in state" - ) + assert ( + data["something_extra"] == "hi there" + ), "nodes can pass extra data to their cond edges, which isn't saved in state" # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: return [Send("tools", tool_call) for tool_call in tool_calls] @@ -4301,7 +4301,7 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: ) -def test_clear_tasks_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: +def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -4450,27 +4450,18 @@ def test_clear_tasks_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: ) # clear the interrupt and next tasks - tool_two.update_state(thread1, None) + tool_two.update_state(thread1, None, as_node=END) - # interrupt is cleared, next task is kept + # interrupt and unresolved tasks are cleared, finished tasks are kept assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, - next=( - "tool_one", - "tool_two", - ), + next=("tool_one",), tasks=( PregelTask( id=AnyStr(), name="tool_one", path=("__pregel_push", 0, False), ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), ), config={ "configurable": { From 005edb297980e1fa7bbc2d7216d0acda7ce7f3bd Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 26 Jun 2025 01:12:37 +0200 Subject: [PATCH 10/15] Async test --- libs/langgraph/tests/test_pregel_async.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3ab761db9..b2d75456d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -860,7 +860,9 @@ async def test_dynamic_interrupt_subgraph( @NEEDS_CONTEXTVARS -async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None: +async def test_partial_pending_checkpoint( + async_checkpointer: BaseCheckpointSaver, +) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -1017,31 +1019,27 @@ async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None: ) # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node="__copy__") - # interrupt is cleared, next task is kept + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared, finished tasks are kept tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_one", "tool_two"), + next=("tool_one",), tasks=( PregelTask( AnyStr(), "tool_one", (PUSH, 0, False), - result=None, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), + error=None, interrupts=(), + state=None, ), ), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, - "source": "fork", + "source": "update", "step": 1, }, parent_config=( From 38ab90217f0a434840b9c48b883ce3ebe6c61869 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 16:16:38 -0700 Subject: [PATCH 11/15] Fix assertion --- libs/langgraph/langgraph/pregel/__init__.py | 6 +++--- libs/langgraph/tests/test_large_cases.py | 6 +++--- libs/langgraph/tests/test_pregel_async.py | 4 +--- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3df2bef35..8fe9fed2e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -219,9 +219,9 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert isinstance( - self._channels, list - ), "Cannot read additional channels when subscribed to single channels" + assert isinstance(self._channels, list), ( + "Cannot read additional channels when subscribed to single channels" + ) self._channels.extend(channels) return self diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index ab8752897..cb24702ff 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -1677,9 +1677,9 @@ def test_state_graph_packets( # Define decision-making logic def should_continue(data: dict) -> str: - assert ( - data["something_extra"] == "hi there" - ), "nodes can pass extra data to their cond edges, which isn't saved in state" + assert data["something_extra"] == "hi there", ( + "nodes can pass extra data to their cond edges, which isn't saved in state" + ) # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: return [Send("tools", tool_call) for tool_call in tool_calls] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b2d75456d..0a0fe58f5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1042,9 +1042,7 @@ async def test_partial_pending_checkpoint( "source": "update", "step": 1, }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) From abc5a5ff4461a5ba18c9afc0454a2e71f6d2c3da Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 16:19:09 -0700 Subject: [PATCH 12/15] Update --- libs/langgraph/tests/test_pregel_async.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0a0fe58f5..8c09a91a3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1021,7 +1021,7 @@ async def test_partial_pending_checkpoint( # clear the interrupt and next tasks await tool_two.aupdate_state(thread1, None, as_node=END) # interrupt and next tasks are cleared, finished tasks are kept - tup = await tool_two.checkpointer.aget_tuple(thread1) + tup_upd = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, next=("tool_one",), @@ -1035,14 +1035,14 @@ async def test_partial_pending_checkpoint( state=None, ), ), - config=tup.config, - created_at=tup.checkpoint["ts"], + config=tup_upd.config, + created_at=tup_upd.checkpoint["ts"], metadata={ "parents": {}, "source": "update", "step": 1, }, - parent_config=None, + parent_config=tup.config, interrupts=(), ) From 8a5519da295eeed209cabac637879f189b3542eb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 17:01:20 -0700 Subject: [PATCH 13/15] Fix update_state bugs --- libs/langgraph/langgraph/pregel/__init__.py | 59 ++++++++++----------- libs/langgraph/tests/test_large_cases.py | 12 ++--- libs/langgraph/tests/test_pregel_async.py | 17 ++---- 3 files changed, 35 insertions(+), 53 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 8fe9fed2e..40f0212cf 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1443,9 +1443,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 3, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=checkpointer, manager=None, ) # apply null writes @@ -1455,10 +1453,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if w[0] == NULL_TASK_ID ]: apply_writes( - saved.checkpoint, + checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) # apply writes from tasks that already ran @@ -1473,19 +1471,22 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou checkpoint, channels, next_tasks.values(), - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) # save checkpoint next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, None, step), + create_checkpoint(checkpoint, channels, step), { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, - {}, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None @@ -1645,11 +1646,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 3, for_execution=True, store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), + checkpointer=checkpointer, manager=None, ) # apply null writes @@ -1657,10 +1654,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( - saved.checkpoint, + checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) # apply writes @@ -1672,7 +1669,11 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes( - checkpoint, channels, tasks, None, self.trigger_to_nodes + checkpoint, + channels, + tasks, + checkpointer.get_next_version, + self.trigger_to_nodes, ) valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: @@ -1901,9 +1902,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 3, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=checkpointer, manager=None, ) # apply null writes @@ -1913,10 +1912,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if w[0] == NULL_TASK_ID ]: apply_writes( - saved.checkpoint, + checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) # apply writes from tasks that already ran @@ -1931,19 +1930,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou checkpoint, channels, next_tasks.values(), - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) # save checkpoint next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, None, step), + create_checkpoint(checkpoint, channels, step), { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, - {}, + get_new_channel_versions( + checkpoint_previous_versions, checkpoint["channel_versions"] + ), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None @@ -2103,11 +2104,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 3, for_execution=True, store=self.store, - checkpointer=( - self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None - ), + checkpointer=checkpointer, manager=None, ) # apply null writes @@ -2115,10 +2112,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( - saved.checkpoint, + checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], - None, + checkpointer.get_next_version, self.trigger_to_nodes, ) for tid, k, v in saved.pending_writes: diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index cb24702ff..e3108c105 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -4454,15 +4454,9 @@ def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> N # interrupt and unresolved tasks are cleared, finished tasks are kept assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_one",), - tasks=( - PregelTask( - id=AnyStr(), - name="tool_one", - path=("__pregel_push", 0, False), - ), - ), + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), config={ "configurable": { "thread_id": "1", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 8c09a91a3..3fb05ed34 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -40,7 +40,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import InMemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START from langgraph.errors import InvalidUpdateError, NodeInterrupt, ParentCommand from langgraph.func import entrypoint, task from langgraph.graph import END, StateGraph @@ -1023,18 +1023,9 @@ async def test_partial_pending_checkpoint( # interrupt and next tasks are cleared, finished tasks are kept tup_upd = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_one",), - tasks=( - PregelTask( - AnyStr(), - "tool_one", - (PUSH, 0, False), - error=None, - interrupts=(), - state=None, - ), - ), + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), config=tup_upd.config, created_at=tup_upd.checkpoint["ts"], metadata={ From 3ddb6b1477bab5d788cb51a0866f2c9f6c441a03 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 17:05:41 -0700 Subject: [PATCH 14/15] One more --- libs/langgraph/langgraph/pregel/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 40f0212cf..48b46c0d3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2126,7 +2126,11 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes( - checkpoint, channels, tasks, None, self.trigger_to_nodes + checkpoint, + channels, + tasks, + checkpointer.get_next_version, + self.trigger_to_nodes, ) valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: From cee9ac0b7aa647af02b0cd562f98f82de33b8775 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 17:07:26 -0700 Subject: [PATCH 15/15] One more --- libs/langgraph/langgraph/pregel/__init__.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 48b46c0d3..67c938611 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1587,9 +1587,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 4, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=checkpointer, manager=None, ) @@ -2045,9 +2043,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou step + 4, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=checkpointer, manager=None, )