diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 8de4b87bf..1f79495d0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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" diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index f62f94c09..9bc0f59b4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 495bddc3e..edc7764dd 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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(),