diff --git a/permchain/pregel.py b/permchain/pregel.py index a0765a227..c5d0179b4 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -120,15 +120,6 @@ class PregelInvoke(RunnableBinding): **other_kwargs, ) - def join(self, channels: Sequence[str]) -> PregelInvoke: - joiner = RunnablePassthrough.assign( - **{chan: PregelRead(chan) for chan in channels} - ) - if isinstance(self.bound, RunnablePassthrough): - return PregelInvoke(channels=self.channels, bound=joiner) - else: - return PregelInvoke(channels=self.channels, bound=self.bound | joiner) - def __or__( self, other: Runnable[Any, Other] @@ -152,8 +143,27 @@ class PregelInvoke(RunnableBinding): class PregelBatch(RunnableEach): channel: str + key: Optional[str] + bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) + def join(self, channels: Sequence[str]) -> PregelBatch: + if self.key is None: + raise ValueError( + "Cannot join() additional channels without a key." + " Pass a key arg to Pregel.subscribe_to_each()." + ) + + joiner = RunnablePassthrough.assign( + **{chan: PregelRead(chan) for chan in channels} + ) + if isinstance(self.bound, RunnablePassthrough): + return PregelBatch(channel=self.channel, key=self.key, bound=joiner) + else: + return PregelBatch( + channel=self.channel, key=self.key, bound=self.bound | joiner + ) + def __or__( self, other: Runnable[Any, Other] @@ -161,9 +171,13 @@ class PregelBatch(RunnableEach): | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], ) -> PregelBatch: if isinstance(self.bound, RunnablePassthrough): - return PregelBatch(channel=self.channel, bound=coerce_to_runnable(other)) + return PregelBatch( + channel=self.channel, key=self.key, bound=coerce_to_runnable(other) + ) else: - return PregelBatch(channel=self.channel, bound=self.bound | other) + return PregelBatch( + channel=self.channel, key=self.key, bound=self.bound | other + ) def __ror__( self, @@ -348,9 +362,9 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): ) @classmethod - def subscribe_to_each(cls, inbox: str) -> PregelBatch: + def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> PregelBatch: """Runs process.batch() with the content of inbox each time it is updated.""" - return PregelBatch(channel=inbox) + return PregelBatch(channel=inbox, key=key) @classmethod def send_to( @@ -665,6 +679,8 @@ def _apply_writes_and_prepare_next_tasks( # Here we don't catch EmptyChannelError because the channel # must be intialized if the previous `if` condition is true val = channels[proc.channel].get() + if proc.key is not None: + val = [{proc.key: v} for v in val] tasks.append((proc, val))