diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 05c157d38..67c938611 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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. @@ -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: @@ -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 @@ -1544,28 +1545,87 @@ 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 or [], + self.nodes, + channels, + managed, + next_config, + step + 2, + step + 4, + for_execution=True, + store=self.store, + checkpointer=checkpointer, + manager=None, + ) + + 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 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] + + 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] @@ -1584,11 +1644,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 @@ -1596,10 +1652,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 @@ -1611,11 +1667,15 @@ 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]] = [] + 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 +1706,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 +1716,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 +1743,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 +1813,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. @@ -1817,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][:2] # no values, just clear all tasks if values is None and as_node == END: if len(updates) > 1: @@ -1837,9 +1900,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 @@ -1849,10 +1910,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 @@ -1867,19 +1928,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 @@ -1939,24 +2002,83 @@ 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 or [], + self.nodes, + channels, + managed, + next_config, + step + 2, + step + 4, + for_execution=True, + store=self.store, + checkpointer=checkpointer, + manager=None, + ) + + 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 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] + + 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 ) @@ -1978,11 +2100,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 @@ -1990,10 +2108,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: @@ -2004,11 +2122,15 @@ 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]] = [] + 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 +2157,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 +2167,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 +2194,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 +2261,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_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 9fd2f40f9..e3108c105 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_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -4450,27 +4450,13 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: ) # clear the interrupt and next tasks - tool_two.update_state(thread1, None, as_node="__copy__") - # interrupt is cleared, next task is kept + tool_two.update_state(thread1, None, as_node=END) + + # 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", - ), - tasks=( - PregelTask( - id=AnyStr(), - name="tool_one", - path=("__pregel_push", 0, False), - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), - ), + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), config={ "configurable": { "thread_id": "1", @@ -4481,7 +4467,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), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bc32e2480..dc732da91 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8047,3 +8047,347 @@ 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.""" + + 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=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({"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 + graph.invoke( + None, + graph.update_state( + history[4].config, + values=[StateUpdate(values={"name": "start*"}, as_node="__start__")], + as_node="__copy__", + ), + ) + + history = list(graph.get_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" + + 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({"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 + 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"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), + ], + "__copy__", + ), + ) + + history = list(graph.get_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" + + 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({"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" + + graph.invoke(None, graph.update_state(history[3].config, None, "__copy__")) + + history = list(graph.get_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"}), + ], + ], + ], + ], + ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b088eb7c5..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 @@ -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,36 +1019,21 @@ 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 - tup = await tool_two.checkpointer.aget_tuple(thread1) + await tool_two.aupdate_state(thread1, None, as_node=END) + # 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", "tool_two"), - tasks=( - PregelTask( - AnyStr(), - "tool_one", - (PUSH, 0, False), - result=None, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), + config=tup_upd.config, + created_at=tup_upd.checkpoint["ts"], metadata={ "parents": {}, - "source": "fork", + "source": "update", "step": 1, }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config - ), + parent_config=tup.config, interrupts=(), ) @@ -8770,3 +8757,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"}), + ], + ], + ], + ], + ]