Change .join() implementation to have non-subscribed channels as input arg

permchain/pregel/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 from collections import defaultdict, deque
+from functools import partial
 from typing import (
     Any,
     AsyncIterator,
@@ -98,7 +99,8 @@ class Channel:
                 {key: channels}
                 if isinstance(channels, str)
                 else {chan: chan for chan in channels},
-            )
+            ),
+            triggers=[channels] if isinstance(channels, str) else channels,
         )
 
     @classmethod
@@ -211,11 +213,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 0,
             )
 
-            def read(chan: str) -> Any:
-                try:
-                    return channels[chan].get()
-                except EmptyChannelError:
-                    return None
+            if not next_tasks:
+                return
+
+            read = partial(_read_channel, channels)
 
             # Similarly to Bulk Synchronous Parallel / Pregel model
             # computation proceeds in steps, while there are channel updates
@@ -309,11 +310,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 0,
             )
 
-            def read(chan: str) -> Any:
-                try:
-                    return channels[chan].get()
-                except EmptyChannelError:
-                    return None
+            if not next_tasks:
+                return
+
+            read = partial(_read_channel, channels)
 
             # Similarly to Bulk Synchronous Parallel / Pregel model
             # computation proceeds in steps, while there are channel updates
@@ -476,6 +476,15 @@ def _interrupt_or_proceed(
         raise TimeoutError(f"Timed out at step {step}")
 
 
+def _read_channel(
+    channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
+) -> Any:
+    try:
+        return channels[chan].get()
+    except EmptyChannelError:
+        return None
+
+
 def _apply_writes_and_prepare_next_tasks(
     processes: Mapping[str, ChannelInvoke | ChannelBatch],
     channels: Mapping[str, BaseChannel],
@@ -514,19 +523,24 @@ def _apply_writes_and_prepare_next_tasks(
     for name, proc in processes.items():
         if isinstance(proc, ChannelInvoke):
             # If any of the channels read by this process were updated
-            if any(chan in updated_channels for chan in proc.channels.values()):
-                # If all channels read by this process have been initialized
+            if any(chan in updated_channels for chan in proc.triggers):
+                # If all channels subscribed by this process have been initialized
                 try:
-                    val = {k: channels[chan].get() for k, chan in proc.channels.items()}
+                    val = {
+                        k: _read_channel(
+                            channels, chan, catch=chan not in proc.triggers
+                        )
+                        for k, chan in proc.channels.items()
+                    }
                 except EmptyChannelError:
                     continue
 
                 # Processes that subscribe to a single keyless channel get
                 # the value directly, instead of a dict
                 if list(proc.channels.keys()) == [None]:
-                    tasks.append((proc, val[None], name))
-                else:
-                    tasks.append((proc, val, name))
+                    val = val[None]
+
+                tasks.append((proc, val, name))
         elif isinstance(proc, ChannelBatch):
             # If the channel read by this process was updated
             if proc.channel in updated_channels:

permchain/pregel/read.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Mapping, Optional, Sequence
+from typing import Any, Callable, List, Mapping, Optional, Sequence
 
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import (
@@ -67,6 +67,10 @@ default_bound = RunnablePassthrough()
 class ChannelInvoke(RunnableBindingBase):
     channels: Mapping[None, str] | Mapping[str, str]
 
+    triggers: List[str] = Field(default_factory=list)
+
+    skip: Optional[Callable[[Any], bool]] = None
+
     bound: Runnable[Any, Any] = Field(default=default_bound)
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
@@ -74,6 +78,7 @@ class ChannelInvoke(RunnableBindingBase):
     def __init__(
         self,
         channels: Mapping[None, str] | Mapping[str, str],
+        triggers: Sequence[str],
         *,
         bound: Optional[Runnable[Any, Any]] = None,
         kwargs: Optional[Mapping[str, Any]] = None,
@@ -82,6 +87,7 @@ class ChannelInvoke(RunnableBindingBase):
     ) -> None:
         super().__init__(
             channels=channels,
+            triggers=triggers,
             bound=bound or default_bound,
             kwargs=kwargs or {},
             config=config,
@@ -92,13 +98,19 @@ class ChannelInvoke(RunnableBindingBase):
         assert isinstance(channels, list) or isinstance(
             channels, tuple
         ), "channels must be a list or tuple"
-        joiner = RunnablePassthrough.assign(
-            **{chan: ChannelRead(chan) for chan in channels}
+        assert all(
+            k is not None for k in self.channels.keys()
+        ), "all channels must be named when using .join()"
+        return ChannelInvoke(
+            channels={
+                **self.channels,
+                **{chan: chan for chan in channels},
+            },
+            triggers=self.triggers,
+            bound=self.bound,
+            kwargs=self.kwargs,
+            config=self.config,
         )
-        if self.bound is default_bound:
-            return ChannelInvoke(channels=self.channels, bound=joiner)
-        else:
-            return ChannelInvoke(channels=self.channels, bound=self.bound | joiner)
 
     def __or__(
         self,
@@ -108,11 +120,21 @@ class ChannelInvoke(RunnableBindingBase):
     ) -> ChannelInvoke:
         if self.bound is default_bound:
             return ChannelInvoke(
-                channels=self.channels, bound=coerce_to_runnable(other)
+                channels=self.channels,
+                triggers=self.triggers,
+                bound=coerce_to_runnable(other),
+                kwargs=self.kwargs,
+                config=self.config,
             )
         else:
-            # delegate to __or__ in self.bound
-            return ChannelInvoke(channels=self.channels, bound=self.bound | other)
+            return ChannelInvoke(
+                channels=self.channels,
+                triggers=self.triggers,
+                # delegate to __or__ in self.bound
+                bound=self.bound | other,
+                kwargs=self.kwargs,
+                config=self.config,
+            )
 
     def __ror__(
         self,