diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 4dd69806d..79470680b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c37c95373..4546ab300 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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() diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 1dd604f02..a55941a30 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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():