more code review

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

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(

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

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