x

libs/langgraph/langgraph/pregel/__init__.py

@@ -1430,7 +1430,7 @@ class Pregel(PregelProtocol):
             if saved and channel_writes:
                 checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
             # apply to checkpoint and save
-            mv_writes, updated_channels = apply_writes(
+            mv_writes, _ = apply_writes(
                 checkpoint, channels, [task], checkpointer.get_next_version
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"
@@ -1716,7 +1716,7 @@ class Pregel(PregelProtocol):
                     checkpoint_config, channel_writes, task_id
                 )
             # apply to checkpoint and save
-            mv_writes, updated_channels = apply_writes(
+            mv_writes, _ = apply_writes(
                 checkpoint, channels, [task], checkpointer.get_next_version
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"

libs/langgraph/langgraph/pregel/algo.py

@@ -381,6 +381,8 @@ def prepare_next_tasks(
     store: Optional[BaseStore] = None,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
+    # Nodes that are known to have been triggered in the previous step
+    triggered_nodes: Optional[set[str]] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
     """Prepare the set of tasks that will make up the next Pregel step.
 
@@ -396,9 +398,7 @@ def prepare_next_tasks(
         store: An instance of BaseStore to make it available for usage within tasks.
         checkpointer: Checkpointer instance used for saving checkpoints.
         manager: The parent run manager to use for the tasks.
-        updated_channels: The set of channels that were updated in the previous step.
-            When available, it allows to efficiently determine which tasks
-            should be executed next instead of having to check all of them.
+        triggered_nodes: The set of nodes that were triggered in the previous step
 
     Returns:
         A dictionary of tasks to be executed. The keys are the task ids and the values
@@ -430,9 +430,11 @@ def prepare_next_tasks(
             tasks.append(task)
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
-    for name in processes:
-        # Check if we know which channels have been updated.
 
+    candidate_nodes: Iterable[str] = processes.keys() if triggered_nodes is None else sorted(triggered_nodes)
+
+    for name in candidate_nodes:
+        # Check if we know which channels have been updated.
         if task := prepare_single_task(
             (PULL, name),
             None,

libs/langgraph/langgraph/pregel/loop.py

@@ -431,7 +431,7 @@ class PregelLoop(LoopProtocol):
                     ),
                 )
             # all tasks have finished
-            mv_writes, updated_channels = apply_writes(
+            mv_writes, _ = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks.values(),
@@ -486,23 +486,18 @@ class PregelLoop(LoopProtocol):
 
         # If updated channels is available, we project only a subset of the nodes.
         # since only those nodes have been updated?
+        triggered_nodes: set[str] = set()
         if updated_channels is not None:
-            candidate_node_ids: set[str] = set()
             # Get all nodes that have triggers associated with an updated channel
             for channel in updated_channels:
-                if triggered_nodes := self.triggers_to_nodes.get(channel):
-                    candidate_node_ids.update(triggered_nodes)
-            candidate_nodes = {
-                node_id: self.nodes[node_id] for node_id in candidate_node_ids
-            }
-        else:
-            candidate_nodes = self.nodes
+                if node_ids := self.triggers_to_nodes.get(channel):
+                    triggered_nodes.update(node_ids)
 
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
             self.checkpoint_pending_writes,
-            candidate_nodes,
+            self.nodes,
             self.channels,
             self.managed,
             self.config,
@@ -511,6 +506,7 @@ class PregelLoop(LoopProtocol):
             manager=self.manager,
             store=self.store,
             checkpointer=self.checkpointer,
+            candidate_node_ids=triggered_nodes or None,
         )
         self.to_interrupt = []
 
@@ -630,7 +626,7 @@ class PregelLoop(LoopProtocol):
         if null_writes := [
             w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
         ]:
-            mv_writes, updated_channels = apply_writes(
+            mv_writes, _ = apply_writes(
                 self.checkpoint,
                 self.channels,
                 [PregelTaskWrites((), INPUT, null_writes, [])],
@@ -679,7 +675,7 @@ class PregelLoop(LoopProtocol):
                 manager=None,
             )
             # apply input writes
-            mv_writes, updated_channels = apply_writes(
+            mv_writes, _ = apply_writes(
                 self.checkpoint,
                 self.channels,
                 [
@@ -794,7 +790,7 @@ class PregelLoop(LoopProtocol):
                 and self.checkpoint_pending_writes
                 and any(task.writes for task in self.tasks.values())
             ):
-                mv_writes, updated_channels = apply_writes(
+                mv_writes, _ = apply_writes(
                     self.checkpoint,
                     self.channels,
                     self.tasks.values(),