This commit is contained in:
Nuno Campos
2025-03-20 08:18:17 -07:00
parent eaa18cc2dd
commit d3bb2b9aa0
3 changed files with 10 additions and 10 deletions
+4 -4
View File
@@ -504,7 +504,7 @@ class Pregel(PregelProtocol):
name: str = "LangGraph"
trigger_to_nodes: dict[str, set[str]] = None
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None
def __init__(
self,
@@ -2749,10 +2749,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)
+3 -3
View File
@@ -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.
+3 -3
View File
@@ -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,
@@ -893,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,
@@ -1036,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,