Update API

README.md

@@ -9,26 +9,32 @@
 ```python
 from permchain import Pregel, channels
 
-value = channels.LastValue[str]("value")
-
 grow_value = (
-    Pregel.subscribe_to(value)
+    Pregel.subscribe_to("value")
     | (lambda x: x + x)
-    | Pregel.send_to({value: lambda x: x if len(x) < 10 else None})
+    | Pregel.send_to(value=lambda x: x if len(x) < 10 else None)
 )
 
-pubsub = Pregel(grow_value, input=value, output=value)
+pubsub = Pregel(
+    grow_value,
+    channels={"value": channels.LastValue[str]()},
+    input="value",
+    output="value",
+)
 
 assert pubsub.invoke("a") == "aaaaaaaa"
+
 ```
 
 Check `examples` for more examples.
 
 ## Near-term Roadmap
 
-- [ ] Iterate on API
-  - [ ] do we want api to receive output from multiple channels in invoke()
-  - [ ] do we want api to send input to multiple channels in invoke()
+- [x] Iterate on API
+  - [x] do we want api to receive output from multiple channels in invoke()
+  - [x] do we want api to send input to multiple channels in invoke()
+  - [ ] Finish updating tests to new API
+- [ ] Implement input_schema and output_schema in Pregel
 - [ ] Implement checkpointing
   - [ ] Save checkpoints at end of each step
   - [ ] Load checkpoint at start of invocation

examples/readme.py

@@ -1,14 +1,16 @@
-import operator
 from permchain import Pregel, channels
 
-value = channels.LastValue[str]("value")
-
 grow_value = (
-    Pregel.subscribe_to(value)
+    Pregel.subscribe_to("value")
     | (lambda x: x + x)
-    | Pregel.send_to({value: lambda x: x if len(x) < 10 else None})
+    | Pregel.send_to(value=lambda x: x if len(x) < 10 else None)
 )
 
-pubsub = Pregel(grow_value, input=value, output=value)
+pubsub = Pregel(
+    grow_value,
+    channels={"value": channels.LastValue[str]()},
+    input="value",
+    output="value",
+)
 
 assert pubsub.invoke("a") == "aaaaaaaa"

examples/runnable-pregel.py

@@ -74,53 +74,61 @@ editor_chain = (
     | JsonOutputFunctionsParser(args_only=False)
 )
 
-# channels
-
-question = channels.LastValue[str]("question")
-
-draft = channels.LastValue[str]("draft")
-
-notes = channels.LastValue[str]("notes")
-
 # application
 
-drafter_node = (
-    Pregel.subscribe_to(question=question) | drafter_chain | Pregel.send_to(draft)
+drafter = (
+    # subscribe to question channel as a dict with a single key, "question"
+    Pregel.subscribe_to(["question"])
+    | drafter_chain
+    | Pregel.send_to("draft")
 )
 
-editor_node = (
-    Pregel.subscribe_to(draft=draft)
+editor = (
+    # subscribe to draft channel as a dict with a single key, "draft"
+    Pregel.subscribe_to(["draft"])
     | editor_chain
     | Pregel.send_to(
-        {notes: lambda x: x["arguments"]["notes"] if x["name"] == "revise" else None}
+        # 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_node = (
-    Pregel.subscribe_to(notes=notes).join(question=question, draft=draft)
+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)
+    | Pregel.send_to("draft")
 )
 
 draft_revise_loop = Pregel(
-    drafter_node,
-    reviser_node,
-    editor_node,
-    input=question,
-    output=draft,
+    [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"],
+    # input can be a dict with any of the channels as keys
+    input=None,
 )
 
 # run
 
-# for draft in draft_revise_loop.stream("What food do turtles eat?"):
-#     print('Draft: "' + draft + '"')
-#     print("---")
+for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}):
+    print(draft)
+    print("---")
 
 
 async def main():
-    async for draft in draft_revise_loop.astream("What food do turtles eat?"):
-        print('Draft: "' + draft + '"')
+    async for draft in draft_revise_loop.astream(
+        {"question": "What food do turtles eat?"}
+    ):
+        print(draft)
         print("---")
 
 
-asyncio.run(main())
+# asyncio.run(main())

permchain/channels.py

@@ -2,7 +2,7 @@ import json
 from abc import ABC, abstractmethod
 from typing import Callable, FrozenSet, Generic, Optional, Sequence, TypeVar
 
-from typing_extensions import Self
+from typing_extensions import Self, get_args
 
 Value = TypeVar("Value")
 Update = TypeVar("Update")
@@ -17,14 +17,20 @@ class InvalidUpdateError(Exception):
 
 
 class Channel(Generic[Value, Update], ABC):
-    def __init__(self, name: str) -> None:
-        self.name = name
+    # TODO: add type hints for ValueType and UpdateType
+    # @property
+    # def ValueType(self) -> type[Value]:
+    #     """The type of the value stored in the channel."""
+    #     type_args = get_args(self.__class__.__orig_bases__[-1])  # type: ignore[attr-defined]
+    #     if type_args and len(type_args) == 2:
+    #         return type_args[0]
 
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self.name})"
-
-    def __str__(self) -> str:
-        return self.name
+    # @property
+    # def UpdateType(self) -> type[Update]:
+    #     """The type of the update received by the channel."""
+    #     type_args = get_args(self.__class__.__orig_bases__[-1])  # type: ignore[attr-defined]
+    #     if type_args and len(type_args) == 2:
+    #         return type_args[1]
 
     @abstractmethod
     def _empty(self, checkpoint: Optional[str] = None) -> Self:
@@ -53,12 +59,12 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
     ```
     """
 
-    def __init__(self, name: str, operator: Callable[[Value, Value], Value]):
-        super().__init__(name)
+    def __init__(self, operator: Callable[[Value, Value], Value]):
+        super().__init__()
         self.operator = operator
 
     def _empty(self, checkpoint: Optional[str] = None) -> Self:
-        empty = self.__class__(self.name, self.operator)
+        empty = self.__class__(self.operator)
         if checkpoint is not None:
             empty.value = json.loads(checkpoint)
         return empty
@@ -85,7 +91,7 @@ class LastValue(Generic[Value], Channel[Value, Value]):
     """Stores the last value received."""
 
     def _empty(self, checkpoint: Optional[str] = None) -> Self:
-        empty = self.__class__(self.name)
+        empty = self.__class__()
         if checkpoint is not None:
             empty.value = json.loads(checkpoint)
         return empty
@@ -110,7 +116,7 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
     """Stores all values received, resets in each step."""
 
     def _empty(self, checkpoint: Optional[str] = None) -> Self:
-        empty = self.__class__(self.name)
+        empty = self.__class__()
         if checkpoint is not None:
             empty.queue = tuple(json.loads(checkpoint))
         return empty
@@ -134,7 +140,7 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
     set: set[Value]
 
     def _empty(self, checkpoint: Optional[str] = None) -> Self:
-        empty = self.__class__(self.name)
+        empty = self.__class__()
         if checkpoint is not None:
             empty.set = set(json.loads(checkpoint))
         return empty

permchain/pregel.py

@@ -13,6 +13,7 @@ from typing import (
     Mapping,
     Optional,
     Sequence,
+    cast,
     overload,
 )
 
@@ -50,11 +51,11 @@ CONFIG_KEY_STEP = "__pregel_step"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 
-TYPE_SEND = Callable[[Sequence[tuple[Channel, Any]]], None]
+TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
 
 class PregelRead(RunnableLambda):
-    channel: Channel
+    channel: str
 
     @property
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
@@ -68,13 +69,13 @@ class PregelRead(RunnableLambda):
             ),
         ]
 
-    def __init__(self, channel: Channel) -> None:
+    def __init__(self, channel: str) -> None:
         super().__init__(func=self._read, afunc=self._aread)  # type: ignore[arg-type]
         self.channel = channel
 
     def _read(self, _: Any, config: RunnableConfig) -> Any:
         try:
-            read: Callable[[Channel], Any] = config["configurable"][CONFIG_KEY_READ]
+            read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
         except KeyError:
             raise RuntimeError(
                 f"Runnable {self} is not configured with a read function"
@@ -84,7 +85,7 @@ class PregelRead(RunnableLambda):
 
     async def _aread(self, _: Any, config: RunnableConfig) -> Any:
         try:
-            read: Callable[[Channel], Any] = config["configurable"][CONFIG_KEY_READ]
+            read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
         except KeyError:
             raise RuntimeError(
                 f"Runnable {self} is not configured with a read function"
@@ -94,15 +95,15 @@ class PregelRead(RunnableLambda):
 
 
 class PregelInvoke(RunnableBinding):
-    channels: Mapping[None, Channel] | Mapping[str, Channel]
+    channels: Mapping[None, str] | Mapping[str, str]
 
     bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
 
-    def join(self, **channels: Channel) -> PregelInvoke:
+    def join(self, *channels: str) -> PregelInvoke:
         joiner = RunnablePassthrough.assign(
-            **{k: PregelRead(chan) for k, chan in channels.items()}
+            **{chan: PregelRead(chan) for chan in channels}
         )
         if isinstance(self.bound, RunnablePassthrough):
             return PregelInvoke(channels=self.channels, bound=joiner)
@@ -130,7 +131,7 @@ class PregelInvoke(RunnableBinding):
 
 
 class PregelBatch(RunnableEach):
-    channel: Inbox
+    channel: str
 
     bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
 
@@ -155,7 +156,7 @@ class PregelBatch(RunnableEach):
 
 
 class PregelSink(RunnableLambda):
-    channels: Sequence[tuple[Channel, Runnable]]
+    channels: Sequence[tuple[str, Runnable]]
     """
     Mapping of write channels to Runnables that return the value to be written,
     or None to skip writing.
@@ -166,7 +167,7 @@ class PregelSink(RunnableLambda):
     def __init__(
         self,
         *,
-        channels: Sequence[tuple[Channel, Runnable]],
+        channels: Sequence[tuple[str, Runnable]],
         max_steps: Optional[int] = None,
     ):
         super().__init__(func=self._write, afunc=self._awrite)
@@ -221,12 +222,14 @@ class PregelSink(RunnableLambda):
         return input
 
 
-class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
-    input: Channel[Any, Input]
+class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]):
+    channels: Mapping[str, Channel]
 
-    output: Channel[Output, Any]
+    chains: Sequence[PregelInvoke | PregelBatch]
 
-    processes: Sequence[PregelInvoke | PregelBatch]
+    output: str | Sequence[str]
+
+    input: str | None
 
     step_timeout: Optional[float] = None
 
@@ -235,42 +238,34 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
 
     def __init__(
         self,
-        *processes: PregelInvoke | PregelBatch,
-        input: Channel[Input, Any],
-        output: Channel[Output, Any],
+        *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch,
+        channels: Mapping[str, Channel],
+        output: str | Sequence[str],
+        input: str | None = None,
         step_timeout: Optional[float] = None,
-        **kwargs: Any,
-    ):
+    ) -> None:
+        chains_flat: list[PregelInvoke | PregelBatch] = []
+        for chain in chains:
+            if isinstance(chain, (list, tuple)):
+                chains_flat.extend(chain)
+            else:
+                chains_flat.append(chain)
         super().__init__(
-            processes=processes,
-            input=input,
+            chains=chains_flat,
+            channels=channels,
             output=output,
+            input=input,
             step_timeout=step_timeout,
-            **kwargs,
         )
 
-    @overload
     @classmethod
-    def subscribe_to(cls, __channel: Channel) -> PregelInvoke:
-        ...
-
-    @overload
-    @classmethod
-    def subscribe_to(
-        cls, __channel: Mapping[str, Channel] | None = None, **kwargs: Channel
-    ) -> PregelInvoke:
-        ...
-
-    @classmethod
-    def subscribe_to(
-        cls, __channel: Channel | Mapping[str, Channel] | None = None, **kwargs: Channel
-    ) -> PregelInvoke:
-        """Runs process.invoke() each time channels are updated."""
-        __channel = __channel or {}
-        return (
-            PregelInvoke(channels={None: __channel})
-            if isinstance(__channel, Channel)
-            else PregelInvoke(channels={**__channel, **kwargs})
+    def subscribe_to(cls, channels: str | Sequence[str]) -> PregelInvoke:
+        """Runs process.invoke() each time channels are updated,
+        with a dict of the channel values as input."""
+        return PregelInvoke(
+            channels={None: channels}
+            if isinstance(channels, str)
+            else {chan: chan for chan in channels}
         )
 
     @classmethod
@@ -281,54 +276,34 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
     @classmethod
     def send_to(
         cls,
-        channels: Channel | Mapping[Channel, RunnableLike],
-        *,
-        max_steps: Optional[int] = None,
+        *channels: str,
+        _max_steps: Optional[int] = None,
+        **kwargs: RunnableLike,
     ) -> PregelSink:
         """Writes to channels the result of the lambda, or None to skip writing."""
         return PregelSink(
             channels=(
-                [(channels, RunnablePassthrough())]
-                if isinstance(channels, Channel)
-                else [(k, coerce_to_runnable(v)) for k, v in channels.items()]
+                [(c, RunnablePassthrough()) for c in channels]
+                + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
             ),
-            max_steps=max_steps,
+            max_steps=_max_steps,
         )
 
-    def _prepare_channels(self) -> Mapping[Channel, Channel]:
-        channels: dict[Channel, Channel] = {self.output: self.output._empty()}
-        for proc in self.processes:
-            if isinstance(proc, PregelInvoke):
-                for chan in proc.channels.values():
-                    if chan not in channels:
-                        channels[chan] = chan._empty()
-            elif isinstance(proc, PregelBatch):
-                if proc.channel not in channels:
-                    channels[proc.channel] = proc.channel._empty()
-            else:
-                raise TypeError(
-                    f"Received process {proc}, expected instance of PregelInvoke or PregelBatch"
-                )
-
-        if not channels:
-            raise ValueError("Found 0 channels for Pregel run")
-
-        if self.input not in channels:
-            raise ValueError("Input channel not being read from")
-
-        return channels
-
     def _transform(
         self,
-        input: Iterator[Input],
+        input: Iterator[dict[str, Any] | Any],
         run_manager: CallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> Iterator[Output]:
-        processes = tuple(self.processes)
+        processes = tuple(self.chains)
         # TODO this is where we'd restore from checkpoint
-        channels = self._prepare_channels()
+        channels = {k: v._empty() for k, v in self.channels.items()}
         next_tasks = _apply_writes_and_prepare_next_tasks(
-            processes, channels, deque((self.input, chunk) for chunk in input)
+            processes,
+            channels,
+            deque((self.input, chunk) for chunk in input)
+            if self.input is not None
+            else deque((k, v) for chunk in input for k, v in chunk.items()),
         )
 
         def read(chan: Channel) -> Any:
@@ -406,14 +381,18 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
 
     async def _atransform(
         self,
-        input: AsyncIterator[Input],
+        input: AsyncIterator[dict[str, Any] | Any],
         run_manager: AsyncCallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> AsyncIterator[Output]:
-        processes = tuple(self.processes)
-        channels = self._prepare_channels()
+        processes = tuple(self.chains)
+        channels = {k: v._empty() for k, v in self.channels.items()}
         next_tasks = _apply_writes_and_prepare_next_tasks(
-            processes, channels, [(self.input, chunk) async for chunk in input]
+            processes,
+            channels,
+            deque((self.input, chunk) async for chunk in input)
+            if self.input is not None
+            else deque((k, v) async for chunk in input for k, v in chunk.items()),
         )
 
         def read(chan: Channel) -> Any:
@@ -480,15 +459,22 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
             )
 
             # if any write to output channel in this step, yield current value
-            if any(chan is self.output for chan, _ in pending_writes):
-                yield channels[self.output]._get()
+            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 no more tasks, we're done
             if not next_tasks:
                 break
 
     def invoke(
-        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
     ) -> Output:
         latest: Output | None = None
         for chunk in self.stream(input, config, **kwargs):
@@ -496,13 +482,16 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
         return latest
 
     def stream(
-        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
     ) -> Iterator[Output]:
         return self.transform(iter([input]), config, **kwargs)
 
     def transform(
         self,
-        input: Iterator[Input],
+        input: Iterator[dict[str, Any] | Any],
         config: RunnableConfig | None = None,
         **kwargs: Any | None,
     ) -> Iterator[Output]:
@@ -511,7 +500,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
         )
 
     async def ainvoke(
-        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
     ) -> Output:
         latest: Output | None = None
         async for chunk in self.astream(input, config, **kwargs):
@@ -519,7 +511,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
         return latest
 
     async def astream(
-        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
     ) -> AsyncIterator[Output]:
         async def input_stream() -> AsyncIterator[Input]:
             yield input
@@ -529,7 +524,7 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
 
     async def atransform(
         self,
-        input: AsyncIterator[Input],
+        input: AsyncIterator[dict[str, Any] | Any],
         config: RunnableConfig | None = None,
         **kwargs: Any | None,
     ) -> AsyncIterator[Output]:
@@ -541,10 +536,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]):
 
 def _apply_writes_and_prepare_next_tasks(
     processes: Sequence[PregelInvoke | PregelBatch],
-    channels: Mapping[Channel, Channel],
-    pending_writes: Sequence[tuple[Channel, Any]],
+    channels: Mapping[str, Channel],
+    pending_writes: Sequence[tuple[str, Any]],
 ) -> list[tuple[Runnable, Any]]:
-    pending_writes_by_channel: dict[Channel, list[Any]] = defaultdict(list)
+    pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
     # Group writes by channel
     for chan, val in pending_writes:
         pending_writes_by_channel[chan].append(val)