diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7f15d01a4..477f5a6d1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -77,6 +77,7 @@ from langgraph.pregel.debug import ( print_step_writes, tasks_w_writes, ) +from langgraph.pregel.get_state import assemble_state_snapshot_hierarchy from langgraph.pregel.io import read_channels from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager @@ -89,7 +90,6 @@ from langgraph.pregel.types import ( StreamMode, ) from langgraph.pregel.utils import ( - assemble_state_snapshot_hierarchy, get_new_channel_versions, ) from langgraph.pregel.validate import validate_graph, validate_keys diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 27777bbab..e3e699518 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -283,9 +283,9 @@ def prepare_next_tasks( "langgraph_task_idx": len(tasks), } checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}" if parent_ns - else packet.node + else f"{packet.node}:{packet.id}" ) task_id = str( uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) @@ -293,21 +293,6 @@ def prepare_next_tasks( if for_execution: proc = processes[packet.node] if node := proc.get_node(): - triggers = [TASKS] - metadata = { - "langgraph_step": step, - "langgraph_node": packet.node, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}" - if parent_ns - else f"{packet.node}:{packet.id}" - ) - task_id = str( - uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) - ) writes = deque() tasks.append( PregelExecutableTask( diff --git a/libs/langgraph/langgraph/pregel/get_state.py b/libs/langgraph/langgraph/pregel/get_state.py new file mode 100644 index 000000000..b8efbe6a2 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/get_state.py @@ -0,0 +1,36 @@ +from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR +from langgraph.pregel.types import StateSnapshot + + +def assemble_state_snapshot_hierarchy( + root_checkpoint_ns: str, + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot], +) -> StateSnapshot: + checkpoint_ns_list_to_visit = sorted( + checkpoint_ns_to_state_snapshots.keys(), + key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)), + ) + while checkpoint_ns_list_to_visit: + checkpoint_ns = checkpoint_ns_list_to_visit.pop() + state_snapshot = checkpoint_ns_to_state_snapshots[checkpoint_ns] + *path, subgraph_node = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) + parent_checkpoint_ns = CHECKPOINT_NAMESPACE_SEPARATOR.join(path) + if subgraph_node and ( + parent_state_snapshot := checkpoint_ns_to_state_snapshots.get( + parent_checkpoint_ns + ) + ): + parent_subgraph_snapshots = { + **(parent_state_snapshot.subgraph_state_snapshots or {}), + subgraph_node: state_snapshot, + } + checkpoint_ns_to_state_snapshots[ + parent_checkpoint_ns + ] = checkpoint_ns_to_state_snapshots[parent_checkpoint_ns]._replace( + subgraph_state_snapshots=parent_subgraph_snapshots + ) + + state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None) + if state_snapshot is None: + raise ValueError(f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'") + return state_snapshot diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 9c76e5f7e..d3d0d989f 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,6 +1,4 @@ from langgraph.checkpoint.base import ChannelVersions -from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR -from langgraph.pregel.types import StateSnapshot def get_new_channel_versions( @@ -19,37 +17,3 @@ def get_new_channel_versions( new_versions = current_versions return new_versions - - -def assemble_state_snapshot_hierarchy( - root_checkpoint_ns: str, - checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot], -) -> StateSnapshot: - checkpoint_ns_list_to_visit = sorted( - checkpoint_ns_to_state_snapshots.keys(), - key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)), - ) - while checkpoint_ns_list_to_visit: - checkpoint_ns = checkpoint_ns_list_to_visit.pop() - state_snapshot = checkpoint_ns_to_state_snapshots[checkpoint_ns] - *path, subgraph_node = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) - parent_checkpoint_ns = CHECKPOINT_NAMESPACE_SEPARATOR.join(path) - if subgraph_node and ( - parent_state_snapshot := checkpoint_ns_to_state_snapshots.get( - parent_checkpoint_ns - ) - ): - parent_subgraph_snapshots = { - **(parent_state_snapshot.subgraph_state_snapshots or {}), - subgraph_node: state_snapshot, - } - checkpoint_ns_to_state_snapshots[ - parent_checkpoint_ns - ] = checkpoint_ns_to_state_snapshots[parent_checkpoint_ns]._replace( - subgraph_state_snapshots=parent_subgraph_snapshots - ) - - state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None) - if state_snapshot is None: - raise ValueError(f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'") - return state_snapshot