0.0.6

Add concept of reserved channels (ie. channels updated by Pregel), Add a reserved channel for is_last_step

LICENSE

@@ -1,21 +1,51 @@
-The MIT License
+# PermChain License
 
-Copyright (c) Harrison Chase
+By using the software, you agree to all of the terms and conditions below.
 
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
+## Copyright License
 
-The above copyright notice and this permission notice shall be included in
-all copies or substantial portions of the Software.
+The licensor grants you a non-exclusive, royalty-free, worldwide, non-sublicensable, non-transferable license to use, copy, distribute, make available, and prepare derivative works of the software, in each case subject to the limitations and conditions below.
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
+## Limitations
+
+You may not provide the software to third parties as a hosted or managed service, where the service provides users with access to any substantial set of the features or functionality of the software.
+
+You may not move, change, disable, or circumvent the license key functionality in the software, and you may not remove or obscure any functionality in the software that is protected by the license key.
+
+You may not alter, remove, or obscure any licensing, copyright, or other notices of the licensor in the software. Any use of the licensor’s trademarks is subject to applicable law.
+
+## Patents
+
+The licensor grants you a license, under any patent claims the licensor can license, or becomes able to license, to make, have made, use, sell, offer for sale, import and have imported the software, in each case subject to the limitations and conditions in this license. This license does not cover any patent claims that you cause to be infringed by modifications or additions to the software. If you or your company make any written claim that the software infringes or contributes to infringement of any patent, your patent license for the software granted under these terms ends immediately. If your company makes such a claim, your patent license ends immediately for work on behalf of your company.
+
+## Notices
+
+You must ensure that anyone who gets a copy of any part of the software from you also gets a copy of these terms.
+
+If you modify the software, you must include in any modified copies of the software prominent notices stating that you have modified the software.
+
+## No Other Rights
+
+These terms do not imply any licenses other than those expressly granted in these terms.
+
+## Termination
+
+If you use the software in violation of these terms, such use is not licensed, and your licenses will automatically terminate. If the licensor provides you with a notice of your violation, and you cease all violation of this license no later than 30 days after you receive that notice, your licenses will be reinstated retroactively. However, if you violate these terms after such reinstatement, any additional violation of these terms will cause your licenses to terminate automatically and permanently.
+
+## No Liability
+
+As far as the law allows, the software comes as is, without any warranty or condition, and the licensor will not be liable to you for any damages arising out of these terms or the use or nature of the software, under any kind of legal claim.
+
+## Definitions
+
+The licensor is the entity offering these terms, and the software is the software the licensor makes available under these terms, including any portion of it.
+
+you refers to the individual or entity agreeing to these terms.
+
+your company is any legal entity, sole proprietorship, or other kind of organization that you work for, plus all organizations that have control over, are under the control of, or are under common control with that organization. control means ownership of substantially all the assets of an entity, or the power to direct its management and policies by vote, contract, or otherwise. Control can be direct or indirect.
+
+your licenses are all the licenses granted to you for the software under these terms.
+
+use means anything you do with the software requiring one of your licenses.
+
+trademark means trademarks, service marks, and similar rights.

Makefile

@@ -31,11 +31,11 @@ lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --diff-filter=d
 
 lint lint_diff:
 	poetry run ruff .
-	poetry run black $(PYTHON_FILES) --check
+	poetry run ruff format $(PYTHON_FILES) --check
 	poetry run mypy $(PYTHON_FILES)
 
 format format_diff:
-	poetry run black $(PYTHON_FILES)
+	poetry run ruff format $(PYTHON_FILES)
 	poetry run ruff --select I --fix $(PYTHON_FILES)
 
 spell_check:

README.md

@@ -22,13 +22,15 @@ Some of the use cases are:
 
 Channels are used to communicate between chains. Each channel has a value type, an update type, and an update function – which takes a sequence of updates and modifies the stored value. Channels can be used to send data from one chain to another, or to send data from a chain to itself in a future step. PermChain provides a number of built-in channels:
 
-- `LastValue`: stores the last value sent to the channel, useful for input values, and single-value outputs
-- `Inbox`: stores an ephemeral sequence of values sent to the channel, useful for sending data from one chain to another
-- `UniqueInbox`: same as Inbox, but deduplicates values sent to the channel
-- `Archive`: stores a persistent sequence of values sent to the channel, useful for accumulating data over multiple steps
-- `UniqueArchive`: same as Archive, but deduplicates values sent to the channel
-- `BinaryOperatorAggregate`: stores a persistent value, updated by applying a binary operator to the current value and each update sent to the channel, useful for computing aggregates over multiple steps. eg. `total = BinaryOperatorAggregate(int, operator.add)`
+#### Basic channels: LastValue and Topic
+
+- `LastValue`: The default channel, stores the last value sent to the channel, useful for input and output values, or for sending data from one step to the next
+- `Topic`: A configurable PubSub Topic, useful for sending multiple values between chains, or for accumulating output. Can be configured to deduplicate values, and/or to accummulate values over the course of multiple steps.
+
+#### Advanced channels: Context and BinaryOperatorAggregate
+
 - `Context`: exposes the value of a context manager, managing its lifecycle. Useful for accessing external resources that require setup and/or teardown. eg. `client = Context(httpx.Client)`
+- `BinaryOperatorAggregate`: stores a persistent value, updated by applying a binary operator to the current value and each update sent to the channel, useful for computing aggregates over multiple steps. eg. `total = BinaryOperatorAggregate(int, operator.add)`
 
 ### Chains
 
@@ -48,7 +50,6 @@ Repeat until no chains are planned for execution, or a maximum number of steps i
 
 ```python
 from permchain import Channel, Pregel
-from permchain.channels import LastValue
 
 grow_value = (
     Channel.subscribe_to("value")
@@ -58,13 +59,11 @@ grow_value = (
 
 app = Pregel(
     chains={"grow_value": grow_value},
-    channels={"value": LastValue(str)},
     input="value",
     output="value",
 )
 
 assert app.invoke("a") == "aaaaaaaa"
-
 ```
 
 Check `examples` for more examples.
@@ -81,16 +80,20 @@ Check `examples` for more examples.
   - [x] Add tests for Stream, UniqueInbox
   - [ ] Add tests for subscribe_to_each().join()
 - [x] Add optional debug logging
-- [ ] Implement checkpointing
-  - [ ] Save checkpoints at end of each step
-  - [ ] Load checkpoint at start of invocation
-  - [ ] API to specify storage backend and save key
+- [ ] Add an optional Diff value for Channels that implements `__add__`, returned by update(), yielded by Pregel for output channels. Add replacing_keys set to AddableDict. use an addabledict for yielding values. channels that dont implement it get marked with replacing_keys
+- [x] Implement checkpointing
+  - [x] Save checkpoints at end of each step/run
+  - [x] Load checkpoint at start of invocation
+  - [x] API to specify storage backend and save key
+  - [x] Tests
 - [ ] Add more examples
+  - [ ] multi agent simulation
   - [ ] human in the loop
   - [ ] combine documents
-  - [ ] agent executor
+  - [ ] agent executor (add current v total iterations info to read/write steps to enable doing a final update at the end)
   - [ ] run over dataset
 - [ ] Fault tolerance
+  - [ ] Expose a unique id to each step, hash of (app, chain, checkpoint) (include input updates for first step)
   - [ ] Retry individual processes in a step
   - [ ] Retry entire step?
 - [ ] Pregel.stream_log to contain additional keys specific to Pregel

examples/combine_docs.ipynb

@@ -25,8 +25,8 @@
     "from langchain.schema.document import Document\n",
     "from langchain.schema import format_document\n",
     "\n",
