Merge pull request #10 from langchain-ai/nc/main-chain

Add an optional "main" chain which has implicit input and output channels

README.md

@@ -12,11 +12,11 @@ from permchain import Pregel, channels
 grow_value = (
     Pregel.subscribe_to("value")
     | (lambda x: x + x)
-    | Pregel.send_to(value=lambda x: x if len(x) < 10 else None)
+    | Pregel.write_to(value=lambda x: x if len(x) < 10 else None)
 )
 
 app = Pregel(
-    grow_value,
+    chains={"grow_value": grow_value},
     channels={"value": channels.LastValue(str)},
     input="value",
     output="value",

examples/draft-revise-loop.py

@@ -4,6 +4,7 @@ from langchain.chat_models.openai import ChatOpenAI
 from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
 from langchain.prompts import SystemMessagePromptTemplate
 from langchain.schema.output_parser import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
 
 from permchain import Pregel, channels
 
@@ -64,28 +65,28 @@ gpt4 = ChatOpenAI(model="gpt-4")
 
 drafter_chain = drafter_prompt | gpt3 | StrOutputParser()
 
-reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
-
 editor_chain = (
     editor_prompt
     | gpt4.bind(functions=editor_functions)
     | JsonOutputFunctionsParser(args_only=False)
 )
 
+reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
+
 # application
 
 drafter = (
     # subscribe to question channel as a dict with a single key, "question"
     Pregel.subscribe_to(["question"])
     | drafter_chain
-    | Pregel.send_to("draft")
+    | Pregel.write_to("draft")
 )
 
 editor = (
     # subscribe to draft channel as a dict with a single key, "draft"
     Pregel.subscribe_to(["draft"])
     | editor_chain
-    | Pregel.send_to(
+    | Pregel.write_to(
         # send to "notes" channel if the editor does not accept the draft
         notes=lambda x: x["arguments"]["notes"]
         if x["name"] == "revise"
@@ -95,23 +96,27 @@ editor = (
 
 reviser = (
     # subscribe to new values of "notes" channel,
-    # and join them with the current values of "question" and "draft"
+    # and join them with the input value (question) and "draft"
     Pregel.subscribe_to(["notes"]).join(["question", "draft"])
     | reviser_chain
-    | Pregel.send_to("draft")
+    | Pregel.write_to("draft")
 )
 
 draft_revise_loop = Pregel(
-    [drafter, reviser, editor],
     channels={
         "question": channels.LastValue(str),
         "draft": channels.LastValue(str),
         "notes": channels.LastValue(str),
     },
-    # output will be a dict with keys "draft" and "notes"
-    output=["draft", "notes"],
+    chains={
+        "drafter": drafter,
+        "editor": editor,
+        "reviser": reviser,
+    },
     # input will be a dict with a single key, "question"
     input=["question"],
+    # output will be a dict with keys "draft" and "notes"
+    output=["draft", "notes"],
     # debug logging
     debug=True,
 )

examples/readme.py

@@ -3,11 +3,11 @@ from permchain import Pregel, channels
 grow_value = (
     Pregel.subscribe_to("value")
     | (lambda x: x + x)
-    | Pregel.send_to(value=lambda x: x if len(x) < 10 else None)
+    | Pregel.write_to(value=lambda x: x if len(x) < 10 else None)
 )
 
 app = Pregel(
-    grow_value,
+    chains={"grow_value": grow_value},
     channels={"value": channels.LastValue(str)},
     input="value",
     output="value",

examples/recursive-web-loader.py

@@ -1,4 +1,5 @@
-from typing import Callable, FrozenSet, Optional, TypedDict
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncGenerator, Callable, FrozenSet, Generator, Optional, TypedDict
 
 import httpx
 from langchain.schema import Document
@@ -10,6 +11,14 @@ from permchain import Pregel, channels
 # Load url with sync httpx client
 
 
+@contextmanager
+def httpx_client() -> Generator[httpx.Client, None, None]:
+    with httpx.HTTPTransport(retries=3) as transport, httpx.Client(
+        transport=transport
+    ) as client:
+        yield client
+
+
 class LoadUrlInput(TypedDict):
     url: str
     visited: FrozenSet[str]
@@ -24,6 +33,14 @@ def load_url(input: LoadUrlInput) -> str:
 # Same as above but with async httpx client
 
 
+@asynccontextmanager
+async def httpx_aclient() -> AsyncGenerator[httpx.AsyncClient, None]:
+    async with httpx.AsyncHTTPTransport(retries=3) as transport, httpx.AsyncClient(
+        transport=transport
+    ) as client:
+        yield client
+
+
 class LoadUrlInputAsync(TypedDict):
     url: str
     visited: FrozenSet[str]
@@ -66,7 +83,7 @@ def recursive_web_loader(
     extractor = extractor or (lambda x: x)
     metadata_extractor = metadata_extractor or _metadata_extractor
     # the main chain that gets executed recursively
-    chain = (
+    visitor = (
         # while there are urls in next_urls
         # run the chain below for each url in next_urls
         # adding the current values of visited set, base_url and httpx client
@@ -75,7 +92,7 @@ def recursive_web_loader(
         )
         # load the url (with sync and async implementations)
         | RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async))
-        | Pregel.send_to(
+        | Pregel.write_to(
             # send this url to the visited set
             visited=lambda x: x["url"],
             # send a new document to the documents stream
@@ -92,31 +109,34 @@ def recursive_web_loader(
                 )
                 if url not in x["visited"] and url != x["url"]
             ],
-            _max_steps=max_depth,
         )
     )
     return Pregel(
-        # use the base_url as the first url to visit
-        Pregel.subscribe_to("base_url") | Pregel.send_to("next_urls"),
-        # add the main chain
-        chain,
+        chains={
+            # use the base_url as the first url to visit
+            "input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"),
+            # add the main chain
+            "visitor": visitor,
+        },
         # define the channels
         channels={
             "base_url": channels.LastValue(str),
             "next_urls": channels.UniqueInbox(str),
             "documents": channels.Stream(Document),
             "visited": channels.Set(str),
-            "client": channels.ContextManager(httpx.Client, httpx.AsyncClient),
+            "client": channels.ContextManager(httpx_client, httpx_aclient),
         },
         # this will accept a string as input
         input="base_url",
         # and return a dict with documents and visited set
         output=["documents", "visited"],
-    )
+        # debug logging
+        debug=True,
+    ).with_config({"recursion_limit": max_depth + 1})
 
 
 loader = recursive_web_loader(max_depth=3)
 
 documents = loader.invoke("https://docs.python.org/3.9/")
 
-print(documents)
+print(len(documents["documents"]))

permchain/pregel/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
-import logging
 from collections import defaultdict, deque
 from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast
 
@@ -11,7 +10,7 @@ from langchain.callbacks.manager import (
     CallbackManagerForChainRun,
 )
 from langchain.globals import get_debug
-from langchain.pydantic_v1 import BaseModel, create_model
+from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator
 from langchain.schema.runnable import (
     Runnable,
     RunnablePassthrough,
@@ -23,7 +22,6 @@ from langchain.schema.runnable.config import (
     get_executor_for_config,
     patch_config,
 )
-from langchain.utils.input import get_bolded_text, get_colored_text
 
 from permchain.channels.base import (
     AsyncChannelsManager,
@@ -31,18 +29,19 @@ from permchain.channels.base import (
     ChannelsManager,
     EmptyChannelError,
 )
-from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP
+from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
+from permchain.pregel.debug import print_checkpoint, print_step_start
+from permchain.pregel.io import map_input, map_output
+from permchain.pregel.log import logger
 from permchain.pregel.read import PregelBatch, PregelInvoke
 from permchain.pregel.validate import validate_chains_channels
 from permchain.pregel.write import PregelSink
 
-logger = logging.getLogger(__name__)
-
 
 class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     channels: Mapping[str, Channel]
 
-    chains: Sequence[PregelInvoke | PregelBatch]
+    chains: Mapping[str, PregelInvoke | PregelBatch]
 
     output: str | Sequence[str]
 
@@ -50,37 +49,17 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
     step_timeout: Optional[float] = None
 
-    debug: bool
+    debug: bool = Field(default_factory=get_debug)
 
     class Config:
         arbitrary_types_allowed = True
 
-    def __init__(
-        self,
-        *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch,
-        channels: Mapping[str, Channel],
-        output: str | Sequence[str],
-        input: str | Sequence[str],
-        step_timeout: Optional[float] = None,
-        debug: Optional[bool] = None,
-    ) -> None:
-        chains_flat: list[PregelInvoke | PregelBatch] = []
-        for chain in chains:
-            if isinstance(chain, (list, tuple)):
-                chains_flat.extend(chain)
-            else:
-                chains_flat.append(cast(PregelInvoke | PregelBatch, chain))
-
-        validate_chains_channels(chains_flat, channels, input, output)
-
-        super().__init__(
-            chains=chains_flat,
-            channels=channels,
-            output=output,
-            input=input,
-            step_timeout=step_timeout,
-            debug=debug if debug is not None else get_debug(),
+    @root_validator(skip_on_failure=True)
+    def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
+        validate_chains_channels(
+            values["chains"], values["channels"], values["input"], values["output"]
         )
+        return values
 
     @property
     def InputType(self) -> Any:
@@ -136,10 +115,9 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         return PregelBatch(channel=inbox, key=key)
 
     @classmethod
-    def send_to(
+    def write_to(
         cls,
         *channels: str,
-        _max_steps: Optional[int] = None,
         **kwargs: RunnableLike,
     ) -> PregelSink:
         """Writes to channels the result of the lambda, or None to skip writing."""
@@ -147,8 +125,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             channels=(
                 [(c, RunnablePassthrough()) for c in channels]
                 + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
-            ),
-            max_steps=_max_steps,
+            )
         )
 
     def _transform(
@@ -157,7 +134,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         run_manager: CallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> Iterator[dict[str, Any] | Any]:
-        processes = tuple(self.chains)
+        processes = {**self.chains}
         # TODO this is where we'd restore from checkpoint
         with ChannelsManager(self.channels) as channels, get_executor_for_config(
             config
@@ -165,14 +142,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
-                deque((self.input, chunk) for chunk in input)
-                if isinstance(self.input, str)
-                else deque(
-                    (k, v)
-                    for chunk in input
-                    for k, v in chunk.items()
-                    if k in self.input
-                ),
+                deque(w for c in input for w in map_input(self.input, c)),
             )
 
             def read(chan: str) -> Any:
@@ -188,16 +158,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             # with channel updates applied only at the transition between steps
             for step in range(config["recursion_limit"]):
                 if self.debug:
-                    from pprint import pformat
-
-                    n_tasks = len(next_tasks)
-                    print(
-                        f"{get_colored_text('[pregel/step]', color='blue')} "
-                        + get_bolded_text(
-                            f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n"
-                        )
-                        + pformat({k: read(k) for k in channels})
-                    )
+                    print_step_start(step, next_tasks)
 
                 # collect all writes to channels, without applying them yet
                 pending_writes = deque[tuple[str, Any]]()
@@ -216,46 +177,29 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: pending_writes.extend,
                                     CONFIG_KEY_READ: read,
-                                    CONFIG_KEY_STEP: step,
                                 },
                             ),
                         )
-                        for proc, input in next_tasks
+                        for proc, input, _ in next_tasks
                     ),
                     return_when=concurrent.futures.FIRST_EXCEPTION,
                     timeout=self.step_timeout,
                 )
 
-                while done:
-                    # if any task failed
-                    if exc := done.pop().exception():
-                        # cancel all pending tasks
-                        while inflight:
-                            inflight.pop().cancel()
-                        # raise the exception
-                        raise exc
-                        # TODO this is where retry of an entire step would happen
-
-                if inflight:
-                    # if we got here means we timed out
-                    while inflight:
-                        # cancel all pending tasks
-                        inflight.pop().cancel()
-                    # raise timeout error
-                    raise TimeoutError(f"Timed out at step {step}")
+                # interrupt on failure or timeout
+                _interrupt_or_proceed(done, inflight, step)
 
                 # apply writes to channels, decide on next step
                 next_tasks = _apply_writes_and_prepare_next_tasks(
                     processes, channels, pending_writes
                 )
 
-                # if any write to output channel in this step, yield current value
-                if isinstance(self.output, str):
-                    if any(chan is self.output for chan, _ in pending_writes):
-                        yield channels[self.output].get()
-                else:
-                    if updated := {c for c, _ in pending_writes if c in self.output}:
-                        yield {chan: channels[chan].get() for chan in updated}
+                if self.debug:
+                    print_checkpoint(step, channels)
+
+                # if any write to output channels in this step, yield current value
+                for output in map_output(self.output, pending_writes, channels):
+                    yield output
 
                 # TODO this is where we'd save checkpoint
 
@@ -269,22 +213,13 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         run_manager: AsyncCallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> AsyncIterator[dict[str, Any] | Any]:
-        processes = tuple(self.chains)
+        processes = {**self.chains}
         # TODO this is where we'd restore from checkpoint
         async with AsyncChannelsManager(self.channels) as channels:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
-                deque([(self.input, chunk) async for chunk in input])
-                if isinstance(self.input, str)
-                else deque(
-                    [
-                        (k, v)
-                        async for chunk in input
-                        for k, v in chunk.items()
-                        if k in self.input
-                    ]
-                ),
+                deque([w async for c in input for w in map_input(self.input, c)]),
             )
 
             def read(chan: str) -> Any:
@@ -300,16 +235,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             # channel updates being applied only at the transition between steps
             for step in range(config["recursion_limit"]):
                 if self.debug:
-                    from pprint import pformat
-
-                    n_tasks = len(next_tasks)
-                    print(
-                        f"{get_colored_text('[pregel/step]', color='blue')} "
-                        + get_bolded_text(
-                            f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n"
-                        )
-                        + pformat({k: read(k) for k in channels})
-                    )
+                    print_step_start(step, next_tasks)
 
                 # collect all writes to channels, without applying them yet
                 pending_writes = deque[tuple[str, Any]]()
@@ -330,47 +256,32 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                                         # deque.extend is thread-safe
                                         CONFIG_KEY_SEND: pending_writes.extend,
                                         CONFIG_KEY_READ: read,
-                                        CONFIG_KEY_STEP: step,
                                     },
                                 ),
                             )
                         )
