From 565d52975cc7da56ca586b49371e8f00b9b87323 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 25 Jun 2025 02:30:04 +0200 Subject: [PATCH] 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"}), + ], + ], + ], + ], + ]