-    "from permchain import Channel, Pregel, PregelRead\n",
-    "from permchain.channels import LastValue, Inbox"
+    "from permchain import Channel, Pregel\n",
+    "from permchain.channels import LastValue, Topic"
    ]
   },
   {
@@ -209,18 +209,15 @@
    "source": [
     "channels = {\n",
     "    # input\n",
-    "    \"question\": LastValue(str),\n",
-    "    \"docs\": Inbox(Document),\n",
+    "    \"docs\": Topic(Document),\n",
     "    # intermediate\n",
-    "    \"docs_to_finalize\": Inbox(Document),\n",
-    "    # output\n",
-    "    \"answer\": LastValue(str),\n",
+    "    \"docs_to_finalize\": Topic(Document),\n",
     "}"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 23,
    "id": "67370694-86f4-4b64-9d4f-38b2e306abeb",
    "metadata": {},
    "outputs": [],
@@ -234,12 +231,16 @@
     "        return Channel.write_to(\"docs_to_finalize\")\n",
     "\n",
     "\n",
+    "def split_docs_with_question(input: dict[str, str | list[Document]]) -> list[dict[str, str | list[Document]]]:\n",
+    "    return [\n",
+    "        {\"docs\": docs, \"question\": input[\"question\"]}\n",
+    "        for docs in _split_list_of_docs(input[\"docs\"])\n",
+    "    ]\n",
+    "\n",
+    "\n",
     "collapse = (\n",
-    "    Channel.subscribe_to(\"docs\")\n",
-    "    | _split_list_of_docs\n",
-    "    | {\"docs_list\": RunnablePassthrough(), \"question\": PregelRead(\"question\")}\n",
-    "    # {docs: list[list[Doc]], question: str} -> list[{docs: list[Doc], question: str}]\n",
-    "    | (lambda x: [{\"docs\": docs, \"question\": x[\"question\"]} for docs in x[\"docs_list\"]])\n",
+    "    Channel.subscribe_to([\"docs\", \"question\"])\n",
+    "    | split_docs_with_question\n",
     "    | stuff_chain.map()  # Collapse each list of docs to a single string\n",
     "    | (lambda x: [Document(page_content=s) for s in x])  # A new (smaller) list of docs\n",
     "    | decide\n",
@@ -255,7 +256,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 24,
    "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
    "metadata": {},
    "outputs": [],
@@ -274,7 +275,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 25,
    "id": "69fcb829-3dae-432a-8db3-11bbb179a7d2",
    "metadata": {},
    "outputs": [
@@ -283,51 +284,52 @@
      "output_type": "stream",
      "text": [
       "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 0 with 1 task. Next tasks:\n",
-      "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho'),\n",
-      " Document(page_content='Ankush worked at Facebook'),\n",
-      " Document(page_content='Harrison used to work at Kensho'),\n",
-      " Document(page_content='Ankush worked at Facebook'),\n",
-      " Document(page_content='Harrison used to work at Kensho'),\n",
-      " Document(page_content='Ankush worked at Facebook'),\n",
-      " Document(page_content='Harrison used to work at Kensho'),\n",
-      " Document(page_content='Ankush worked at Facebook'),\n",
-      " Document(page_content='Harrison used to work at Kensho'),\n",
-      " Document(page_content='Ankush worked at Facebook')))\n",
+      "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho'),\n",
+      "          Document(page_content='Ankush worked at Facebook'),\n",
+      "          Document(page_content='Harrison used to work at Kensho'),\n",
+      "          Document(page_content='Ankush worked at Facebook'),\n",
+      "          Document(page_content='Harrison used to work at Kensho'),\n",
+      "          Document(page_content='Ankush worked at Facebook'),\n",
+      "          Document(page_content='Harrison used to work at Kensho'),\n",
+      "          Document(page_content='Ankush worked at Facebook'),\n",
+      "          Document(page_content='Harrison used to work at Kensho'),\n",
+      "          Document(page_content='Ankush worked at Facebook')],\n",
+      " 'question': 'where did harrison work'})\n",
       "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 0. Channel values:\n",
-      "\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
+      "\u001b[0m{'docs': [...], 'docs_to_finalize': [], 'question': 'where did harrison work'}\n",
       "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 1 with 1 task. Next tasks:\n",
-      "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.')))\n",
+      "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.')],\n",
+      " 'question': 'where did harrison work'})\n",
       "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 1. Channel values:\n",
-      "\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n",
+      "\u001b[0m{'docs': [...], 'docs_to_finalize': [], 'question': 'where did harrison work'}\n",
       "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 2 with 1 task. Next tasks:\n",
-      "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.'),\n",
-      " Document(page_content='Harrison used to work at Kensho.')))\n",
+      "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.')],\n",
+      " 'question': 'where did harrison work'})\n",
       "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 2. Channel values:\n",
-      "\u001b[0m{'docs': (...),\n",
-      " 'docs_to_finalize': (...),\n",
-      " 'question': 'where did harrison work'}\n",
+      "\u001b[0m{'docs': [], 'docs_to_finalize': [...], 'question': 'where did harrison work'}\n",
       "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 3 with 1 task. Next tasks:\n",
-      "\u001b[0m- finalize({'docs': (Document(page_content='Harrison used to work at Kensho.'),\n",
-      "          Document(page_content='Harrison used to work at Kensho.'))})\n",
+      "\u001b[0m- finalize({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n",
+      "          Document(page_content='Harrison used to work at Kensho.')]})\n",
       "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 3. Channel values:\n",
-      "\u001b[0m{'answer': 'Harrison worked at Kensho.',\n",
-      " 'docs': (...),\n",
-      " 'docs_to_finalize': (...),\n",
+      "\u001b[0m{'answer': 'Harrison used to work at Kensho.',\n",
+      " 'docs': [],\n",
+      " 'docs_to_finalize': [],\n",
       " 'question': 'where did harrison work'}\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "'Harrison worked at Kensho.'"
+       "'Harrison used to work at Kensho.'"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 25,
      "metadata": {},
      "output_type": "execute_result"
     }

examples/draft-revise-loop.py

@@ -6,7 +6,6 @@ from langchain.prompts import SystemMessagePromptTemplate
 from langchain.schema.output_parser import StrOutputParser
 
 from permchain import Channel, Pregel
-from permchain.channels import LastValue
 
 # prompts
 
@@ -75,17 +74,9 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
 
 # application
 
-channels = {
-    "question": LastValue(str),
-    "draft": LastValue(str),
-    "notes": LastValue(str),
-}
-
 drafter = (
     # subscribe to question channel as a dict with a single key, "question"
-    Channel.subscribe_to(["question"])
-    | drafter_chain
-    | Channel.write_to("draft")
+    Channel.subscribe_to(["question"]) | drafter_chain | Channel.write_to("draft")
 )
 
 editor = (
@@ -94,9 +85,7 @@ editor = (
     | editor_chain
     | Channel.write_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
+        notes=lambda x: x["arguments"]["notes"] if x["name"] == "revise" else None
     )
 )
 
@@ -109,7 +98,6 @@ reviser = (
 )
 
 draft_revise_loop = Pregel(
-    channels=channels,
     chains={
         "drafter": drafter,
         "editor": editor,
@@ -117,27 +105,12 @@ draft_revise_loop = Pregel(
     },
     # 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"],
+    # output will be the value of "draft"
+    output="draft",
     # debug logging
     debug=True,
 )
 
 # run
 
-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(
-        {"question": "What food do turtles eat?"}
-    ):
-        print(draft)
-        print("---")
-
-
-# import asyncio
-
-# asyncio.run(main())
+print(draft_revise_loop.invoke({"question": "What food do turtles eat?"}))

examples/rag.py

@@ -0,0 +1,58 @@
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.schema.messages import AIMessage, AnyMessage, FunctionMessage
+from langchain.vectorstores import FAISS
+
+from permchain import Channel, Pregel
+from permchain.channels import Topic
+
+texts = ["harrison went to kensho"]
+embeddings = OpenAIEmbeddings()
+db = FAISS.from_texts(texts, embeddings)
+
+retriever = db.as_retriever()
+
+
+prompt = PromptTemplate.from_template(
+    """Answer the question "{question}"  based on the following context: {context}"""
+)
+
+model = ChatOpenAI()
+
+chain = (
+    Channel.subscribe_to(["question"])
+    | {
+        "context": (lambda x: x["question"])
+        | Channel.write_to(
+            messages=lambda _input: AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": "retrieval",
+                    "arguments": {"question": _input},
+                },
+            )
+        )
+        | retriever
+        | Channel.write_to(
+            messages=lambda documents: FunctionMessage.construct(
+                content=documents,  # function message requires content to be str
+                name="retrieval",
+            )
+        ),
+        "question": lambda x: x["question"],
+    }
+    | prompt
+    | model
+    | Channel.write_to(messages=lambda message: [message])
+)
+
+app = Pregel(
+    chains={"chain": chain},
+    channels={"messages": Topic(AnyMessage)},
+    input=["question"],
+    output=["messages"],
+)
+
+for s in app.stream({"question": "where did harrison go"}):
+    print(s)

examples/readme.py

@@ -1,5 +1,4 @@
 from permchain import Channel, Pregel
-from permchain.channels import LastValue
 
 grow_value = (
     Channel.subscribe_to("value")
@@ -9,7 +8,6 @@ grow_value = (
 
 app = Pregel(
     chains={"grow_value": grow_value},
-    channels={"value": LastValue(str)},
     input="value",
     output="value",
 )

examples/recursive-web-loader.py

@@ -7,7 +7,8 @@ from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
 from langchain.utils.html import extract_sub_links
 
 from permchain import Channel, Pregel
-from permchain.channels import Archive, Context, LastValue, UniqueArchive, UniqueInbox
+from permchain.channels.context import Context
+from permchain.channels.topic import Topic
 
 # Load url with sync httpx client
 
@@ -85,39 +86,28 @@ def recursive_web_loader(
     metadata_extractor = metadata_extractor or _metadata_extractor
     # define the channels
     channels = {
-        "base_url": LastValue(str),
-        "next_urls": UniqueInbox(str),
-        "documents": Archive(Document),
-        "visited": UniqueArchive(str),
+        "next_urls": Topic(str, unique=True),
+        "documents": Topic(Document, accumulate=True),
         "client": Context(httpx_client, httpx_aclient),
     }
     # the main chain that gets executed recursively
+    # while there are urls in next_urls
     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
-        Channel.subscribe_to_each("next_urls", key="url").join(
-            ["visited", "client", "base_url"]
-        )
+        # adding the current values of base_url and httpx client
+        Channel.subscribe_to_each("next_urls", key="url").join(["client", "base_url"])
         # load the url (with sync and async implementations)
         | RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async))
         | Channel.write_to(
-            # send this url to the visited set
-            visited=lambda x: x["url"],
             # send a new document to the documents stream
             documents=lambda x: Document(
                 page_content=extractor(x["body"]),
                 metadata=metadata_extractor(x["body"], x["url"]),
             ),
-            # send the next urls to the next_urls set
-            # only if not visited already
-            next_urls=lambda x: [
-                url
-                for url in extract_sub_links(
-                    x["body"], x["url"], base_url=x["base_url"]
-                )
-                if url not in x["visited"] and url != x["url"]
-            ],
+            # send the next urls to the next_urls topic
+            next_urls=lambda x: extract_sub_links(
+                x["body"], x["url"], base_url=x["base_url"]
+            ),
         )
     )
     return Pregel(

permchain/__init__.py

@@ -1,4 +1,4 @@
-from permchain.pregel import Channel, Pregel
-from permchain.pregel.read import PregelRead
+from permchain.pregel import Channel, Pregel, ReservedChannels
+from permchain.pregel.read import ChannelRead
 
-__all__ = ["Channel", "Pregel", "PregelRead"]
+__all__ = ["Channel", "Pregel", "ReservedChannels", "ChannelRead"]

permchain/channels/__init__.py

@@ -1,15 +1,11 @@
-from permchain.channels.archive import Archive, UniqueArchive
 from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
-from permchain.channels.inbox import Inbox, UniqueInbox
 from permchain.channels.last_value import LastValue
+from permchain.channels.topic import Topic
 
 __all__ = [
     "LastValue",
-    "Inbox",
-    "UniqueInbox",
-    "Archive",
-    "UniqueArchive",
-    "BinaryOperatorAggregate",
+    "Topic",
     "Context",
+    "BinaryOperatorAggregate",
 ]

permchain/channels/archive.py

@@ -1,94 +0,0 @@
-import json
-from contextlib import contextmanager
-from typing import Any, FrozenSet, Generator, Generic, Optional, Sequence, Type
-
-from typing_extensions import Self
-
-from permchain.channels.base import BaseChannel, EmptyChannelError, Value
-from permchain.channels.inbox import flatten
-
-
-class Archive(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]):
-    """Stores all unique values received, persists across steps."""
-
-    def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
-        self.set = list[Value]()
-
-    @property
-    def ValueType(self) -> Any:
-        """The type of the value stored in the channel."""
-        return Sequence[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 = json.loads(checkpoint)
-        try:
-            yield empty
-        finally:
-            pass
-
-    def update(self, values: Sequence[Value | list[Value]]) -> None:
-        self.set.extend(flatten(values))
-
-    def get(self) -> Sequence[Value]:
-        try:
-            return tuple(self.set)
-        except AttributeError:
-            raise EmptyChannelError()
-
-    def checkpoint(self) -> str:
-        try:
-            return json.dumps(self.set)
-        except AttributeError:
-            raise EmptyChannelError()
-
-
-class UniqueArchive(Generic[Value], BaseChannel[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

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
 from typing import (
     Any,
     AsyncGenerator,
@@ -13,8 +14,11 @@ from typing import (
 
 from typing_extensions import Self
 
+from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
+Checkpoint = TypeVar("Checkpoint")
 
 
 class EmptyChannelError(Exception):
@@ -30,7 +34,7 @@ class InvalidUpdateError(Exception):
     pass
 
 
-class BaseChannel(Generic[Value, Update], ABC):
+class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -43,7 +47,9 @@ class BaseChannel(Generic[Value, Update], ABC):
 
     @contextmanager
     @abstractmethod
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(
+        self, checkpoint: Optional[Checkpoint] = None
+    ) -> Generator[Self, None, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
 
     @asynccontextmanager
@@ -68,19 +74,22 @@ class BaseChannel(Generic[Value, Update], ABC):
         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.
+    def checkpoint(self) -> Checkpoint | None:
+        """Return a string representation of the channel's current state.
 
-        Raises EmptyChannelError if the channel is empty (never updated yet)."""
+        Raises EmptyChannelError if the channel is empty (never updated yet),
+        or doesn't supportcheckpoints."""
 
 
 @contextmanager
 def ChannelsManager(
-    channels: Mapping[str, BaseChannel]
+    channels: Mapping[str, BaseChannel],
+    checkpoint: Optional[Mapping[str, Any]],
 ) -> Generator[Mapping[str, BaseChannel], None, None]:
     """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    empty = {k: v.empty() for k, v in channels.items()}
+    # TODO use https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
+    checkpoint = checkpoint or {}
+    empty = {k: v.empty(checkpoint.get(k)) for k, v in channels.items()}
     try:
         yield {k: v.__enter__() for k, v in empty.items()}
     finally:
@@ -90,12 +99,28 @@ def ChannelsManager(
 
 @asynccontextmanager
 async def AsyncChannelsManager(
-    channels: Mapping[str, BaseChannel]
+    channels: Mapping[str, BaseChannel],
+    checkpoint: Optional[Mapping[str, Any]],
 ) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
     """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    empty = {k: v.aempty() for k, v in channels.items()}
+    checkpoint = checkpoint or {}
+    empty = {k: v.aempty(checkpoint.get(k)) for k, v in channels.items()}
     try:
         yield {k: await v.__aenter__() for k, v in empty.items()}
     finally:
         for v in empty.values():
             await v.__aexit__(None, None, None)
+
+
+def create_checkpoint(channels: Mapping[str, BaseChannel]) -> Mapping[str, Any]:
+    """Create a checkpoint for the given channels."""
+    checkpoint = {
+        CHECKPOINT_KEY_VERSION: 1,
+        CHECKPOINT_KEY_TS: datetime.utcnow().isoformat(),
+    }
+    for k, v in channels.items():
+        try:
+            checkpoint[k] = v.checkpoint()
+        except EmptyChannelError:
+            pass
+    return checkpoint

permchain/channels/binop.py

@@ -1,4 +1,3 @@
-import json
 from contextlib import contextmanager
 from typing import Callable, Generator, Generic, Optional, Sequence, Type
 
@@ -7,7 +6,7 @@ from typing_extensions import Self
 from permchain.channels.base import BaseChannel, EmptyChannelError, Value
 
 
-class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
+class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the result of applying a binary operator to the current value and each new value.
 
     ```python
@@ -20,6 +19,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
     def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
         self.typ = typ
         self.operator = operator
+        try:
+            self.value = typ()
+        except Exception:
+            pass
 
     @property
     def ValueType(self) -> Type[Value]:
@@ -32,10 +35,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
         return self.typ
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ, self.operator)
         if checkpoint is not None:
-            empty.value = json.loads(checkpoint)
+            empty.value = checkpoint
         try:
             yield empty
         finally:
@@ -45,6 +48,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
                 pass
 
     def update(self, values: Sequence[Value]) -> None:
+        if not values:
+            return
         if not hasattr(self, "value"):
             self.value = values[0]
             values = values[1:]
@@ -58,8 +63,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
         except AttributeError:
             raise EmptyChannelError()
 
-    def checkpoint(self) -> str:
+    def checkpoint(self) -> Value:
         try:
-            return json.dumps(self.value)
+            return self.value
         except AttributeError:
             raise EmptyChannelError()

permchain/channels/context.py

@@ -22,7 +22,7 @@ from permchain.channels.base import (
 )
 
 
-class Context(Generic[Value], BaseChannel[Value, None]):
+class Context(Generic[Value], BaseChannel[Value, None, 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
@@ -66,7 +66,7 @@ class Context(Generic[Value], BaseChannel[Value, None]):
         raise InvalidUpdateError()
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: None = None) -> Generator[Self, None, None]:
         if self.ctx is None:
             raise ValueError("Cannot enter sync context manager.")
 
@@ -97,7 +97,8 @@ class Context(Generic[Value], BaseChannel[Value, None]):
                 yield empty
 
     def update(self, values: Sequence[None]) -> None:
-        raise InvalidUpdateError()
+        if values:
+            raise InvalidUpdateError()
 
     def get(self) -> Value:
         try:
@@ -106,4 +107,4 @@ class Context(Generic[Value], BaseChannel[Value, None]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> None:
-        return None
+        raise EmptyChannelError()

permchain/channels/inbox.py

@@ -1,115 +0,0 @@
-import json
-from contextlib import contextmanager
-from typing import (
-    Any,
-    FrozenSet,
-    Generator,
-    Generic,
-    Iterator,
-    Optional,
-    Sequence,
-    Type,
-    Union,
-)
-
-from typing_extensions import Self
-
-from permchain.channels.base import BaseChannel, EmptyChannelError, 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], BaseChannel[Sequence[Value], Value | list[Value]]):
-    """Stores all values received, resets in each step."""
-
-    def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
-
-    @property
-    def ValueType(self) -> Type[Sequence[Value]]:
-        """The type of the value stored in the channel."""
-        return Sequence[self.typ]  # type: ignore[name-defined]
-
-    @property
-    def UpdateType(self) -> Any:
-        """The type of the update received by the channel."""
-        return Union[self.typ, Sequence[self.typ]]  # type: ignore[name-defined]
-
-    @contextmanager
-    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))
-        try:
-            yield empty
-        finally:
-            try:
-                del empty.queue
-            except AttributeError:
-                pass
-
-    def update(self, values: Sequence[Value | list[Value]]) -> None:
-        self.queue = tuple(flatten(values))
-
-    def get(self) -> Sequence[Value]:
-        try:
-            return self.queue
-        except AttributeError:
-            raise EmptyChannelError()
-
-    def checkpoint(self) -> str:
-        try:
-            return json.dumps(self.queue)
-        except AttributeError:
-            raise EmptyChannelError()
-
-
-class UniqueInbox(Generic[Value], BaseChannel[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[FrozenSet[Value]]:
-        """The type of the value stored in the channel."""
-        return FrozenSet[self.typ]  # type: ignore[name-defined]
-
-    @property
-    def UpdateType(self) -> Any:
-        """The type of the update received by the channel."""
-        return Union[self.typ, Sequence[self.typ]]  # type: ignore[name-defined]
-
-    @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
-        empty = self.__class__(self.typ)
-        if checkpoint is not None:
-            empty.queue = frozenset(json.loads(checkpoint))
-        try:
-            yield empty
-        finally:
-            try:
-                del empty.queue
-            except AttributeError:
-                pass
-
-    def update(self, values: Sequence[Value | list[Value]]) -> None:
-        self.queue = frozenset(flatten(values))
-
-    def get(self) -> FrozenSet[Value]:
-        try:
-            return self.queue
-        except AttributeError:
-            raise EmptyChannelError()
-
-    def checkpoint(self) -> str:
-        try:
-            return json.dumps(self.queue)
-        except AttributeError:
-            raise EmptyChannelError()

permchain/channels/last_value.py

@@ -1,4 +1,3 @@
-import json
 from contextlib import contextmanager
 from typing import Generator, Generic, Optional, Sequence, Type
 
@@ -12,7 +11,7 @@ from permchain.channels.base import (
 )
 
 
-class LastValue(Generic[Value], BaseChannel[Value, Value]):
+class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the last value received, can receive at most one value per step."""
 
     def __init__(self, typ: Type[Value]) -> None:
@@ -29,10 +28,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]):
         return self.typ
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ)
         if checkpoint is not None:
-            empty.value = json.loads(checkpoint)
+            empty.value = checkpoint
         try:
             yield empty
         finally:
@@ -42,6 +41,8 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]):
                 pass
 
     def update(self, values: Sequence[Value]) -> None:
+        if len(values) == 0:
+            return
         if len(values) != 1:
             raise InvalidUpdateError()
 
@@ -53,8 +54,8 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]):
         except AttributeError:
             raise EmptyChannelError()
 
-    def checkpoint(self) -> str:
+    def checkpoint(self) -> Value:
         try:
-            return json.dumps(self.value)
+            return self.value
         except AttributeError:
             raise EmptyChannelError()

permchain/channels/topic.py

@@ -0,0 +1,79 @@
+from contextlib import contextmanager
+from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, Union
+
+from typing_extensions import Self
+
+from permchain.channels.base import BaseChannel, 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 Topic(
+    Generic[Value],
+    BaseChannel[Sequence[Value], Value | list[Value], tuple[set[Value], list[Value]]],
+):
+    """A configurable PubSub Topic.
+
+    Args:
+        typ: The type of the value stored in the channel.
+        unique: Whether to discard duplicate values.
+        accumulate: Whether to accummulate values across steps. If False, the channel will be emptied after each step.
+    """
+
+    def __init__(
+        self, typ: Type[Value], unique: bool = False, accumulate: bool = False
+    ) -> None:
+        # attrs
+        self.typ = typ
+        self.unique = unique
+        self.accumulate = accumulate
+        # state
+        self.seen = set[Value]()
+        self.values = list[Value]()
+
+    @property
+    def ValueType(self) -> Any:
+        """The type of the value stored in the channel."""
+        return Sequence[self.typ]  # type: ignore[name-defined]
+
+    @property
+    def UpdateType(self) -> Type[Value]:
+        """The type of the update received by the channel."""
+        return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
+
+    @contextmanager
+    def empty(
+        self, checkpoint: Optional[tuple[set[Value], list[Value]]] = None
+    ) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ, self.unique, self.accumulate)
+        if checkpoint is not None:
+            empty.seen = checkpoint[0]
+            empty.values = checkpoint[1]
+        try:
+            yield empty
+        finally:
+            pass
+
+    def update(self, values: Sequence[Value | list[Value]]) -> None:
+        if not self.accumulate:
+            self.values = list[Value]()
+        if flat_values := flatten(values):
+            if self.unique:
+                for value in flat_values:
+                    if value not in self.seen:
+                        self.seen.add(value)
+                        self.values.append(value)
+            else:
+                self.values.extend(flat_values)
+
+    def get(self) -> Sequence[Value]:
+        return list(self.values)
+
+    def checkpoint(self) -> tuple[set[Value], list[Value]]:
+        return (self.seen, self.values)

permchain/checkpoint/base.py

@@ -0,0 +1,38 @@
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Sequence
+
+from langchain.load.serializable import Serializable
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+from permchain.utils import StrEnum
+
+
+class CheckpointAt(StrEnum):
+    END_OF_STEP = "end_of_step"
+    END_OF_RUN = "end_of_run"
+
+
+class BaseCheckpointAdapter(Serializable, ABC):
+    at: CheckpointAt = CheckpointAt.END_OF_RUN
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return []
+
+    @abstractmethod
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        ...
+
+    @abstractmethod
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        ...
+
+    async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return await asyncio.get_running_loop().run_in_executor(None, self.get, config)
+
+    async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.put, config, checkpoint
+        )

permchain/checkpoint/memory.py

@@ -0,0 +1,29 @@
+from typing import Any, Dict, Mapping, Sequence
+
+from langchain.pydantic_v1 import Field
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+from permchain.checkpoint.base import BaseCheckpointAdapter
+
+
+class MemoryCheckpoint(BaseCheckpointAdapter):
+    storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict)
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id="thread_id",
+                annotation=str,
+                name="Thread ID",
+                description=None,
+                default="",
+            ),
+        ]
+
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return self.storage.get(config["configurable"]["thread_id"], None)
+
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return self.storage.update({config["configurable"]["thread_id"]: checkpoint})

permchain/constants.py

@@ -0,0 +1,4 @@
+CONFIG_KEY_SEND = "__pregel_send"
+CONFIG_KEY_READ = "__pregel_read"
+CHECKPOINT_KEY_VERSION = "__pregel_version"
+CHECKPOINT_KEY_TS = "__pregel_ts"

permchain/pregel/__init__.py

@@ -6,11 +6,14 @@ from collections import defaultdict, deque
 from typing import (
     Any,
     AsyncIterator,
+    Awaitable,
+    Callable,
     Iterator,
     Mapping,
     Optional,
     Sequence,
     Type,
+    Union,
     cast,
     overload,
 )
@@ -26,50 +29,70 @@ from langchain.schema.runnable import (
     RunnablePassthrough,
     RunnableSerializable,
 )
-from langchain.schema.runnable.base import RunnableLike, coerce_to_runnable
+from langchain.schema.runnable.base import Input, Output, coerce_to_runnable
 from langchain.schema.runnable.config import (
     RunnableConfig,
     get_executor_for_config,
     patch_config,
 )
+from langchain.schema.runnable.utils import (
+    ConfigurableFieldSpec,
+    get_unique_config_specs,
+)
 
 from permchain.channels.base import (
     AsyncChannelsManager,
     BaseChannel,
     ChannelsManager,
     EmptyChannelError,
+    create_checkpoint,
 )
-from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
+from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
+from permchain.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.read import ChannelBatch, ChannelInvoke
+from permchain.pregel.reserved import ReservedChannels
 from permchain.pregel.validate import validate_chains_channels
-from permchain.pregel.write import PregelSink
+from permchain.pregel.write import ChannelWrite
+
+WriteValue = Union[
+    Runnable[Input, Output],
+    Callable[[Input], Output],
+    Callable[[Input], Awaitable[Output]],
+    Any,
+]
+
+
+def _coerce_write_value(value: WriteValue) -> Runnable[Input, Output]:
+    if not isinstance(value, Runnable) and not callable(value):
+        return coerce_to_runnable(lambda _: value)
+    return coerce_to_runnable(value)
 
 
 class Channel:
     @overload
     @classmethod
-    def subscribe_to(cls, channels: str, key: Optional[str] = None) -> PregelInvoke:
+    def subscribe_to(cls, channels: str, key: Optional[str] = None) -> ChannelInvoke:
         ...
 
     @overload
     @classmethod
-    def subscribe_to(cls, channels: Sequence[str], key: None = None) -> PregelInvoke:
+    def subscribe_to(cls, channels: Sequence[str], key: None = None) -> ChannelInvoke:
         ...
 
     @classmethod
     def subscribe_to(
         cls, channels: str | Sequence[str], key: Optional[str] = None
-    ) -> PregelInvoke:
+    ) -> ChannelInvoke:
         """Runs process.invoke() each time channels are updated,
         with a dict of the channel values as input."""
         if not isinstance(channels, str) and key is not None:
             raise ValueError(
                 "Can't specify a key when subscribing to multiple channels"
             )
-        return PregelInvoke(
+        return ChannelInvoke(
             channels=cast(
                 Mapping[None, str] | Mapping[str, str],
                 {key: channels}
@@ -79,38 +102,40 @@ class Channel:
         )
 
     @classmethod
-    def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> PregelBatch:
+    def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> ChannelBatch:
         """Runs process.batch() with the content of inbox each time it is updated."""
-        return PregelBatch(channel=inbox, key=key)
+        return ChannelBatch(channel=inbox, key=key)
 
     @classmethod
     def write_to(
         cls,
         *channels: str,
-        **kwargs: RunnableLike,
-    ) -> PregelSink:
+        **kwargs: WriteValue,
+    ) -> ChannelWrite:
         """Writes to channels the result of the lambda, or None to skip writing."""
-        return PregelSink(
+        return ChannelWrite(
             channels=(
                 [(c, RunnablePassthrough()) for c in channels]
-                + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
+                + [(k, _coerce_write_value(v)) for k, v in kwargs.items()]
             )
         )
 
 
 class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
-    channels: Mapping[str, BaseChannel]
+    chains: Mapping[str, ChannelInvoke | ChannelBatch]
 
-    chains: Mapping[str, PregelInvoke | PregelBatch]
+    channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
 
-    output: str | Sequence[str]
+    output: str | Sequence[str] = "output"
 
-    input: str | Sequence[str]
+    input: str | Sequence[str] = "input"
 
     step_timeout: Optional[float] = None
 
     debug: bool = Field(default_factory=get_debug)
 
+    checkpoint: Optional[BaseCheckpointAdapter] = None
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -121,6 +146,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         )
         return values
 
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return get_unique_config_specs(
+            [spec for chain in self.chains.values() for spec in chain.config_specs]
+            + self.checkpoint.config_specs
+            if self.checkpoint is not None
+            else []
+        )
+
     @property
     def InputType(self) -> Any:
         if isinstance(self.input, str):
@@ -162,15 +196,21 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         run_manager: CallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> Iterator[dict[str, Any] | Any]:
+        if config["recursion_limit"] < 1:
+            raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        with ChannelsManager(self.channels) as channels, get_executor_for_config(
-            config
-        ) as executor:
+        checkpoint = (
+            self.checkpoint.get(config) if self.checkpoint is not None else None
+        )
+        with ChannelsManager(
+            self.channels, checkpoint
+        ) as channels, get_executor_for_config(config) as executor:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
                 deque(w for c in input for w in map_input(self.input, c)),
+                config,
+                0,
             )
 
             def read(chan: str) -> Any:
@@ -194,7 +234,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 # execute tasks, and wait for one to fail or all to finish.
                 # each task is independent from all other concurrent tasks
                 done, inflight = concurrent.futures.wait(
-                    (
+                    [
                         executor.submit(
                             proc.invoke,
                             input,
@@ -209,7 +249,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                             ),
                         )
                         for proc, input, _ in next_tasks
-                    ),
+                    ],
                     return_when=concurrent.futures.FIRST_EXCEPTION,
                     timeout=self.step_timeout,
                 )
@@ -218,8 +258,9 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 _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
+                    processes, channels, pending_writes, config, step + 1
                 )
 
                 if self.debug:
@@ -229,25 +270,45 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # save end of step checkpoint
+                if (
+                    self.checkpoint is not None
+                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                ):
+                    checkpoint = create_checkpoint(channels)
+                    self.checkpoint.put(config, checkpoint)
 
                 # if no more tasks, we're done
                 if not next_tasks:
                     break
 
+            # save end of run checkpoint
+            if (
+                self.checkpoint is not None
+                and self.checkpoint.at == CheckpointAt.END_OF_RUN
+            ):
+                checkpoint = create_checkpoint(channels)
+                self.checkpoint.put(config, checkpoint)
+
     async def _atransform(
         self,
         input: AsyncIterator[dict[str, Any] | Any],
         run_manager: AsyncCallbackManagerForChainRun,
         config: RunnableConfig,
     ) -> AsyncIterator[dict[str, Any] | Any]:
+        if config["recursion_limit"] < 1:
+            raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        async with AsyncChannelsManager(self.channels) as channels:
+        checkpoint = (
+            await self.checkpoint.aget(config) if self.checkpoint is not None else None
+        )
+        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
                 deque([w async for c in input for w in map_input(self.input, c)]),
+                config,
+                0,
             )
 
             def read(chan: str) -> Any:
@@ -299,7 +360,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
                 # apply writes to channels, decide on next step
                 next_tasks = _apply_writes_and_prepare_next_tasks(
-                    processes, channels, pending_writes
+                    processes, channels, pending_writes, config, step + 1
                 )
 
                 if self.debug:
@@ -309,12 +370,26 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # save end of step checkpoint
+                if (
+                    self.checkpoint is not None
+                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                ):
+                    checkpoint = create_checkpoint(channels)
+                    await self.checkpoint.aput(config, checkpoint)
 
                 # if no more tasks, we're done
                 if not next_tasks:
                     break
 
+            # save end of run checkpoint
+            if (
+                self.checkpoint is not None
+                and self.checkpoint.at == CheckpointAt.END_OF_RUN
+            ):
+                checkpoint = create_checkpoint(channels)
+                await self.checkpoint.aput(config, checkpoint)
+
     def invoke(
         self,
         input: dict[str, Any] | Any,
@@ -404,15 +479,24 @@ def _interrupt_or_proceed(
 
 
 def _apply_writes_and_prepare_next_tasks(
-    processes: Mapping[str, PregelInvoke | PregelBatch],
+    processes: Mapping[str, ChannelInvoke | ChannelBatch],
     channels: Mapping[str, BaseChannel],
     pending_writes: Sequence[tuple[str, Any]],
+    config: RunnableConfig,
+    for_step: int,
 ) -> 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:
+        if chan in [c.value for c in ReservedChannels]:
+            raise ValueError(f"Can't write to reserved channel {chan}")
         pending_writes_by_channel[chan].append(val)
 
+    # Update reserved channels
+    pending_writes_by_channel[ReservedChannels.is_last_step] = [
+        for_step + 1 == config["recursion_limit"]
+    ]
+
     updated_channels: set[str] = set()
     # Apply writes to channels
     for chan, vals in pending_writes_by_channel.items():
@@ -421,12 +505,16 @@ def _apply_writes_and_prepare_next_tasks(
             updated_channels.add(chan)
         else:
             logger.warning(f"Skipping write for channel {chan} which has no readers")
+    # Channels that weren't updated in this step are notified of a new step
+    for chan in channels:
+        if chan not in updated_channels:
+            channels[chan].update([])
 
     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 name, proc in processes.items():
-        if isinstance(proc, PregelInvoke):
+        if isinstance(proc, ChannelInvoke):
             # If any of the channels read by this process were updated
             if any(chan in updated_channels for chan in proc.channels.values()):
                 # If all channels read by this process have been initialized
@@ -441,7 +529,7 @@ def _apply_writes_and_prepare_next_tasks(
                     tasks.append((proc, val[None], name))
                 else:
                     tasks.append((proc, val, name))
-        elif isinstance(proc, PregelBatch):
+        elif isinstance(proc, ChannelBatch):
             # If the channel read by this process was updated
             if proc.channel in updated_channels:
                 # Here we don't catch EmptyChannelError because the channel

permchain/pregel/constants.py

@@ -1,2 +0,0 @@
-CONFIG_KEY_SEND = "__pregel_send"
-CONFIG_KEY_READ = "__pregel_read"

permchain/pregel/read.py

@@ -5,19 +5,23 @@ from typing import Any, Callable, Mapping, Optional, Sequence
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import (
     Runnable,
-    RunnableBinding,
     RunnableConfig,
     RunnableLambda,
     RunnablePassthrough,
 )
-from langchain.schema.runnable.base import Other, RunnableEach, coerce_to_runnable
+from langchain.schema.runnable.base import (
+    Other,
+    RunnableBindingBase,
+    RunnableEach,
+    coerce_to_runnable,
+)
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
 from permchain.channels.base import BaseChannel
-from permchain.pregel.constants import CONFIG_KEY_READ
+from permchain.constants import CONFIG_KEY_READ
 
 
-class PregelRead(RunnableLambda):
+class ChannelRead(RunnableLambda):
     channel: str
 
     @property
@@ -33,8 +37,7 @@ class PregelRead(RunnableLambda):
         ]
 
     def __init__(self, channel: str) -> None:
-        # TODO remove type ignore after updating langchain
-        super().__init__(func=self._read, afunc=self._aread)  # type: ignore[arg-type]
+        super().__init__(func=self._read, afunc=self._aread)
         self.channel = channel
 
     def _read(self, _: Any, config: RunnableConfig) -> Any:
@@ -58,10 +61,13 @@ class PregelRead(RunnableLambda):
         return read(self.channel)
 
 
-class PregelInvoke(RunnableBinding):
+default_bound = RunnablePassthrough()
+
+
+class ChannelInvoke(RunnableBindingBase):
     channels: Mapping[None, str] | Mapping[str, str]
 
-    bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
+    bound: Runnable[Any, Any] = Field(default=default_bound)
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
 
@@ -76,32 +82,37 @@ class PregelInvoke(RunnableBinding):
     ) -> None:
         super().__init__(
             channels=channels,
-            bound=bound or RunnablePassthrough(),
+            bound=bound or default_bound,
             kwargs=kwargs or {},
             config=config,
             **other_kwargs,
         )
 
-    def join(self, channels: Sequence[str]) -> PregelInvoke:
+    def join(self, channels: Sequence[str]) -> ChannelInvoke:
+        assert isinstance(channels, list) or isinstance(
+            channels, tuple
+        ), "channels must be a list or tuple"
         joiner = RunnablePassthrough.assign(
-            **{chan: PregelRead(chan) for chan in channels}
+            **{chan: ChannelRead(chan) for chan in channels}
         )
-        if isinstance(self.bound, RunnablePassthrough):
-            return PregelInvoke(channels=self.channels, bound=joiner)
+        if self.bound is default_bound:
+            return ChannelInvoke(channels=self.channels, bound=joiner)
         else:
-            return PregelInvoke(channels=self.channels, bound=self.bound | joiner)
+            return ChannelInvoke(channels=self.channels, bound=self.bound | joiner)
 
     def __or__(
         self,
         other: Runnable[Any, Other]
         | Callable[[Any], Other]
         | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
-    ) -> PregelInvoke:
-        if isinstance(self.bound, RunnablePassthrough):
-            return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other))
+    ) -> ChannelInvoke:
+        if self.bound is default_bound:
+            return ChannelInvoke(
+                channels=self.channels, bound=coerce_to_runnable(other)
+            )
         else:
             # delegate to __or__ in self.bound
-            return PregelInvoke(channels=self.channels, bound=self.bound | other)
+            return ChannelInvoke(channels=self.channels, bound=self.bound | other)
 
     def __ror__(
         self,
@@ -112,14 +123,14 @@ class PregelInvoke(RunnableBinding):
         raise NotImplementedError()
 
 
-class PregelBatch(RunnableEach):
+class ChannelBatch(RunnableEach):
     channel: str
 
     key: Optional[str]
 
-    bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
+    bound: Runnable[Any, Any] = Field(default=default_bound)
 
-    def join(self, channels: Sequence[str]) -> PregelBatch:
+    def join(self, channels: Sequence[str]) -> ChannelBatch:
         if self.key is None:
             raise ValueError(
                 "Cannot join() additional channels without a key."
@@ -127,12 +138,12 @@ class PregelBatch(RunnableEach):
             )
 
         joiner = RunnablePassthrough.assign(
-            **{chan: PregelRead(chan) for chan in channels}
+            **{chan: ChannelRead(chan) for chan in channels}
         )
-        if isinstance(self.bound, RunnablePassthrough):
-            return PregelBatch(channel=self.channel, key=self.key, bound=joiner)
+        if self.bound is default_bound:
+            return ChannelBatch(channel=self.channel, key=self.key, bound=joiner)
         else:
-            return PregelBatch(
+            return ChannelBatch(
                 channel=self.channel, key=self.key, bound=self.bound | joiner
             )
 
@@ -141,14 +152,14 @@ class PregelBatch(RunnableEach):
         other: Runnable[Any, Other]
         | Callable[[Any], Other]
         | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
-    ) -> PregelBatch:
-        if isinstance(self.bound, RunnablePassthrough):
-            return PregelBatch(
+    ) -> ChannelBatch:
+        if self.bound is default_bound:
+            return ChannelBatch(
                 channel=self.channel, key=self.key, bound=coerce_to_runnable(other)
             )
         else:
             # delegate to __or__ in self.bound
-            return PregelBatch(
+            return ChannelBatch(
                 channel=self.channel, key=self.key, bound=self.bound | other
             )
 

permchain/pregel/reserved.py

@@ -0,0 +1,8 @@
+from enum import StrEnum
+
+
+class ReservedChannels(StrEnum):
+    """Channels managed by the framework."""
+
+    is_last_step = "is_last_step"
+    """A channel that is True if the current step is the last step, False otherwise."""

permchain/pregel/validate.py

@@ -1,20 +1,28 @@
-from typing import Mapping, Sequence
+from typing import Any, Mapping, Sequence
 
 from permchain.channels.base import BaseChannel
-from permchain.pregel.read import PregelBatch, PregelInvoke
+from permchain.channels.last_value import LastValue
+from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
+from permchain.pregel.read import ChannelBatch, ChannelInvoke
+from permchain.pregel.reserved import ReservedChannels
+
+FORBIDDEN_CHANNEL_NAMES = {
+    CHECKPOINT_KEY_TS,
+    CHECKPOINT_KEY_VERSION,
+}
 
 
 def validate_chains_channels(
-    chains: Mapping[str, PregelInvoke | PregelBatch],
+    chains: Mapping[str, ChannelInvoke | ChannelBatch],
     channels: Mapping[str, BaseChannel],
     input: str | Sequence[str],
     output: str | Sequence[str],
 ) -> None:
     subscribed_channels = set[str]()
     for chain in chains.values():
-        if isinstance(chain, PregelInvoke):
+        if isinstance(chain, ChannelInvoke):
             subscribed_channels.update(chain.channels.values())
-        elif isinstance(chain, PregelBatch):
+        elif isinstance(chain, ChannelBatch):
             subscribed_channels.add(chain.channel)
         else:
             raise TypeError(
@@ -23,12 +31,17 @@ def validate_chains_channels(
 
     for chan in subscribed_channels:
         if chan not in channels:
-            raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
+            channels[chan] = LastValue(Any)
 
     if isinstance(input, str):
+        if input not in channels:
+            channels[input] = LastValue(Any)
         if input not in subscribed_channels:
             raise ValueError(f"Input channel {input} is not subscribed to by any chain")
     else:
+        for chan in input:
+            if chan not in channels:
+                channels[chan] = LastValue(Any)
         if all(chan not in subscribed_channels for chan in input):
             raise ValueError(
                 f"None of the input channels {input} are subscribed to by any chain"
@@ -36,8 +49,16 @@ def validate_chains_channels(
 
     if isinstance(output, str):
         if output not in channels:
-            raise ValueError(f"Output channel {output} is not initialized")
+            channels[output] = LastValue(Any)
     else:
         for chan in output:
             if chan not in channels:
-                raise ValueError(f"Output channel {chan} is not initialized")
+                channels[chan] = LastValue(Any)
+
+    for name in FORBIDDEN_CHANNEL_NAMES:
+        if name in channels:
+            raise ValueError(f"Channel name {name} is reserved")
+
+    for chan in ReservedChannels:
+        if chan not in channels:
+            channels[chan] = LastValue(Any)

permchain/pregel/write.py

@@ -5,30 +5,31 @@ from typing import Any, Callable, Sequence
 from langchain.schema.runnable import (
     Runnable,
     RunnableConfig,
-    RunnableLambda,
+    RunnablePassthrough,
 )
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
-from permchain.pregel.constants import CONFIG_KEY_SEND
+from permchain.constants import CONFIG_KEY_SEND
 
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
 
-# TODO switch to RunnablePassthrough after updating langchain
-class PregelSink(RunnableLambda):
+class ChannelWrite(RunnablePassthrough):
     channels: Sequence[tuple[str, Runnable]]
     """
     Mapping of write channels to Runnables that return the value to be written,
     or None to skip writing.
     """
 
+    class Config:
+        arbitrary_types_allowed = True
+
     def __init__(
         self,
         *,
         channels: Sequence[tuple[str, Runnable]],
     ):
-        super().__init__(func=self._write, afunc=self._awrite)  # type: ignore[arg-type]
-        self.channels = channels
+        super().__init__(func=self._write, afunc=self._awrite, channels=channels)
 
     @property
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
@@ -43,19 +44,16 @@ class PregelSink(RunnableLambda):
         ]
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
-        write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
-
         values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
 
-        write([(chan, val) for chan, val in values if val is not None])
-
-        return input
+        self.do_write(config, **dict(values))
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
-        write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
-
         values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels]
 
-        write([(chan, val) for chan, val in values if val is not None])
+        self.do_write(config, **dict(values))
 
-        return input
+    @staticmethod
+    def do_write(config: RunnableConfig, **values: Any) -> None:
+        write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
+        write([(chan, val) for chan, val in values.items() if val is not None])

permchain/utils.py

@@ -0,0 +1,8 @@
+import enum
+
+
+# Before Python 3.11 native StrEnum is not available
+class StrEnum(str, enum.Enum):
+    """A string enum."""
+
+    pass

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "permchain"
-version = "0.0.3"
+version = "0.0.6"
 description = "permchain"
 authors = []
 license = "MIT"
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/permchain"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-langchain = ">=0.0.313"
+langchain = ">=0.0.335"
 
 
 [tool.poetry.group.test.dependencies]
@@ -25,8 +25,7 @@ syrupy = "^4.0.2"
 pytest-watch = "^4.2.0"
 
 [tool.poetry.group.lint.dependencies]
-ruff = "^0.0.249"
-black = {extras = ["jupyter"], version = "^23.7.0"}
+ruff = "^0.1.4"
 
 [tool.poetry.group.typing.dependencies]
 mypy = "^1.6.0"
@@ -36,8 +35,6 @@ optional = true
 
 [tool.poetry.group.dev.dependencies]
 jupyter = "^1.0.0"
-playwright = "^1.28.0"
-setuptools = "^67.6.1"
 openai = "^0.27.8"
 
 [tool.ruff]

tests/test_channels.py

@@ -1,17 +1,16 @@
 import operator
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, FrozenSet, Generator, Sequence, Union
+from typing import AsyncGenerator, Generator, Sequence, Union
 
 import httpx
 import pytest
 from pytest_mock import MockerFixture
 
-from permchain.channels.archive import UniqueArchive
 from permchain.channels.base import EmptyChannelError, InvalidUpdateError
 from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
-from permchain.channels.inbox import Inbox
 from permchain.channels.last_value import LastValue
+from permchain.channels.topic import Topic
 
 
 def test_last_value() -> None:
@@ -28,6 +27,9 @@ def test_last_value() -> None:
         assert channel.get() == 3
         channel.update([4])
         assert channel.get() == 4
+        checkpoint = channel.checkpoint()
+    with LastValue(int).empty(checkpoint) as channel:
+        assert channel.get() == 4
 
 
 async def test_last_value_async() -> None:
@@ -44,59 +46,157 @@ async def test_last_value_async() -> None:
         assert channel.get() == 3
         channel.update([4])
         assert channel.get() == 4
+        checkpoint = channel.checkpoint()
+    async with LastValue(int).aempty(checkpoint) as channel:
+        assert channel.get() == 4
 
 
-def test_inbox() -> None:
-    with Inbox(str).empty() as channel:
+def test_topic() -> None:
+    with Topic(str).empty() as channel:
         assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, Sequence[str]]
-
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.UpdateType is Union[str, list[str]]
 
         channel.update(["a", "b"])
-        assert channel.get() == ("a", "b")
-        channel.update([["c"], "d"])
-        assert channel.get() == ("c", "d")
+        assert channel.get() == ["a", "b"]
+        channel.update([["c", "d"], "d"])
+        assert channel.get() == ["c", "d", "d"]
+        channel.update([])
+        assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    with Topic(str).empty(checkpoint) as channel:
+        assert channel.get() == ["e"]
 
 
-async def test_inbox_async() -> None:
-    async with Inbox(str).aempty() as channel:
+async def test_topic_async() -> None:
+    async with Topic(str).aempty() as channel:
         assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is Union[str, Sequence[str]]
-
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.UpdateType is Union[str, list[str]]
 
         channel.update(["a", "b"])
-        assert channel.get() == ("a", "b")
-        channel.update(["c"])
-        channel.update([["c"], "d"])
-        assert channel.get() == ("c", "d")
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["b", "c", "d", "d"]
+        channel.update([])
+        assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str).aempty(checkpoint) as channel:
+        assert channel.get() == ["e"]
 
 
-def test_set() -> None:
-    with UniqueArchive(str).empty() as channel:
-        assert channel.ValueType is FrozenSet[str]
-        assert channel.UpdateType is str
+def test_topic_unique() -> None:
+    with Topic(str, unique=True).empty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
 
-        assert channel.get() == frozenset()
         channel.update(["a", "b"])
-        assert channel.get() == frozenset(("a", "b"))
-        channel.update(["b", "c"])
-        assert channel.get() == frozenset(("a", "b", "c"))
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
+        channel.update([])
+        assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, unique=True).empty(checkpoint) as channel:
+        assert channel.get() == ["e"]
+        channel.update(["d", "f"])
+        assert channel.get() == ["f"], "de-dupes from checkpoint"
 
 
-async def test_set_async() -> None:
-    async with UniqueArchive(str).aempty() as channel:
-        assert channel.ValueType is FrozenSet[str]
-        assert channel.UpdateType is str
+async def test_topic_unique_async() -> None:
+    async with Topic(str, unique=True).aempty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
 
-        assert channel.get() == frozenset()
         channel.update(["a", "b"])
-        assert channel.get() == frozenset(("a", "b"))
-        channel.update(["b", "c"])
-        assert channel.get() == frozenset(("a", "b", "c"))
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
+        channel.update([])
+        assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, unique=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["e"]
+        channel.update(["d", "f"])
+        assert channel.get() == ["f"], "de-dupes from checkpoint"
+
+
+def test_topic_accumulate() -> None:
+    with Topic(str, accumulate=True).empty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
+
+        channel.update(["a", "b"])
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update([])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, accumulate=True).empty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update(["e"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
+
+
+async def test_topic_accumulate_async() -> None:
+    async with Topic(str, accumulate=True).aempty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
+
+        channel.update(["a", "b"])
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update([])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, accumulate=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update(["e"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
+
+
+def test_topic_unique_accumulate() -> None:
+    with Topic(str, unique=True, accumulate=True).empty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
+
+        channel.update(["a", "b"])
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update([])
+        assert channel.get() == ["a", "b", "c", "d"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, unique=True, accumulate=True).empty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update(["d", "e"])
+        assert channel.get() == ["a", "b", "c", "d", "e"]
+
+
+async def test_topic_unique_accumulate_async() -> None:
+    async with Topic(str, unique=True, accumulate=True).aempty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, list[str]]
+
+        channel.update(["a", "b"])
+        assert channel.get() == ["a", "b"]
+        channel.update(["b", ["c", "d"], "d"])
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update([])
+        assert channel.get() == ["a", "b", "c", "d"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, unique=True, accumulate=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update(["d", "e"])
+        assert channel.get() == ["a", "b", "c", "d", "e"]
 
 
 def test_binop() -> None:
@@ -104,13 +204,15 @@ def test_binop() -> None:
         assert channel.ValueType is int
         assert channel.UpdateType is int
 
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.get() == 0
 
         channel.update([1, 2, 3])
         assert channel.get() == 6
         channel.update([4])
         assert channel.get() == 10
+        checkpoint = channel.checkpoint()
+    with BinaryOperatorAggregate(int, operator.add).empty(checkpoint) as channel:
+        assert channel.get() == 10
 
 
 async def test_binop_async() -> None:
@@ -118,13 +220,15 @@ async def test_binop_async() -> None:
         assert channel.ValueType is int
         assert channel.UpdateType is int
 
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.get() == 0
 
         channel.update([1, 2, 3])
         assert channel.get() == 6
         channel.update([4])
         assert channel.get() == 10
+        checkpoint = channel.checkpoint()
+    async with BinaryOperatorAggregate(int, operator.add).aempty(checkpoint) as channel:
+        assert channel.get() == 10
 
 
 def test_ctx_manager(mocker: MockerFixture) -> None:
@@ -167,6 +271,9 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
         with pytest.raises(InvalidUpdateError):
             channel.update([5])  # type: ignore
 
+        with pytest.raises(EmptyChannelError):
+            channel.checkpoint()
+
 
 async def test_ctx_manager_async(mocker: MockerFixture) -> None:
     setup = mocker.Mock()

tests/test_pregel.py

@@ -1,3 +1,4 @@
+import operator
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
@@ -9,9 +10,12 @@ from pytest_mock import MockerFixture
 
 from permchain import Channel, Pregel
 from permchain.channels.base import InvalidUpdateError
+from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
-from permchain.channels.inbox import Inbox
 from permchain.channels.last_value import LastValue
+from permchain.channels.topic import Topic
+from permchain.checkpoint.memory import MemoryCheckpoint
+from permchain.pregel.reserved import ReservedChannels
 
 
 def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -35,6 +39,57 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     assert app.invoke(2) == 3
 
 
+def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert app.invoke(2) == 3
+
+
+def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = (
+        Channel.subscribe_to("input")
+        | add_one
+        | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
+    )
+
+    app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"])
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {
+            "output": {"title": "Output"},
+            "fixed": {"title": "Fixed"},
+            "output_plus_one": {"title": "Output Plus One"},
+        },
+    }
+    assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
+
+
+def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
+
+    chain = (
+        Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
+        | add_one
+        | Channel.write_to("output")
+    )
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert app.invoke(2) == {"input": 3, "is_last_step": False}
+    assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True}
+
+
 def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -43,19 +98,14 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
         chains={
             "one": chain,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
         output=["output"],
     )
 
-    assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
+    assert app.input_schema.schema() == {"title": "PregelInput"}
     assert app.output_schema.schema() == {
         "title": "PregelOutput",
         "type": "object",
-        "properties": {"output": {"title": "Output", "type": "integer"}},
+        "properties": {"output": {"title": "Output"}},
     }
     assert app.invoke(2) == {"output": 3}
 
@@ -68,10 +118,6 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
         chains={
             "one": chain,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
         input=["input"],
         output=["output"],
     )
@@ -79,12 +125,12 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
     assert app.input_schema.schema() == {
         "title": "PregelInput",
         "type": "object",
-        "properties": {"input": {"title": "Input", "type": "integer"}},
+        "properties": {"input": {"title": "Input"}},
     }
     assert app.output_schema.schema() == {
         "title": "PregelOutput",
         "type": "object",
-        "properties": {"output": {"title": "Output", "type": "integer"}},
+        "properties": {"output": {"title": "Output"}},
     }
     assert app.invoke({"input": 2}) == {"output": 3}
 
@@ -98,13 +144,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"inbox": Topic(int)},
     )
 
     assert app.invoke(2) == 4
@@ -119,13 +159,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
+        channels={"inbox": Topic(int)},
         input=["input", "inbox"],
-        output="output",
     )
 
     assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4]  # [12 + 1, 2 + 1 + 1]
@@ -143,16 +178,7 @@ def test_batch_two_processes_in_out() -> None:
         Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
     )
 
-    app = Pregel(
-        chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "one": LastValue(int),
-        },
-        input="input",
-        output="output",
-    )
+    app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
 
     assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
 
@@ -161,20 +187,14 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
     test_size = 100
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chans = {
-        "input": LastValue(int),
-        "output": LastValue(int),
-        "-1": LastValue(int),
-    }
     chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
     for i in range(test_size - 2):
-        chans[str(i)] = LastValue(int)
         chains[str(i)] = (
             Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
         )
     chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
 
-    app = Pregel(chains=chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains)
 
     for _ in range(10):
         assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
@@ -189,22 +209,16 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
     test_size = 100
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chans = {
-        "input": LastValue(int),
-        "output": LastValue(int),
-        "-1": LastValue(int),
-    }
     chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
     for i in range(test_size - 2):
-        chans[str(i)] = LastValue(int)
         chains[str(i)] = (
             Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
         )
     chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
 
-    app = Pregel(chains=chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains)
 
-    for _ in range(10):
+    for _ in range(3):
         assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
             2 + test_size,
             1 + test_size,
@@ -216,11 +230,11 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
     with ThreadPoolExecutor() as executor:
         assert [
             *executor.map(
-                app.batch, [[2, 1, 3, 4, 5]] * 10, [{"recursion_limit": test_size}] * 10
+                app.batch, [[2, 1, 3, 4, 5]] * 3, [{"recursion_limit": test_size}] * 3
             )
         ] == [
             [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
-        ] * 10
+        ] * 3
 
 
 def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
@@ -229,15 +243,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
     chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
     chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
 
-    app = Pregel(
-        chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
-        output="output",
-    )
+    app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
 
     with pytest.raises(InvalidUpdateError):
         # LastValue channels can only be updated once per iteration
@@ -252,16 +258,49 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"output": Topic(int)},
     )
 
     # An Inbox channel accumulates updates into a sequence
-    assert app.invoke(2) == (3, 3)
+    assert app.invoke(2) == [3, 3]
+
+
+def test_invoke_checkpoint(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
+
+    def raise_if_above_10(input: int) -> int:
+        if input > 10:
+            raise ValueError("Input is too large")
+        return input
+
+    chain_one = (
+        Channel.subscribe_to(["input"]).join(["total"])
+        | add_one
+        | Channel.write_to("output", "total")
+        | raise_if_above_10
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpoint=MemoryCheckpoint(),
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 2
+    # total is now 2, so output is 2+3=5
+    assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        app.invoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    assert app.checkpoint.get({"configurable": {"thread_id": "2"}}).get("total") == 5
 
 
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
@@ -280,13 +319,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
             "chain_three": chain_three,
             "chain_four": chain_four,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"inbox": Topic(int)},
     )
 
     # Then invoke app
@@ -306,13 +339,7 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
     inner_app = Pregel(
         chains={
             "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
-        },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
-        output="output",
+        }
     )
 
     chain_one = (
@@ -334,14 +361,7 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
             "chain_two": chain_two,
             "chain_three": chain_three,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox_one": Inbox(int),
-            "outbox_one": LastValue(int),
-        },
-        input="input",
-        output="output",
+        channels={"inbox_one": Topic(int)},
     )
 
     for _ in range(10):
@@ -363,13 +383,6 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "between": LastValue(int),
-        },
-        input="input",
-        output="output",
     )
 
     assert [c for c in app.stream(2)] == [3, 4]
@@ -382,13 +395,6 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "between": LastValue(int),
-        },
-        input="input",
-        output="output",
     )
 
     # It finishes executing (once no more messages being published)
@@ -405,13 +411,6 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
     with pytest.raises(ValueError):
         Pregel(
             chains={"chain_one": chain_one, "chain_two": chain_two},
-            channels={
-                "input": LastValue(int),
-                "output": LastValue(int),
-                "between": LastValue(int),
-            },
-            input="input",
-            output="output",
         )
 
 
@@ -436,12 +435,9 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
+            "inbox": Topic(int),
             "ctx": Context(an_int, typ=int),
         },
-        input="input",
         output=["inbox", "output"],
     )
 
@@ -451,7 +447,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
         assert setup.call_count == 1, "Expected setup to be called once"
         assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
-            assert chunk == {"inbox": (3,)}
+            assert chunk == {"inbox": [3]}
         elif i == 1:
             assert chunk == {"output": 4}
         else:

tests/test_pregel_async.py

@@ -1,4 +1,5 @@
 import asyncio
+import operator
 from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncGenerator, AsyncIterator, Generator
 
@@ -8,9 +9,12 @@ from pytest_mock import MockerFixture
 
 from permchain import Channel, Pregel
 from permchain.channels.base import InvalidUpdateError
+from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
-from permchain.channels.inbox import Inbox
 from permchain.channels.last_value import LastValue
+from permchain.channels.topic import Topic
+from permchain.checkpoint.memory import MemoryCheckpoint
+from permchain.pregel.reserved import ReservedChannels
 
 
 async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -29,9 +33,69 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
         output="output",
     )
 
+    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_implicit_channels(
+    mocker: MockerFixture
+) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert await app.ainvoke(2) == 3
+
+
+async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = (
+        Channel.subscribe_to("input")
+        | add_one
+        | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
+    )
+
+    app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"])
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {
+            "output": {"title": "Output"},
+            "fixed": {"title": "Fixed"},
+            "output_plus_one": {"title": "Output Plus One"},
+        },
+    }
+    assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
+
+
+async def test_invoke_single_process_in_out_reserved_is_last(
+    mocker: MockerFixture
+) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
+
+    chain = (
+        Channel.subscribe_to(["input"]).join([ReservedChannels.is_last_step])
+        | add_one
+        | Channel.write_to("output")
+    )
+
+    app = Pregel(chains={"one": chain})
+
+    assert app.input_schema.schema() == {"title": "PregelInput"}
+    assert app.output_schema.schema() == {"title": "PregelOutput"}
+    assert await app.ainvoke(2) == {"input": 3, "is_last_step": False}
+    assert await app.ainvoke(2, {"recursion_limit": 1}) == {
+        "input": 3,
+        "is_last_step": True,
+    }
+
+
 async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -40,19 +104,14 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
         chains={
             "one": chain,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
         output=["output"],
     )
 
-    assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
+    assert app.input_schema.schema() == {"title": "PregelInput"}
     assert app.output_schema.schema() == {
         "title": "PregelOutput",
         "type": "object",
-        "properties": {"output": {"title": "Output", "type": "integer"}},
+        "properties": {"output": {"title": "Output"}},
     }
     assert await app.ainvoke(2) == {"output": 3}
 
@@ -65,10 +124,6 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
         chains={
             "one": chain,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
         input=["input"],
         output=["output"],
     )
@@ -76,12 +131,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
     assert app.input_schema.schema() == {
         "title": "PregelInput",
         "type": "object",
-        "properties": {"input": {"title": "Input", "type": "integer"}},
+        "properties": {"input": {"title": "Input"}},
     }
     assert app.output_schema.schema() == {
         "title": "PregelOutput",
         "type": "object",
-        "properties": {"output": {"title": "Output", "type": "integer"}},
+        "properties": {"output": {"title": "Output"}},
     }
     assert await app.ainvoke({"input": 2}) == {"output": 3}
 
@@ -95,13 +150,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"inbox": Topic(int)},
     )
 
     assert await app.ainvoke(2) == 4
@@ -116,13 +165,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
 
     pubsub = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
+        channels={"inbox": Topic(int)},
         input=["input", "inbox"],
-        output="output",
     )
 
     # [12 + 1, 2 + 1 + 1]
@@ -143,13 +187,7 @@ async def test_batch_two_processes_in_out() -> None:
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "one": LastValue(int),
-        },
-        input="input",
-        output="output",
+        channels={"one": LastValue(int)},
     )
 
     assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
@@ -159,20 +197,14 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
     test_size = 100
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chans = {
-        "input": LastValue(int),
-        "output": LastValue(int),
-        "-1": LastValue(int),
-    }
     chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
     for i in range(test_size - 2):
-        chans[str(i)] = LastValue(int)
         chains[str(i)] = (
             Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
         )
     chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
 
-    app = Pregel(chains=chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains)
 
     # No state is left over from previous invocations
     for _ in range(10):
@@ -188,23 +220,17 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
     test_size = 100
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
 
-    chans = {
-        "input": LastValue(int),
-        "output": LastValue(int),
-        "-1": LastValue(int),
-    }
     chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
     for i in range(test_size - 2):
-        chans[str(i)] = LastValue(int)
         chains[str(i)] = (
             Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
         )
     chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
 
-    app = Pregel(chains=chains, channels=chans, input="input", output="output")
+    app = Pregel(chains=chains)
 
     # No state is left over from previous invocations
-    for _ in range(10):
+    for _ in range(3):
         # Then invoke pubsub
         assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
             2 + test_size,
@@ -216,13 +242,10 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
 
     # Concurrent invocations do not interfere with each other
     assert await asyncio.gather(
-        *(
-            app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size})
-            for _ in range(10)
-        )
+        *(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3))
     ) == [
         [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
-        for _ in range(10)
+        for _ in range(3)
     ]
 
 
@@ -234,15 +257,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
     chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
     chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
 
-    app = Pregel(
-        chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
-        output="output",
-    )
+    app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
 
     with pytest.raises(InvalidUpdateError):
         # LastValue channels can only be updated once per iteration
@@ -257,16 +272,59 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
 
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"output": Topic(int)},
     )
 
-    # An Inbox channel accumulates updates into a sequence
-    assert await app.ainvoke(2) == (3, 3)
+    # An Topic channel accumulates updates into a sequence
+    assert await app.ainvoke(2) == [3, 3]
+
+
+async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
+
+    def raise_if_above_10(input: int) -> int:
+        if input > 10:
+            raise ValueError("Input is too large")
+        return input
+
+    chain_one = (
+        Channel.subscribe_to(["input"]).join(["total"])
+        | add_one
+        | Channel.write_to("output", "total")
+        | raise_if_above_10
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpoint=MemoryCheckpoint(),
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 2
+    # total is now 2, so output is 2+3=5
+    assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "2"}})).get(
+        "total"
+    ) == 5
 
 
 async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
@@ -285,13 +343,7 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
             "chain_three": chain_three,
             "chain_four": chain_four,
         },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
-        },
-        input="input",
-        output="output",
+        channels={"inbox": Topic(int)},
     )
 
     # Then invoke app
@@ -312,13 +364,7 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
     inner_app = Pregel(
         chains={
             "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
-        },
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-        },
-        input="input",
-        output="output",
+        }
     )
 
     chain_one = (
@@ -341,13 +387,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
             "chain_three": chain_three,
         },
         channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox_one": Inbox(int),
+            "inbox_one": Topic(int),
             "outbox_one": LastValue(int),
         },
-        input="input",
-        output="output",
     )
 
     # Then invoke pubsub
@@ -369,16 +411,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
     )
     chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
 
-    app = Pregel(
-        chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "between": LastValue(int),
-        },
-        input="input",
-        output="output",
-    )
+    app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
 
     # Then invoke pubsub
     assert [c async for c in app.astream(2)] == [3, 4]
@@ -389,20 +422,10 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
     chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
     chain_two = Channel.subscribe_to("between") | add_one
 
-    app = Pregel(
-        chains={"chain_one": chain_one, "chain_two": chain_two},
-        channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "between": LastValue(int),
-        },
-        input="input",
-        output="output",
-    )
+    app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
 
-    # Then invoke pubsub
     # It finishes executing (once no more messages being published)
-    # but returns nothing, as nothing was published to OUT topic
+    # but returns nothing, as nothing was published to "output" topic
     assert await app.ainvoke(2) is None
 
 
@@ -437,12 +460,9 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     app = Pregel(
         chains={"chain_one": chain_one, "chain_two": chain_two},
         channels={
-            "input": LastValue(int),
-            "output": LastValue(int),
-            "inbox": Inbox(int),
+            "inbox": Topic(int),
             "ctx": Context(an_int, an_int_async, typ=int),
         },
-        input="input",
         output=["inbox", "output"],
     )
 
@@ -462,7 +482,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
         assert setup_async.call_count == 1, "Expected setup to be called once"
         assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
-            assert chunk == {"inbox": (3,)}
+            assert chunk == {"inbox": [3]}
         elif i == 1:
             assert chunk == {"output": 4}
         else: