diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 256ddbd6c..4dd69806d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2751,10 +2751,10 @@ class Pregel(PregelProtocol): return chunks -def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> dict[str, list[str]]: +def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, list[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 trigger_to_nodes + return cast(Mapping[str, list[str]], trigger_to_nodes) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index cc9a45e59..5574f4777 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, set[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[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, set[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[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, set[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[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. diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 518372834..0cd109f23 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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, - trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, ) -> None: super().__init__( step=0, @@ -889,7 +889,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, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, ) -> None: super().__init__( input, @@ -1032,7 +1032,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, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, ) -> None: super().__init__( input,