Add doc strings for all channels, small update to names/semantics

examples/combine_docs.ipynb

@@ -25,7 +25,7 @@
     "from langchain.schema.document import Document\n",
     "from langchain.schema import format_document\n",
     "\n",
-    "from permchain import Pregel, PregelRead, channels\n"
+    "from permchain import Pregel, PregelRead, channels"
    ]
   },
   {
@@ -45,7 +45,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain.schema.runnable import RunnableLambda\n"
+    "from langchain.schema.runnable import RunnableLambda"
    ]
   },
   {
@@ -59,7 +59,7 @@
     "\n",
     "_combine_documents = RunnableLambda(\n",
     "    lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n",
-    ").map() | (lambda x: \"\\n\\n\".join(x))\n"
+    ").map() | (lambda x: \"\\n\\n\".join(x))"
    ]
   },
   {
@@ -72,7 +72,7 @@
     "docs = [\n",
     "    Document(page_content=\"Harrison used to work at Kensho\"),\n",
     "    Document(page_content=\"Ankush worked at Facebook\"),\n",
-    "]\n"
+    "]"
    ]
   },
   {
@@ -98,7 +98,7 @@
     "    )\n",
     "    | ChatOpenAI()\n",
     "    | StrOutputParser()\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -119,7 +119,7 @@
     }
    ],
    "source": [
-    "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})\n"
+    "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})"
    ]
   },
   {
@@ -139,7 +139,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "many_docs = docs * 5\n"
+    "many_docs = docs * 5"
    ]
   },
   {
@@ -164,7 +164,7 @@
     "            new_result_doc_list.append(_sub_result_docs[:-1])\n",
     "            _sub_result_docs = _sub_result_docs[-1:]\n",
     "    new_result_doc_list.append(_sub_result_docs)\n",
-    "    return new_result_doc_list\n"
+    "    return new_result_doc_list"
    ]
   },
   {
@@ -196,7 +196,7 @@
    "source": [
     "# Just to show what its like split\n",
     "split_docs = _split_list_of_docs(many_docs)\n",
-    "split_docs\n"
+    "split_docs"
    ]
   },
   {
@@ -214,7 +214,7 @@
     "    \"docs_to_finalize\": channels.Inbox(Document),\n",
     "    # output\n",
     "    \"answer\": channels.LastValue(str),\n",
-    "}\n"
+    "}"
    ]
   },
   {
@@ -249,7 +249,7 @@
     "    Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
     "    | stuff_chain\n",
     "    | Pregel.write_to(\"answer\")\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -268,7 +268,7 @@
     "    input=[\"question\", \"docs\"],\n",
     "    output=\"answer\",\n",
     "    debug=True,\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -332,7 +332,7 @@
     }
    ],
    "source": [
-    "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})\n"
+    "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})"
    ]
   },
   {

examples/recursive-web-loader.py

@@ -86,8 +86,8 @@ def recursive_web_loader(
     channels = {
         "base_url": Channels.LastValue(str),
         "next_urls": Channels.UniqueInbox(str),
-        "documents": Channels.Stream(Document),
-        "visited": Channels.Set(str),
+        "documents": Channels.Archive(Document),
+        "visited": Channels.UniqueArchive(str),
         "client": Channels.ContextManager(httpx_client, httpx_aclient),
     }
     # the main chain that gets executed recursively

permchain/channels/__init__.py

@@ -1,15 +1,15 @@
+from permchain.channels.archive import Archive, UniqueArchive
 from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import ContextManager
 from permchain.channels.inbox import Inbox, UniqueInbox
 from permchain.channels.last_value import LastValue
-from permchain.channels.stream import Set, Stream
 
 __all__ = [
     "LastValue",
     "Inbox",
     "UniqueInbox",
+    "Archive",
+    "UniqueArchive",
     "BinaryOperatorAggregate",
-    "Set",
-    "Stream",
     "ContextManager",
 ]

permchain/channels/archive.py

@@ -5,50 +5,11 @@ from typing import Any, FrozenSet, Generator, Generic, Optional, Sequence, Type
 from typing_extensions import Self
 
 from permchain.channels.base import Channel, EmptyChannelError, Value
+from permchain.channels.inbox import flatten
 
 
-class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
-    """Stores all unique values received."""
-
-    def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
-        self.set = set[Value]()
-
-    @property
-    def ValueType(self) -> Type[FrozenSet[Value]]:
-        """The type of the value stored in the channel."""
-        return FrozenSet[self.typ]  # type: ignore[name-defined]
-
-    @property
-    def UpdateType(self) -> Type[Value]:
-        """The type of the update received by the channel."""
-        return self.typ
-
-    @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
-        empty = self.__class__(self.typ)
-        if checkpoint is not None:
-            empty.set = set(json.loads(checkpoint))
-        try:
-            yield empty
-        finally:
-            pass
-
-    def update(self, values: Sequence[Value]) -> None:
-        self.set.update(values)
-
-    def get(self) -> FrozenSet[Value]:
-        try:
-            return frozenset(self.set)
-        except AttributeError:
-            raise EmptyChannelError()
-
-    def checkpoint(self) -> str:
-        return json.dumps(list(self.set))
-
-
-class Stream(Generic[Value], Channel[Sequence[Value], Value]):
-    """Stores all unique values received."""
+class Archive(Generic[Value], Channel[Sequence[Value], Value | list[Value]]):
+    """Stores all unique values received, persists across steps."""
 
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
@@ -74,8 +35,8 @@ class Stream(Generic[Value], Channel[Sequence[Value], Value]):
         finally:
             pass
 
-    def update(self, values: Sequence[Value]) -> None:
-        self.set.extend(values)
+    def update(self, values: Sequence[Value | list[Value]]) -> None:
+        self.set.extend(flatten(values))
 
     def get(self) -> Sequence[Value]:
         try:
@@ -84,4 +45,50 @@ class Stream(Generic[Value], Channel[Sequence[Value], Value]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> str:
-        return json.dumps(self.set)
+        try:
+            return json.dumps(self.set)
+        except AttributeError:
+            raise EmptyChannelError()
+
+
+class UniqueArchive(Generic[Value], Channel[FrozenSet[Value], Value]):
+    """Stores all unique values received, persists across steps."""
+
+    def __init__(self, typ: Type[Value]) -> None:
+        self.typ = typ
+        self.set = set[Value]()
+
+    @property
+    def ValueType(self) -> Type[FrozenSet[Value]]:
+        """The type of the value stored in the channel."""
+        return FrozenSet[self.typ]  # type: ignore[name-defined]
+
+    @property
+    def UpdateType(self) -> Type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ)
+        if checkpoint is not None:
+            empty.set = set(json.loads(checkpoint))
+        try:
+            yield empty
+        finally:
+            pass
+
+    def update(self, values: Sequence[Value | list[Value]]) -> None:
+        self.set.update(flatten(values))
+
+    def get(self) -> FrozenSet[Value]:
+        try:
+            return frozenset(self.set)
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> str:
+        try:
+            return json.dumps(list(self.set))
+        except AttributeError:
+            raise EmptyChannelError()

permchain/channels/base.py

@@ -18,10 +18,15 @@ Update = TypeVar("Update")
 
 
 class EmptyChannelError(Exception):
+    """Raised when attempting to get the value of a channel that hasn't been updated
+    for the first time yet."""
+
     pass
 
 
 class InvalidUpdateError(Exception):
+    """Raised when attempting to update a channel with an invalid sequence of updates."""
+
     pass
 
 
@@ -51,21 +56,30 @@ class Channel(Generic[Value, Update], ABC):
 
     @abstractmethod
     def update(self, values: Sequence[Update]) -> None:
-        ...
+        """Update the channel's value with the given sequence of updates.
+        The order of the updates in the sequence is arbitrary.
+
+        Raises InvalidUpdateError if the sequence of updates is invalid."""
 
     @abstractmethod
     def get(self) -> Value:
-        ...
+        """Return the current value of the channel.
+
+        Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
     @abstractmethod
     def checkpoint(self) -> str | None:
-        ...
+        """Return a string representation of the channel's current state,
+        or None if the channel doesn't support checkpoints.
+
+        Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
 
 @contextmanager
 def ChannelsManager(
     channels: Mapping[str, Channel]
 ) -> Generator[Mapping[str, Channel], None, None]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
     empty = {k: v.empty() for k, v in channels.items()}
     try:
         yield {k: v.__enter__() for k, v in empty.items()}
@@ -78,6 +92,7 @@ def ChannelsManager(
 async def AsyncChannelsManager(
     channels: Mapping[str, Channel]
 ) -> AsyncGenerator[Mapping[str, Channel], None]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
     empty = {k: v.aempty() for k, v in channels.items()}
     try:
         yield {k: await v.__aenter__() for k, v in empty.items()}

permchain/channels/binop.py

@@ -59,4 +59,7 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> str:
-        return json.dumps(self.value)
+        try:
+            return json.dumps(self.value)
+        except AttributeError:
+            raise EmptyChannelError()

permchain/channels/context.py

@@ -23,6 +23,18 @@ from permchain.channels.base import (
 
 
 class ContextManager(Generic[Value], Channel[Value, None]):
+    """Exposes the value of a context manager, for the duration of an invocation.
+    Context manager is entered before the first step, and exited after the last step.
+    Optionally, provide an equivalent async context manager, which will be used
+    instead for async invocations.
+
+    ```python
+    import httpx
+
+    client = ContextManager(httpx.Client, httpx.AsyncClient)
+    ```
+    """
+
     value: Value
 
     def __init__(

permchain/channels/inbox.py

@@ -1,17 +1,31 @@
 import json
 from contextlib import contextmanager
-from typing import Any, Generator, Generic, Optional, Sequence, Type, Union, cast
+from typing import (
+    Any,
+    FrozenSet,
+    Generator,
+    Generic,
+    Iterator,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+)
 
 from typing_extensions import Self
 
-from permchain.channels.base import (
-    Channel,
-    EmptyChannelError,
-    Value,
-)
+from permchain.channels.base import Channel, EmptyChannelError, Value
 
 
-class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
+def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]:
+    for value in values:
+        if isinstance(value, list):
+            yield from value
+        else:
+            yield value
+
+
+class Inbox(Generic[Value], Channel[Sequence[Value], Value | list[Value]]):
     """Stores all values received, resets in each step."""
 
     def __init__(self, typ: Type[Value]) -> None:
@@ -40,16 +54,8 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
             except AttributeError:
                 pass
 
-    def update(self, values: Sequence[Value | Sequence[Value]]) -> None:
-        self.queue = tuple(
-            cast(Value, v)
-            for value in values
-            for v in (
-                (value,)
-                if isinstance(value, self.typ)
-                else cast(Sequence[Value], value)
-            )
-        )
+    def update(self, values: Sequence[Value | list[Value]]) -> None:
+        self.queue = tuple(flatten(values))
 
     def get(self) -> Sequence[Value]:
         try:
@@ -58,19 +64,22 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> str:
-        return json.dumps(self.queue)
+        try:
+            return json.dumps(self.queue)
+        except AttributeError:
+            raise EmptyChannelError()
 
 
-class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
+class UniqueInbox(Generic[Value], Channel[FrozenSet[Value], Value | list[Value]]):
     """Stores all unique values received, resets in each step."""
 
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
 
     @property
-    def ValueType(self) -> Type[Sequence[Value]]:
+    def ValueType(self) -> Type[FrozenSet[Value]]:
         """The type of the value stored in the channel."""
-        return Sequence[self.typ]  # type: ignore[name-defined]
+        return FrozenSet[self.typ]  # type: ignore[name-defined]
 
     @property
     def UpdateType(self) -> Any:
@@ -81,7 +90,7 @@ class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Valu
     def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ)
         if checkpoint is not None:
-            empty.queue = tuple(json.loads(checkpoint))
+            empty.queue = frozenset(json.loads(checkpoint))
         try:
             yield empty
         finally:
@@ -90,24 +99,17 @@ class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Valu
             except AttributeError:
                 pass
 
-    def update(self, values: Sequence[Value | Sequence[Value]]) -> None:
-        self.queue = tuple(
-            set(
-                cast(Value, v)
-                for value in values
-                for v in (
-                    (value,)
-                    if isinstance(value, self.typ)
-                    else cast(Sequence[Value], value)
-                )
-            )
-        )
+    def update(self, values: Sequence[Value | list[Value]]) -> None:
+        self.queue = frozenset(flatten(values))
 
-    def get(self) -> Sequence[Value]:
+    def get(self) -> FrozenSet[Value]:
         try:
             return self.queue
         except AttributeError:
             raise EmptyChannelError()
 
     def checkpoint(self) -> str:
-        return json.dumps(self.queue)
+        try:
+            return json.dumps(self.queue)
+        except AttributeError:
+            raise EmptyChannelError()

permchain/channels/last_value.py

@@ -13,7 +13,7 @@ from permchain.channels.base import (
 
 
 class LastValue(Generic[Value], Channel[Value, Value]):
-    """Stores the last value received."""
+    """Stores the last value received, can receive at most one value per step."""
 
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
@@ -54,4 +54,7 @@ class LastValue(Generic[Value], Channel[Value, Value]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> str:
-        return json.dumps(self.value)
+        try:
+            return json.dumps(self.value)
+        except AttributeError:
+            raise EmptyChannelError()

tests/test_channels.py

@@ -72,7 +72,7 @@ async def test_inbox_async() -> None:
 
 
 def test_set() -> None:
-    with Channels.Set(str).empty() as channel:
+    with Channels.UniqueArchive(str).empty() as channel:
         assert channel.ValueType is FrozenSet[str]
         assert channel.UpdateType is str
 
@@ -84,7 +84,7 @@ def test_set() -> None:
 
 
 async def test_set_async() -> None:
-    async with Channels.Set(str).aempty() as channel:
+    async with Channels.UniqueArchive(str).aempty() as channel:
         assert channel.ValueType is FrozenSet[str]
         assert channel.UpdateType is str