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

examples/draft-revise-loop.py

@@ -4,8 +4,10 @@ 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
+from permchain.pregel import PregelIO
 
 # prompts
 
@@ -62,9 +64,12 @@ gpt4 = ChatOpenAI(model="gpt-4")
 
 # chains
 
-drafter_chain = drafter_prompt | gpt3 | StrOutputParser()
-
-reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
+drafter_chain = (
+    {"question": RunnablePassthrough(input_type=str)}
+    | drafter_prompt
+    | gpt3
+    | StrOutputParser()
+)
 
 editor_chain = (
     editor_prompt
@@ -72,61 +77,51 @@ editor_chain = (
     | 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")
-)
-
-editor = (
-    # subscribe to draft channel as a dict with a single key, "draft"
-    Pregel.subscribe_to(["draft"])
-    | editor_chain
-    | Pregel.send_to(
-        # send to "notes" channel if the editor does not accept the draft
-        notes=lambda x: x["arguments"]["notes"]
-        if x["name"] == "revise"
-        else None
-    )
-)
-
-reviser = (
-    # subscribe to new values of "notes" channel,
-    # and join them with the current values of "question" and "draft"
-    Pregel.subscribe_to(["notes"]).join(["question", "draft"])
-    | reviser_chain
-    | Pregel.send_to("draft")
-)
-
 draft_revise_loop = Pregel(
-    [drafter, reviser, editor],
+    drafter_chain | Pregel.send_to("draft"),
+    chains={
+        "editor": (
+            # subscribe to draft channel as a dict with a single key, "draft"
+            Pregel.subscribe_to(["draft"])
+            | editor_chain
+            | Pregel.send_to(
+                # send to "notes" channel if the editor does not accept the draft
+                notes=lambda x: x["arguments"]["notes"]
+                if x["name"] == "revise"
+                else None
+            )
+        ),
+        "reviser": (
+            # subscribe to new values of "notes" channel,
+            # and join them with the input value (question) and "draft"
+            Pregel.subscribe_to(["notes"]).join([PregelIO.IN, "draft"])
+            | reviser_chain
+            | Pregel.send_to("draft")
+        ),
+    },
     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"],
-    # input will be a dict with a single key, "question"
-    input=["question"],
     # debug logging
     debug=True,
 )
 
 # run
 
-for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}):
+for draft in draft_revise_loop.stream("What food do turtles eat?"):
     print(draft)
     print("---")
 
 
 async def main():
-    async for draft in draft_revise_loop.astream(
-        {"question": "What food do turtles eat?"}
-    ):
+    async for draft in draft_revise_loop.astream("What food do turtles eat?"):
         print(draft)
         print("---")
 

examples/readme.py

@@ -7,7 +7,7 @@ grow_value = (
 )
 
 app = Pregel(
-    grow_value,
+    chains={"grow_value": grow_value},
     channels={"value": channels.LastValue(str)},
     input="value",
     output="value",

examples/recursive-web-loader.py

@@ -66,7 +66,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
@@ -96,10 +96,12 @@ def recursive_web_loader(
         )
     )
     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.send_to("next_urls"),
+            # add the main chain
+            "visitor": visitor,
+        },
         # define the channels
         channels={
             "base_url": channels.LastValue(str),

permchain/pregel/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import enum
 import logging
 from collections import defaultdict, deque
 from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast
@@ -23,7 +24,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,7 +31,14 @@ from permchain.channels.base import (
     ChannelsManager,
     EmptyChannelError,
 )
-from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP
+from permchain.channels.last_value import LastValue
+from permchain.pregel.constants import (
+    CHAINS_MAIN,
+    CONFIG_KEY_READ,
+    CONFIG_KEY_SEND,
+    CONFIG_KEY_STEP,
+)
+from permchain.pregel.debug import print_step_start
 from permchain.pregel.read import PregelBatch, PregelInvoke
 from permchain.pregel.validate import validate_chains_channels
 from permchain.pregel.write import PregelSink
@@ -39,10 +46,24 @@ from permchain.pregel.write import PregelSink
 logger = logging.getLogger(__name__)
 
 
+# Before Python 3.11 native StrEnum is not available
+class StrEnum(str, enum.Enum):
+    """A string enum."""
+
+    pass
+
+
+class PregelIO(StrEnum):
+    """Pregel IO channels."""
+
+    IN = "__pregel_input"
+    OUT = "__pregel_output"
+
+
 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]
 
@@ -57,24 +78,40 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
     def __init__(
         self,
-        *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch,
+        main: Optional[Runnable] = None,
+        *,
+        chains: Mapping[str, PregelInvoke | PregelBatch],
         channels: Mapping[str, Channel],
-        output: str | Sequence[str],
-        input: str | Sequence[str],
+        output: str | Sequence[str] = PregelIO.OUT,
+        input: str | Sequence[str] = PregelIO.IN,
         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))
+        chains = {**chains}
+        channels = {**channels}
 
-        validate_chains_channels(chains_flat, channels, input, output)
+        if main is not None:
+            chains[CHAINS_MAIN] = (
+                Pregel.subscribe_to(PregelIO.IN) | main | Pregel.send_to(PregelIO.OUT)
+            )
+        elif output is PregelIO.OUT:
+            raise ValueError(
+                f"When no main runnable is provided, output must be one or more of the channels in {channels.keys()}"
+            )
+
+        if input is PregelIO.IN:
+            channels[PregelIO.IN] = LastValue(
+                main.input_schema if main is not None else Any  # type: ignore[arg-type]
+            )
+        if output is PregelIO.OUT:
+            channels[PregelIO.OUT] = LastValue(
+                main.output_schema if main is not None else Any  # type: ignore[arg-type]
+            )
+
+        validate_chains_channels(chains, channels, input, output)
 
         super().__init__(
-            chains=chains_flat,
+            chains=chains,
             channels=channels,
             output=output,
             input=input,
@@ -157,7 +194,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
@@ -188,16 +225,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]]()
@@ -220,7 +248,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                                 },
                             ),
                         )
-                        for proc, input in next_tasks
+                        for proc, input, _ in next_tasks
                     ),
                     return_when=concurrent.futures.FIRST_EXCEPTION,
                     timeout=self.step_timeout,
@@ -251,7 +279,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
                 # 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):
+                    if any(chan == 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}:
@@ -269,7 +297,7 @@ 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(
@@ -300,16 +328,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]]()
@@ -335,7 +354,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                                 ),
                             )
                         )
-                        for proc, input in next_tasks
+                        for proc, input, _ in next_tasks
                     ],
                     return_when=asyncio.FIRST_EXCEPTION,
                     timeout=self.step_timeout,
@@ -366,7 +385,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
                 # 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):
+                    if any(chan == 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}:
@@ -441,10 +460,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
 
 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:
@@ -456,13 +475,13 @@ def _apply_writes_and_prepare_next_tasks(
         if chan in channels:
             channels[chan].update(vals)
             updated_channels.add(chan)
-        else:
+        elif chan != PregelIO.OUT:
             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 +494,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 +506,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,5 @@
+CHAINS_MAIN = "__pregel_main"
+
 CONFIG_KEY_STEP = "__pregel_step"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"

permchain/pregel/debug.py

@@ -0,0 +1,28 @@
+from pprint import pformat
+from typing import Any
+
+from langchain.schema.runnable import Runnable
+from langchain.utils.input import get_bolded_text, get_colored_text
+
+from permchain.pregel.constants import CHAINS_MAIN
+
+
+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"
+        )
+        + pformat(
+            {name if name != CHAINS_MAIN else _main: val for _, val, name in next_tasks}
+        )
+    )
+
+
+class MainChainStr:
+    def __repr__(cls) -> str:
+        return "main"
+
+
+_main = MainChainStr()

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):

tests/test_pregel.py

@@ -4,11 +4,10 @@ from contextlib import contextmanager
 from typing import Generator
 
 import pytest
-from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema.runnable import RunnableLambda, 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:
@@ -16,7 +15,9 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -30,12 +31,26 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     assert app.invoke(2) == 3
 
 
+def test_invoke_single_process_in_out_implicit(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = RunnableLambda(add_one)
+
+    app = Pregel(chain, channels={}, chains={})
+
+    # TODO add back when langchain is updated with with_types()
+    # assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
+    # assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
+    assert app.invoke(2) == 3
+
+
 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")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -58,7 +73,9 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
     chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -86,7 +103,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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),
@@ -105,7 +122,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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),
@@ -131,8 +148,7 @@ def test_batch_two_processes_in_out() -> None:
     )
 
     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 +170,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.send_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
+        chains[str(i)] = (
             Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_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.send_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 +198,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.send_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
+        chains[str(i)] = (
             Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_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.send_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}) == [
@@ -222,8 +234,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
     chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_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),
@@ -244,8 +255,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
     chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_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,9 +277,11 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
     chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_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 +306,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.send_to("output")
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -315,9 +329,11 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
     chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_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),
@@ -346,8 +362,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_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),
@@ -366,8 +381,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
     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),
@@ -390,8 +404,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
 
     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),
@@ -419,7 +432,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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

@@ -3,11 +3,10 @@ from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncGenerator, AsyncIterator, Generator
 
 import pytest
-from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema.runnable import RunnableLambda, 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:
@@ -15,7 +14,9 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -27,12 +28,26 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     assert await app.ainvoke(2) == 3
 
 
+async def test_invoke_single_process_in_out_implicit(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = RunnableLambda(add_one)
+
+    app = Pregel(chain, channels={}, chains={})
+
+    # TODO add back when langchain is updated with with_types()
+    # assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
+    # assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
+    assert await app.ainvoke(2) == 3
+
+
 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")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -55,7 +70,9 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
     chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
 
     app = Pregel(
-        (chain,),
+        chains={
+            "one": chain,
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -83,7 +100,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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),
@@ -102,7 +119,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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),
@@ -129,8 +146,7 @@ async def test_batch_two_processes_in_out() -> None:
     )
 
     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 +168,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.send_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
+        chains[str(i)] = (
             Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_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.send_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 +197,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.send_to("-1")}
     for i in range(test_size - 2):
         chans[str(i)] = channels.LastValue(int)
-        chains.append(
+        chains[str(i)] = (
             Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_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.send_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):
@@ -227,8 +239,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
     chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_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),
@@ -249,8 +260,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
     chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_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),
@@ -272,9 +282,11 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
     chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_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 +312,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.send_to("output")
+        },
         channels={
             "input": channels.LastValue(int),
             "output": channels.LastValue(int),
@@ -321,9 +335,11 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
     chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_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),
@@ -354,8 +370,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
     chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_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),
@@ -375,8 +390,7 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
     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),
@@ -419,7 +433,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_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),