Use fast path for prepare_next_tasks on input

- When there are no values in checkpoint no need to run through all the PULL candidates
- When there are input writes save updated_channels to use on the next call to prepare_next_tasks

libs/langgraph/langgraph/pregel/__init__.py

@@ -504,6 +504,8 @@ class Pregel(PregelProtocol):
 
     name: str = "LangGraph"
 
+    trigger_to_nodes: dict[str, set[str]] = None
+
     def __init__(
         self,
         *,
@@ -577,6 +579,7 @@ class Pregel(PregelProtocol):
             self.interrupt_after_nodes,
             self.interrupt_before_nodes,
         )
+        self.trigger_to_nodes = _trigger_to_nodes(self.nodes)
         return self
 
     @property
@@ -2276,12 +2279,7 @@ class Pregel(PregelProtocol):
                 interrupt_after=interrupt_after_,
                 manager=run_manager,
                 debug=debug,
-                # `self.nodes` can be modified after creation of `Pregel`. For example,
-                # that's how StateGraph compilation currently works.
-                # For now, we recompute the trigger_to_nodes mapping every time the
-                # loop is created. We could potentially memoize this if it becomes a
-                # performance issue.
-                trigger_to_nodes=_trigger_to_nodes(self.nodes),
+                trigger_to_nodes=self.trigger_to_nodes,
             ) as loop:
                 # create runner
                 runner = PregelRunner(

libs/langgraph/langgraph/pregel/algo.py

@@ -452,6 +452,8 @@ def prepare_next_tasks(
                 triggered_nodes.update(node_ids)
         # Sort the nodes to ensure deterministic order
         candidate_nodes: Iterable[str] = sorted(triggered_nodes)
+    elif not checkpoint["channel_versions"]:
+        candidate_nodes = ()
     else:
         candidate_nodes = processes.keys()
 

libs/langgraph/langgraph/pregel/loop.py

@@ -412,7 +412,7 @@ class PregelLoop(LoopProtocol):
         updated_channels: set[str] | None = None
 
         if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE):
-            self._first(input_keys=input_keys)
+            updated_channels = self._first(input_keys=input_keys)
         elif self.to_interrupt:
             # if we need to interrupt, do so
             self.status = "interrupt_before"
@@ -582,7 +582,7 @@ class PregelLoop(LoopProtocol):
                 else:
                     task.writes.append((k, v))
 
-    def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
+    def _first(self, *, input_keys: Union[str, Sequence[str]]) -> set[str] | None:
         # resuming from previous checkpoint requires
         # - finding a previous checkpoint
         # - receiving None input (outer graph) or RESUMING flag (subgraph)
@@ -599,6 +599,8 @@ class PregelLoop(LoopProtocol):
                 ),
             )
         )
+        # this can be set only when there are input_writes
+        updated_channels: set[str] | None = None
 
         # map command to writes
         if isinstance(self.input, Command):
@@ -668,7 +670,7 @@ class PregelLoop(LoopProtocol):
                 manager=None,
             )
             # apply input writes
-            mv_writes, _ = apply_writes(
+            mv_writes, updated_channels = apply_writes(
                 self.checkpoint,
                 self.channels,
                 [
@@ -698,6 +700,7 @@ class PregelLoop(LoopProtocol):
             self.config = patch_configurable(
                 self.config, {CONFIG_KEY_RESUMING: is_resuming}
             )
+        return updated_channels
 
     def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
         for k, v in self.config["metadata"].items():