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"}}, + ), + ], + ], + ], + ], + ]