diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 4dd69806d..45efe5c17 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: Optional[Mapping[str, Sequence[str]]] = None + def __init__( self, *, @@ -525,6 +527,7 @@ class Pregel(PregelProtocol): config_type: Optional[Type[Any]] = None, input_model: Optional[Type[BaseModel]] = None, config: Optional[RunnableConfig] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, name: str = "LangGraph", ) -> None: self.nodes = nodes @@ -544,25 +547,28 @@ class Pregel(PregelProtocol): self.config_type = config_type self.input_model = input_model self.config = config + self.trigger_to_nodes = trigger_to_nodes self.name = name if auto_validate: self.validate() def get_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError async def aget_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError - def copy(self, update: dict[str, Any] | None = None) -> Self: + def copy(self, update: Optional[dict[str, Any]] = None) -> Self: attrs = {**self.__dict__, **(update or {})} return self.__class__(**attrs) - def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: + def with_config( + self, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Self: return self.copy( {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} ) @@ -577,6 +583,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 +2283,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( @@ -2751,10 +2753,10 @@ class Pregel(PregelProtocol): return chunks -def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, list[str]]: +def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: """Index from a trigger to nodes that depend on it.""" trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list) for name, node in nodes.items(): for trigger in node.triggers: - trigger_to_nodes.setdefault(trigger, []).append(name) - return cast(Mapping[str, list[str]], trigger_to_nodes) + trigger_to_nodes[trigger].append(name) + return dict(trigger_to_nodes) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index df8ec3330..19430f026 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -347,7 +347,7 @@ def prepare_next_tasks( store: Literal[None] = None, checkpointer: Literal[None] = None, manager: Literal[None] = None, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelTask]: ... @@ -366,7 +366,7 @@ def prepare_next_tasks( store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], manager: Union[None, ParentRunManager, AsyncParentRunManager], - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelExecutableTask]: ... @@ -384,7 +384,7 @@ def prepare_next_tasks( store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[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. @@ -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..94faff87f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -210,7 +210,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( step=0, @@ -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]]) -> Optional[set[str]]: # 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: Optional[set[str]] = 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(): @@ -890,7 +893,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -1033,7 +1036,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input,