From 60fc49b448c25847e3119cd0e3f9a549a01923af Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Mon, 17 Mar 2025 21:56:58 -0700
Subject: [PATCH] Speed up prepare_single_task

- sequential(2000) goes from 8.4s to 4.7s
- replace UUID(str).bytes with simpler str.encode()
- find only the first active trigger, instead of the full list
- use a dedicated function for checking active trigger
---
 libs/langgraph/langgraph/pregel/algo.py | 40 ++++++++++++++++++-------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index ba028d7bd..614f61dc8 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -18,7 +18,6 @@ from typing import (
     cast,
     overload,
 )
-from uuid import UUID
 
 from langchain_core.callbacks import Callbacks
 from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
@@ -27,6 +26,7 @@ from langchain_core.runnables.config import RunnableConfig
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
+    ChannelVersions,
     Checkpoint,
     PendingWrite,
     V,
@@ -432,7 +432,7 @@ def prepare_single_task(
 ) -> Union[None, PregelTask, PregelExecutableTask]:
     """Prepares a single task for the next Pregel step, given a task path, which
     uniquely identifies a PUSH or PULL task within the graph."""
-    checkpoint_id = UUID(checkpoint["id"]).bytes
+    checkpoint_id = checkpoint["id"].encode()
     configurable = config.get(CONF, {})
     parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
 
@@ -641,18 +641,18 @@ def prepare_single_task(
         if name not in processes:
             return
         proc = processes[name]
-        version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
+        versions = checkpoint["channel_versions"]
+        version_type = type(next(iter(versions.values()), None))
         null_version = version_type()  # type: ignore[misc]
         if null_version is None:
             return
-        seen = checkpoint["versions_seen"].get(name, {})
         # If any of the channels read by this process were updated
-        if triggers := sorted(
-            chan
-            for chan in proc.triggers
-            if channels[chan].is_available()
-            and checkpoint["channel_versions"].get(chan, null_version)  # type: ignore[operator]
-            > seen.get(chan, null_version)
+        if triggers := _triggers(
+            channels,
+            versions,
+            checkpoint["versions_seen"].get(name),
+            null_version,
+            proc,
         ):
             try:
                 val = next(
@@ -761,6 +761,26 @@ def prepare_single_task(
                 return PregelTask(task_id, name, task_path[:3])
 
 
+def _triggers(
+    channels: Mapping[str, BaseChannel],
+    versions: ChannelVersions,
+    seen: Optional[ChannelVersions],
+    null_version: V,
+    proc: PregelNode,
+) -> Sequence[str]:
+    if seen is None:
+        for chan in proc.triggers:
+            if channels[chan].is_available():
+                return (chan,)
+    else:
+        for chan in proc.triggers:
+            if channels[chan].is_available() and versions.get(
+                chan, null_version
+            ) > seen.get(chan, null_version):  # type: ignore[operator]
+                return (chan,)
+    return EMPTY_SEQ
+
+
 def _scratchpad(
     config: RunnableConfig,
     pending_writes: list[PendingWrite],