-                        for proc, input in next_tasks
+                        for proc, input, _ in next_tasks
                     ],
                     return_when=asyncio.FIRST_EXCEPTION,
                     timeout=self.step_timeout,
                 )
 
-                while done:
-                    # if any task failed
-                    if exc := done.pop().exception():
-                        # cancel all pending tasks
-                        while inflight:
-                            inflight.pop().cancel()
-                        # raise the exception
-                        raise exc
-                        # TODO this is where retry of an entire step would happen
-
-                if inflight:
-                    # if we got here means we timed out
-                    while inflight:
-                        # cancel all pending tasks
-                        inflight.pop().cancel()
-                    # raise timeout error
-                    raise TimeoutError(f"Timed out at step {step}")
+                # interrupt on failure or timeout
+                _interrupt_or_proceed(done, inflight, step)
 
                 # apply writes to channels, decide on next step
                 next_tasks = _apply_writes_and_prepare_next_tasks(
                     processes, channels, pending_writes
                 )
 
-                # if any write to output channel in this step, yield current value
-                if isinstance(self.output, str):
-                    if any(chan is self.output for chan, _ in pending_writes):
-                        yield channels[self.output].get()
-                else:
-                    if updated := {c for c, _ in pending_writes if c in self.output}:
-                        yield {chan: channels[chan].get() for chan in updated}
+                if self.debug:
+                    print_checkpoint(step, channels)
+
+                # if any write to output channels in this step, yield current value
+                for output in map_output(self.output, pending_writes, channels):
+                    yield output
+
+                # TODO this is where we'd save checkpoint
 
                 # if no more tasks, we're done
                 if not next_tasks:
@@ -440,11 +351,35 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             yield chunk
 
 
+def _interrupt_or_proceed(
+    done: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
+    inflight: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
+    step: int,
+) -> None:
+    while done:
+        # if any task failed
+        if exc := done.pop().exception():
+            # cancel all pending tasks
+            while inflight:
+                inflight.pop().cancel()
+            # raise the exception
+            raise exc
+            # TODO this is where retry of an entire step would happen
+
+    if inflight:
+        # if we got here means we timed out
+        while inflight:
+            # cancel all pending tasks
+            inflight.pop().cancel()
+        # raise timeout error
+        raise TimeoutError(f"Timed out at step {step}")
+
+
 def _apply_writes_and_prepare_next_tasks(
-    processes: Sequence[PregelInvoke | PregelBatch],
+    processes: Mapping[str, PregelInvoke | PregelBatch],
     channels: Mapping[str, Channel],
     pending_writes: Sequence[tuple[str, Any]],
-) -> list[tuple[Runnable, Any]]:
+) -> list[tuple[Runnable, Any, str]]:
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
     # Group writes by channel
     for chan, val in pending_writes:
@@ -459,10 +394,10 @@ def _apply_writes_and_prepare_next_tasks(
         else:
             logger.warning(f"Skipping write for channel {chan} which has no readers")
 
-    tasks: list[tuple[Runnable, Any]] = []
+    tasks: list[tuple[Runnable, Any, str]] = []
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
-    for proc in processes:
+    for name, proc in processes.items():
         if isinstance(proc, PregelInvoke):
             # If any of the channels read by this process were updated
             if any(chan in updated_channels for chan in proc.channels.values()):
@@ -475,9 +410,9 @@ def _apply_writes_and_prepare_next_tasks(
                 # 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]))
+                    tasks.append((proc, val[None], name))
                 else:
-                    tasks.append((proc, val))
+                    tasks.append((proc, val, name))
         elif isinstance(proc, PregelBatch):
             # If the channel read by this process was updated
             if proc.channel in updated_channels:
@@ -487,6 +422,6 @@ def _apply_writes_and_prepare_next_tasks(
                 if proc.key is not None:
                     val = [{proc.key: v} for v in val]
 
-                tasks.append((proc, val))
+                tasks.append((proc, val, name))
 
     return tasks

permchain/pregel/constants.py

@@ -1,3 +1,2 @@
-CONFIG_KEY_STEP = "__pregel_step"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"

permchain/pregel/debug.py

@@ -0,0 +1,34 @@
+from pprint import pformat
+from typing import Any, Iterator, Mapping
+
+from langchain.schema.runnable import Runnable
+from langchain.utils.input import get_bolded_text, get_colored_text
+
+from permchain.channels.base import Channel, EmptyChannelError
+
+
+def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> None:
+    n_tasks = len(next_tasks)
+    print(
+        f"{get_colored_text('[pregel/step]', color='blue')} "
+        + get_bolded_text(
+            f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n"
+        )
+        + "\n".join(f"- {name}({pformat(val)})" for _, val, name in next_tasks)
+    )
+
+
+def print_checkpoint(step: int, channels: Mapping[str, Channel]) -> None:
+    print(
+        f"{get_colored_text('[pregel/checkpoint]', color='blue')} "
+        + get_bolded_text(f"Finishing step {step}. Channel values:\n")
+        + pformat({name: val for name, val in _read_channels(channels)}, depth=1)
+    )
+
+
+def _read_channels(channels: Mapping[str, Channel]) -> Iterator[tuple[str, Any]]:
+    for name, channel in channels.items():
+        try:
+            yield (name, channel.get())
+        except EmptyChannelError:
+            pass

permchain/pregel/io.py

@@ -0,0 +1,34 @@
+from typing import Any, Iterator, Mapping, Sequence
+
+from permchain.channels.base import Channel
+from permchain.pregel.log import logger
+
+
+def map_input(
+    input_channels: str | Sequence[str], chunk: dict[str, Any] | Any
+) -> Iterator[tuple[str, Any]]:
+    """Map input chunk to a sequence of pending writes in the form (channel, value)."""
+    if isinstance(input_channels, str):
+        yield (input_channels, chunk)
+    else:
+        if not isinstance(chunk, dict):
+            raise TypeError(f"Expected chunk to be a dict, got {type(chunk).__name__}")
+        for k in chunk:
+            if k in input_channels:
+                yield (k, chunk[k])
+            else:
+                logger.warning(f"Input channel {k} not found in {input_channels}")
+
+
+def map_output(
+    output_channels: str | Sequence[str],
+    pending_writes: Sequence[tuple[str, Any]],
+    channels: Mapping[str, Channel],
+) -> Iterator[dict[str, Any] | Any]:
+    """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
+    if isinstance(output_channels, str):
+        if any(chan == output_channels for chan, _ in pending_writes):
+            yield channels[output_channels].get()
+    else:
+        if updated := {c for c, _ in pending_writes if c in output_channels}:
+            yield {chan: channels[chan].get() for chan in updated}

permchain/pregel/log.py

@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger(__name__)

permchain/pregel/validate.py

@@ -5,13 +5,13 @@ from permchain.pregel.read import PregelBatch, PregelInvoke
 
 
 def validate_chains_channels(
-    chains: Sequence[PregelInvoke | PregelBatch],
+    chains: Mapping[str, PregelInvoke | PregelBatch],
     channels: Mapping[str, Channel],
     input: str | Sequence[str],
     output: str | Sequence[str],
 ) -> None:
     subscribed_channels = set[str]()
-    for chain in chains:
+    for chain in chains.values():
         if isinstance(chain, PregelInvoke):
             subscribed_channels.update(chain.channels.values())
         elif isinstance(chain, PregelBatch):

permchain/pregel/write.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, Sequence
 
 from langchain.schema.runnable import (
     Runnable,
@@ -9,7 +9,7 @@ from langchain.schema.runnable import (
 )
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
-from permchain.pregel.constants import CONFIG_KEY_SEND, CONFIG_KEY_STEP
+from permchain.pregel.constants import CONFIG_KEY_SEND
 
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
@@ -22,28 +22,17 @@ class PregelSink(RunnableLambda):
     or None to skip writing.
     """
 
-    max_steps: Optional[int]
-
     def __init__(
         self,
         *,
         channels: Sequence[tuple[str, Runnable]],
-        max_steps: Optional[int] = None,
     ):
         super().__init__(func=self._write, afunc=self._awrite)  # type: ignore[arg-type]
         self.channels = channels
-        self.max_steps = max_steps
 
     @property
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
         return [
-            ConfigurableFieldSpec(
-                id=CONFIG_KEY_STEP,
-                name=CONFIG_KEY_STEP,
-                description=None,
-                default=None,
-                annotation=int,
-            ),
             ConfigurableFieldSpec(
                 id=CONFIG_KEY_SEND,
                 name=CONFIG_KEY_SEND,
@@ -54,11 +43,6 @@ class PregelSink(RunnableLambda):
         ]
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
-        step: int = config["configurable"][CONFIG_KEY_STEP]
-
-        if self.max_steps is not None and step >= self.max_steps:
-            return
-
         write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
 
         values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
@@ -68,11 +52,6 @@ class PregelSink(RunnableLambda):
         return input
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
-        step: int = config["configurable"][CONFIG_KEY_STEP]
-
-        if self.max_steps is not None and step >= self.max_steps:
-            return
-
         write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
 
         values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels]

tests/test_pregel.py

@@ -8,15 +8,16 @@ from langchain.schema.runnable import RunnablePassthrough
 from pytest_mock import MockerFixture
 
 from permchain import Pregel, channels
-from permchain.pregel import PregelInvoke
 
 
 def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -32,10 +33,12 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
 
 def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -55,10 +58,12 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
 
 def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -82,11 +87,11 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
 
 def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -101,11 +106,11 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -124,15 +129,14 @@ def test_batch_two_processes_in_out() -> None:
         return inp + 1
 
     chain_one = (
-        Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one")
+        Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one")
     )
     chain_two = (
-        Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output")
+        Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output")
     )
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -154,17 +158,15 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
         "output": channels.LastValue(int),
         "-1": channels.LastValue(int),
     }
-    chains: list[PregelInvoke] = [
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
-    ]
+    chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
-            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
+        chains[str(i)] = (
+            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
         )
-    chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
+    chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
 
-    app = Pregel(*chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains, channels=chans, input="input", output="output")
 
     for _ in range(10):
         assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
@@ -184,17 +186,15 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
         "output": channels.LastValue(int),
         "-1": channels.LastValue(int),
     }
-    chains: list[PregelInvoke] = [
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
-    ]
+    chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
-            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
+        chains[str(i)] = (
+            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
         )
-    chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
+    chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
 
-    app = Pregel(*chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains, channels=chans, input="input", output="output")
 
     for _ in range(10):
         assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
@@ -218,12 +218,11 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
 def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
-    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -240,12 +239,11 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
 def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
-    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.Inbox(int),
@@ -262,14 +260,16 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_three = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_three,
-        chain_four,
+        chains={
+            "chain_one": chain_one,
+            "chain_three": chain_three,
+            "chain_four": chain_four,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -294,7 +294,9 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
     add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
 
     inner_app = Pregel(
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("output"),
+        chains={
+            "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -304,20 +306,22 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
     )
 
     chain_one = (
-        Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map()
+        Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map()
     )
     chain_two = (
         Pregel.subscribe_to("inbox_one")
         | inner_app.map()
         | sorted
-        | Pregel.send_to("outbox_one")
+        | Pregel.write_to("outbox_one")
     )
-    chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output")
+    chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
-        chain_three,
+        chains={
+            "chain_one": chain_one,
+            "chain_two": chain_two,
+            "chain_three": chain_three,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -341,13 +345,12 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
     chain_one = (
         Pregel.subscribe_to("input")
         | add_one
-        | Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough())
+        | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
     )
-    chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output")
+    chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -362,12 +365,11 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
 
 def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("between")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between")
     chain_two = Pregel.subscribe_to("between") | add_one
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -385,13 +387,12 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
 def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chain_one = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
     chain_two = Pregel.subscribe_to("between") | add_one
 
     with pytest.raises(ValueError):
         Pregel(
-            chain_one,
-            chain_two,
+            chains={"chain_one": chain_one, "chain_two": chain_two},
             channels={
                 "input": channels.LastValue(int),
                 "output": channels.LastValue(int),
@@ -415,11 +416,11 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
             cleanup()
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),

tests/test_pregel_async.py

@@ -7,15 +7,16 @@ from langchain.schema.runnable import RunnablePassthrough
 from pytest_mock import MockerFixture
 
 from permchain import Pregel, channels
-from permchain.pregel import PregelInvoke
 
 
 async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -29,10 +30,12 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
 
 async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -52,10 +55,12 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
 
 async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -79,11 +84,11 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
 
 async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -98,11 +103,11 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     pubsub = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -122,15 +127,14 @@ async def test_batch_two_processes_in_out() -> None:
         return inp + 1
 
     chain_one = (
-        Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one")
+        Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one")
     )
     chain_two = (
-        Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output")
+        Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output")
     )
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -152,17 +156,15 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
         "output": channels.LastValue(int),
         "-1": channels.LastValue(int),
     }
-    chains: list[PregelInvoke] = [
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
-    ]
+    chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
-            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
+        chains[str(i)] = (
+            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
         )
-    chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
+    chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
 
-    app = Pregel(*chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains, channels=chans, input="input", output="output")
 
     # No state is left over from previous invocations
     for _ in range(10):
@@ -183,17 +185,15 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
         "output": channels.LastValue(int),
         "-1": channels.LastValue(int),
     }
-    chains: list[PregelInvoke] = [
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
-    ]
+    chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
-            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
+        chains[str(i)] = (
+            Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
         )
-    chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
+    chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
 
-    app = Pregel(*chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains, channels=chans, input="input", output="output")
 
     # No state is left over from previous invocations
     for _ in range(10):
@@ -223,12 +223,11 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
 ) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
-    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -245,12 +244,11 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
 async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
-    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+    chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.Inbox(int),
@@ -267,14 +265,16 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
 
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_three = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_three,
-        chain_four,
+        chains={
+            "chain_one": chain_one,
+            "chain_three": chain_three,
+            "chain_four": chain_four,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -300,7 +300,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
     add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
 
     inner_app = Pregel(
-        Pregel.subscribe_to("input") | add_one | Pregel.send_to("output"),
+        chains={
+            "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -310,20 +312,22 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
     )
 
     chain_one = (
-        Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map()
+        Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map()
     )
     chain_two = (
         Pregel.subscribe_to("inbox_one")
         | inner_app.map()
         | sorted
-        | Pregel.send_to("outbox_one")
+        | Pregel.write_to("outbox_one")
     )
-    chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output")
+    chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
-        chain_three,
+        chains={
+            "chain_one": chain_one,
+            "chain_two": chain_two,
+            "chain_three": chain_three,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -349,13 +353,12 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
     chain_one = (
         Pregel.subscribe_to("input")
         | add_one
-        | Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough())
+        | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
     )
-    chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output")
+    chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -371,12 +374,11 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
 
 async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("between")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between")
     chain_two = Pregel.subscribe_to("between") | add_one
 
     app = Pregel(
-        chain_one,
-        chain_two,
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -415,11 +417,11 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
             cleanup_async()
 
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
-    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
+    chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
+    chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
 
     app = Pregel(
-        [chain_one, chain_two],
+        chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),