x

libs/langgraph/langgraph/pregel/__init__.py

@@ -543,12 +543,6 @@ class Pregel(PregelProtocol):
         self.config_type = config_type
         self.input_model = input_model
         self.config = config
-        # Index from a trigger to nodes that depend on it
-        trigger_to_node = {}
-        for name, node in self.nodes.items():
-            for trigger in node.triggers:
-                trigger_to_node.setdefault(trigger, []).append(name)
-        self.trigger_to_node = trigger_to_node
 
         self.name = name
         if auto_validate:
@@ -2006,7 +2000,12 @@ class Pregel(PregelProtocol):
                 interrupt_after=interrupt_after_,
                 manager=run_manager,
                 debug=debug,
-                triggers_to_nodes=self.trigger_to_node,
+                # `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),
             ) as loop:
                 # create runner
                 runner = PregelRunner(
@@ -2285,6 +2284,7 @@ class Pregel(PregelProtocol):
                         stream.put_nowait, ((), "custom", c)
                     )
                 )
+
             async with AsyncPregelLoop(
                 input,
                 input_model=self.input_model,
@@ -2300,7 +2300,12 @@ class Pregel(PregelProtocol):
                 interrupt_after=interrupt_after_,
                 manager=run_manager,
                 debug=debug,
-                triggers_to_nodes=self.trigger_to_node,
+                # `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),
             ) as loop:
                 # create runner
                 runner = PregelRunner(
@@ -2469,3 +2474,12 @@ class Pregel(PregelProtocol):
             return latest
         else:
             return chunks
+
+
+def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> dict[str, list[str]]:
+    """Index from a trigger to nodes that depend on it."""
+    trigger_to_nodes = {}
+    for name, node in nodes.items():
+        for trigger in node.triggers:
+            trigger_to_nodes.setdefault(trigger, []).append(name)
+    return trigger_to_nodes

libs/langgraph/langgraph/pregel/algo.py

@@ -329,7 +329,6 @@ def apply_writes(
                         max_version,
                         channels[chan],
                     )
-
     # Return managed values writes to be applied externally
     return pending_writes_by_managed, updated_channels
 
@@ -348,6 +347,8 @@ def prepare_next_tasks(
     store: Literal[None] = None,
     checkpointer: Literal[None] = None,
     manager: Literal[None] = None,
+    trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> dict[str, PregelTask]: ...
 
 
@@ -365,6 +366,8 @@ def prepare_next_tasks(
     store: Optional[BaseStore],
     checkpointer: Optional[BaseCheckpointSaver],
     manager: Union[None, ParentRunManager, AsyncParentRunManager],
+    trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> dict[str, PregelExecutableTask]: ...
 
 
@@ -381,8 +384,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,
+    trigger_to_nodes: Optional[Mapping[str, set[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
     """Prepare the set of tasks that will make up the next Pregel step.
 
@@ -398,7 +401,12 @@ 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.
-        triggered_nodes: The set of nodes that were triggered in the previous step
+        trigger_to_nodes: Optional: Mapping of channel names to the set of nodes
+            that are can be triggered by that channel.
+        updated_channels: Optional. Set of channel names that have been updated during
+            the previous step. Using in conjunction with trigger_to_nodes to speed
+            up the process of determining which nodes should be triggered in the next
+            step.
 
     Returns:
         A dictionary of tasks to be executed. The keys are the task ids and the values
@@ -428,11 +436,21 @@ def prepare_next_tasks(
             manager=manager,
         ):
             tasks.append(task)
+    # If updated channels is available, we project only a subset of the nodes.
+    # since only those nodes have been updated?
+    if updated_channels and trigger_to_nodes:
+        triggered_nodes: set[str] = set()
+        # Get all nodes that have triggers associated with an updated channel
+        for channel in updated_channels:
+            if node_ids := trigger_to_nodes.get(channel):
+                triggered_nodes.update(node_ids)
+        # Sort the nodes to ensure deterministic order
+        candidate_nodes: Iterable[str] = sorted(triggered_nodes)
+    else:
+        candidate_nodes = processes.keys()
+
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
-
-    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(

libs/langgraph/langgraph/pregel/loop.py

@@ -209,7 +209,7 @@ class PregelLoop(LoopProtocol):
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
-        triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
+        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     ) -> None:
         super().__init__(
             step=0,
@@ -233,6 +233,7 @@ class PregelLoop(LoopProtocol):
             CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
             or CONFIG_KEY_DEDUPE_TASKS in config[CONF]
         )
+        self.trigger_to_nodes = trigger_to_nodes
         self.debug = debug
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
@@ -279,7 +280,6 @@ class PregelLoop(LoopProtocol):
             if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
             else ()
         )
-        self.triggers_to_nodes = triggers_to_nodes
         self.prev_checkpoint_config = None
 
     def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
@@ -431,7 +431,7 @@ class PregelLoop(LoopProtocol):
                     ),
                 )
             # all tasks have finished
-            mv_writes, _ = apply_writes(
+            mv_writes, updated_channels = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks.values(),
@@ -484,15 +484,6 @@ class PregelLoop(LoopProtocol):
             self.status = "out_of_steps"
             return False
 
-        # 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:
-            # Get all nodes that have triggers associated with an updated channel
-            for channel in updated_channels:
-                if node_ids := self.triggers_to_nodes.get(channel):
-                    triggered_nodes.update(node_ids)
-
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
@@ -506,7 +497,8 @@ class PregelLoop(LoopProtocol):
             manager=self.manager,
             store=self.store,
             checkpointer=self.checkpointer,
-            triggered_nodes=triggered_nodes or None,
+            trigger_to_nodes=self.trigger_to_nodes,
+            updated_channels=updated_channels,
         )
         self.to_interrupt = []
 
@@ -897,7 +889,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
-        triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
+        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     ) -> None:
         super().__init__(
             input,
@@ -914,7 +906,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             interrupt_before=interrupt_before,
             manager=manager,
             debug=debug,
-            triggers_to_nodes=triggers_to_nodes,
+            trigger_to_nodes=trigger_to_nodes,
         )
         self.stack = ExitStack()
         if checkpointer:
@@ -1040,7 +1032,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
-        triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
+        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     ) -> None:
         super().__init__(
             input,
@@ -1057,7 +1049,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             interrupt_before=interrupt_before,
             manager=manager,
             debug=debug,
-            triggers_to_nodes=triggers_to_nodes,
+            trigger_to_nodes=trigger_to_nodes,
         )
         self.stack = AsyncExitStack()
         if checkpointer: