Lint

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

Makefile

@@ -27,12 +27,12 @@ test_watch:
 # Define a variable for Python and notebook files.
 PYTHON_FILES=.
 lint format: PYTHON_FILES=.
-lint_diff format_diff: PYTHON_FILES=$(shell git diff --relative=libs/langchain --name-only --diff-filter=d master | grep -E '\.py$$|\.ipynb$$')
+lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --diff-filter=d master | grep -E '\.py$$|\.ipynb$$')
 
 lint lint_diff:
 	poetry run ruff .
 	poetry run black $(PYTHON_FILES) --check
-	# poetry run mypy $(PYTHON_FILES)
+	poetry run mypy $(PYTHON_FILES)
 
 format format_diff:
 	poetry run black $(PYTHON_FILES)

README.md

@@ -4,39 +4,96 @@
 
 `pip install permchain`
 
-## Usage
+## Overview
+
+PermChain is an alpha-stage library for building stateful, multi-actor applications with LLMs. It extends the [LangChain Expression Language](https://python.langchain.com/docs/expression_language/) with the ability to coordinate multiple chains (or actors) across multiple steps of computation. It is inspired by [Pregel](https://research.google/pubs/pub37252/) and [Apache Beam](https://beam.apache.org/).
+
+Some of the use cases are:
+
+- Recursive/iterative LLM chains
+- LLM chains with persistent state/memory
+- LLM agents
+- Multi-agent simulations
+- ...and more!
+
+## How it works
+
+### Channels
+
+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)`
+- `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)`
+
+### Chains
+
+Chains are LCEL Runnables which subscribe to one or more channels, and write to one or more channels. Any valid LCEL expression can be used as a chain. Chains can be combined into a Pregel application, which coordinates the execution of the chains across multiple steps.
+
+### Pregel
+
+Pregel combines multiple chains (or actors) into a single application. It coordinates the execution of the chains across multiple steps, following the Pregel/Bulk Synchronous Parallel model. Each step consists of three phases:
+
+- **Plan**: Determine which chains to execute in this step, ie. the chains that subscribe to channels updated in the previous step (or, in the first step, chains that subscribe to input channels)
+- **Execution**: Execute those chains in parallel, until all complete, or one fails, or a timeout is reached. Any channel updates are invisible to other chains until the next step.
+- **Update**: Update the channels with the values written by the chains in this step.
+
+Repeat until no chains are planned for execution, or a maximum number of steps is reached.
+
+## Example
 
 ```python
-from permchain import InMemoryPubSubConnection, PubSub, Topic
+from permchain import Channel, Pregel
+from permchain.channels import LastValue
 
-topic_one = Topic("one")
-chain_one = Topic.IN.subscribe() | (lambda x: x + 'b') | topic_one.publish()
-chain_two = topic_one.subscribe() | (lambda x: x + 'c') | Topic.OUT.publish()
+grow_value = (
+    Channel.subscribe_to("value")
+    | (lambda x: x + x)
+    | Channel.write_to(value=lambda x: x if len(x) < 10 else None)
+)
 
-conn = InMemoryPubSubConnection()
-pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+app = Pregel(
+    chains={"grow_value": grow_value},
+    channels={"value": LastValue(str)},
+    input="value",
+    output="value",
+)
+
+assert app.invoke("a") == "aaaaaaaa"
 
-assert pubsub.invoke('a') == ['abc']
 ```
 
-Check `tests` and `examples` for more examples.
+Check `examples` for more examples.
 
 ## Near-term Roadmap
 
-- [x] Add initial retry support (pending changes in `langchain`)
-- [x] Implement OUT as regular topic
-- [x] Implement IN as regular topic
-- [x] Add Connection.peek() to monitor past messages from all topics
-- [x] Enable resuming PubSub from the "middle" of the computation
-- [x] Add test for .peek()
-- [x] Add "wait until topic X is done" pattern, aka. `Topic.join()`
-- [ ] Move tracking of inflight processes/messages to Connection
-  - [ ] Use this to build retry mechanism, where any inflight messages are moved back to the respective topics when restarting
-  - [ ] But this would require being able to replay a message for a single listener only, which maybe requires a larger redesign of PubSub<>Connection contract than what I wanted to do here
-- [ ] Detect cycles (aka. infinite loops) and throw an error
-  - [ ] Allow user to catch that error (by subcribing to an error topic?)
-- [ ] Add example for "human in the loop" pattern, one of the two below
-  - [ ] Example with one permchain, which runs until it produces either 1. request for input or 2. output. The consumer code then gets the needed info, and restarts the permchain with answer, and same state id
-  - [ ] Allow interrupting execution by breaking out of the iterator returned by .stream()
-    - [ ] Build example showing a simple "human in the loop" pattern using this, ie. if a certain message asking for input is published the consumer of the iterator breaks out, does something and then restarts it
-- [ ] Add Redis-backed Connection implementation
+- [x] Iterate on API
+  - [x] do we want api to receive output from multiple channels in invoke()
+  - [x] do we want api to send input to multiple channels in invoke()
+  - [x] Finish updating tests to new API
+- [x] Implement input_schema and output_schema in Pregel
+- [ ] More tests
+  - [x] Test different input and output types (str, str sequence)
+  - [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 more examples
+  - [ ] human in the loop
+  - [ ] combine documents
+  - [ ] agent executor
+  - [ ] run over dataset
+- [ ] Fault tolerance
+  - [ ] Retry individual processes in a step
+  - [ ] Retry entire step?
+- [ ] Pregel.stream_log to contain additional keys specific to Pregel
+  - [ ] tasks: inputs of each chain in each step, keyed by {name}:{step}
+  - [ ] task_results: same as above but outputs
+  - [ ] channels: channel values at end of each step, keyed by {name}:{step}

examples/combine_docs.ipynb

@@ -0,0 +1,369 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "780c1001-557c-4b03-8ebd-a2a381d5f85d",
+   "metadata": {},
+   "source": [
+    "# Combine Docs\n",
+    "\n",
+    "PermChain is a great choice for implementating workflows that involve operating over longer documents because of its recursive nature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "624c452c-ddd5-4390-9065-7ec55dc64b96",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.chat_models.openai import ChatOpenAI\n",
+    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
+    "from langchain.schema.output_parser import StrOutputParser\n",
+    "from langchain.schema.runnable import Runnable, RunnablePassthrough\n",
+    "from langchain.schema.output_parser import StrOutputParser\n",
+    "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "271728d7-b3c8-4ec6-a728-19835e282ec3",
+   "metadata": {},
+   "source": [
+    "## Stuff Documents\n",
+    "\n",
+    "Stuff documents is simple - just a chain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "0462aff0-1b88-49cc-bfe2-3c169d5e1d63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.schema.runnable import RunnableLambda"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "59d6430b-c113-4498-9ffc-f4623f7a0b5c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template=\"{page_content}\")\n",
+    "\n",
+    "_combine_documents = RunnableLambda(\n",
+    "    lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n",
+    ").map() | (lambda x: \"\\n\\n\".join(x))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "29b2668d-e4a6-4876-9b04-bdc841774c62",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "docs = [\n",
+    "    Document(page_content=\"Harrison used to work at Kensho\"),\n",
+    "    Document(page_content=\"Ankush worked at Facebook\"),\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "17da58b7-8685-4d0a-9a47-c398c085d477",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "stuff_chain = (\n",
+    "    {\n",
+    "        \"question\": lambda x: x[\"question\"],\n",
+    "        \"context\": (lambda x: x[\"docs\"]) | _combine_documents,\n",
+    "    }\n",
+    "    | ChatPromptTemplate.from_messages(\n",
+    "        [\n",
+    "            (\n",
+    "                \"system\",\n",
+    "                \"Answer user questions based on the following documents:\\n\\n{context}\",\n",
+    "            ),\n",
+    "            (\"human\", \"{question}\"),\n",
+    "        ]\n",
+    "    )\n",
+    "    | ChatOpenAI()\n",
+    "    | StrOutputParser()\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "87295b71-0afc-4901-b57c-a7b945aa4bd9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Harrison used to work at Kensho.'"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fff324c1-7fbf-41e5-861f-a10ba0112dbd",
+   "metadata": {},
+   "source": [
+    "## Reduce Documents\n",
+    "\n",
+    "Reduce documents tries to merge documents recursively."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b15f5abb-1cfe-4965-a021-c891506c5dd2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "many_docs = docs * 5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "ccad04a3-fd3f-4e73-b895-29e53535f000",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _split_list_of_docs(docs, max_length=70):\n",
+    "    new_result_doc_list = []\n",
+    "    _sub_result_docs = []\n",
+    "    for doc in docs:\n",
+    "        _sub_result_docs.append(doc)\n",
+    "        _num_tokens = sum([len(d.page_content) for d in _sub_result_docs])\n",
+    "        if _num_tokens > max_length:\n",
+    "            if len(_sub_result_docs) == 1:\n",
+    "                raise ValueError(\n",
+    "                    \"A single document was longer than the context length,\"\n",
+    "                    \" we cannot handle this.\"\n",
+    "                )\n",
+    "            new_result_doc_list.append(_sub_result_docs[:-1])\n",
+    "            _sub_result_docs = _sub_result_docs[-1:]\n",
+    "    new_result_doc_list.append(_sub_result_docs)\n",
+    "    return new_result_doc_list"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "11cfd337-9f3b-4b26-ba30-251e17b18994",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[[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')]]"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Just to show what its like split\n",
+    "split_docs = _split_list_of_docs(many_docs)\n",
+    "split_docs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "8d524ba6-0939-4a5d-8db0-4fa1ef06eaeb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "channels = {\n",
+    "    # input\n",
+    "    \"question\": LastValue(str),\n",
+    "    \"docs\": Inbox(Document),\n",
+    "    # intermediate\n",
+    "    \"docs_to_finalize\": Inbox(Document),\n",
+    "    # output\n",
+    "    \"answer\": LastValue(str),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "67370694-86f4-4b64-9d4f-38b2e306abeb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def decide(docs: list[Document]) -> Runnable:\n",
+    "    if len(_split_list_of_docs(docs)) > 1:\n",
+    "        # send back to the beginning if we still need to collapse more\n",
+    "        return Channel.write_to(\"docs\")\n",
+    "    else:\n",
+    "        # send to the finalizer if we're ready to produce final answer\n",
+    "        return Channel.write_to(\"docs_to_finalize\")\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",
+    "    | 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",
+    ")\n",
+    "\n",
+    "# Convert final set of docs to an answer\n",
+    "finalize = (\n",
+    "    Channel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
+    "    | stuff_chain\n",
+    "    | Channel.write_to(\"answer\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reduce_chain = Pregel(\n",
+    "    chains={\n",
+    "        \"collapse\": collapse,\n",
+    "        \"finalize\": finalize,\n",
+    "    },\n",
+    "    channels=channels,\n",
+    "    input=[\"question\", \"docs\"],\n",
+    "    output=\"answer\",\n",
+    "    debug=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "69fcb829-3dae-432a-8db3-11bbb179a7d2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "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[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[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[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[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[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[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[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",
+      " 'question': 'where did harrison work'}\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'Harrison worked at Kensho.'"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "265b29cd-d4f4-4e48-8d4e-b759e909ac2e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/draft-revise-loop.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from langchain.chat_models.openai import ChatOpenAI
+from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain.prompts import SystemMessagePromptTemplate
+from langchain.schema.output_parser import StrOutputParser
+
+from permchain import Channel, Pregel
+from permchain.channels import LastValue
+
+# prompts
+
+drafter_prompt = (
+    SystemMessagePromptTemplate.from_template(
+        "You are an expert on turtles, who likes to write in pirate-speak. You have been tasked by your editor with drafting a 100-word article answering the following question."
+    )
+    + "Question:\n\n{question}"
+)
+
+reviser_prompt = (
+    SystemMessagePromptTemplate.from_template(
+        "You are an expert on turtles. You have been tasked by your editor with revising the following draft, which was written by a non-expert. You may follow the editor's notes or not, as you see fit."
+    )
+    + "Draft:\n\n{draft}"
+    + "Editor's notes:\n\n{notes}"
+)
+
+editor_prompt = (
+    SystemMessagePromptTemplate.from_template(
+        "You are an editor. You have been tasked with editing the following draft, which was written by a non-expert. Please accept the draft if it is good enough to publish, or send it for revision, along with your notes to guide the revision."
+    )
+    + "Draft:\n\n{draft}"
+)
+
+editor_functions = [
+    {
+        "name": "revise",
+        "description": "Sends the draft for revision",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "notes": {
+                    "type": "string",
+                    "description": "The editor's notes to guide the revision.",
+                },
+            },
+        },
+    },
+    {
+        "name": "accept",
+        "description": "Accepts the draft",
+        "parameters": {
+            "type": "object",
+            "properties": {"ready": {"const": True}},
+        },
+    },
+]
+
+# llms
+
+gpt3 = ChatOpenAI(model="gpt-3.5-turbo")
+gpt4 = ChatOpenAI(model="gpt-4")
+
+# chains
+
+drafter_chain = drafter_prompt | gpt3 | StrOutputParser()
+
+editor_chain = (
+    editor_prompt
+    | gpt4.bind(functions=editor_functions)
+    | JsonOutputFunctionsParser(args_only=False)
+)
+
+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")
+)
+
+editor = (
+    # subscribe to draft channel as a dict with a single key, "draft"
+    Channel.subscribe_to(["draft"])
+    | 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
+    )
+)
+
+reviser = (
+    # subscribe to new values of "notes" channel,
+    # and join them with the input value (question) and "draft"
+    Channel.subscribe_to(["notes"]).join(["question", "draft"])
+    | reviser_chain
+    | Channel.write_to("draft")
+)
+
+draft_revise_loop = Pregel(
+    channels=channels,
+    chains={
+        "drafter": drafter,
+        "editor": editor,
+        "reviser": reviser,
+    },
+    # input will be a dict with a single key, "question"
+    input=["question"],
+    # output will be a dict with keys "draft" and "notes"
+    output=["draft", "notes"],
+    # debug logging
+    debug=True,
+)
+
+# 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())

examples/readme.py

@@ -0,0 +1,17 @@
+from permchain import Channel, Pregel
+from permchain.channels import LastValue
+
+grow_value = (
+    Channel.subscribe_to("value")
+    | (lambda x: x + x)
+    | Channel.write_to(value=lambda x: x if len(x) < 10 else None)
+)
+
+app = Pregel(
+    chains={"grow_value": grow_value},
+    channels={"value": LastValue(str)},
+    input="value",
+    output="value",
+)
+
+assert app.invoke("a") == "aaaaaaaa"

examples/recursive-web-loader.py

@@ -0,0 +1,144 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncGenerator, Callable, FrozenSet, Generator, Optional, TypedDict
+
+import httpx
+from langchain.schema import Document
+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
+
+# Load url with sync httpx client
+
+
+@contextmanager
+def httpx_client() -> Generator[httpx.Client, None, None]:
+    with httpx.HTTPTransport(retries=3) as transport, httpx.Client(
+        transport=transport
+    ) as client:
+        yield client
+
+
+class LoadUrlInput(TypedDict):
+    url: str
+    visited: FrozenSet[str]
+    client: httpx.Client
+
+
+def load_url(input: LoadUrlInput) -> str:
+    response = input["client"].get(input["url"])
+    return response.text
+
+
+# Same as above but with async httpx client
+
+
+@asynccontextmanager
+async def httpx_aclient() -> AsyncGenerator[httpx.AsyncClient, None]:
+    async with httpx.AsyncHTTPTransport(retries=3) as transport, httpx.AsyncClient(
+        transport=transport
+    ) as client:
+        yield client
+
+
+class LoadUrlInputAsync(TypedDict):
+    url: str
+    visited: FrozenSet[str]
+    client: httpx.AsyncClient
+
+
+async def load_url_async(input: LoadUrlInputAsync) -> str:
+    response = await input["client"].get(input["url"])
+    return response.text
+
+
+# default metadata extractor copied from langchain.document_loaders
+
+
+def _metadata_extractor(raw_html: str, url: str) -> dict:
+    """Extract metadata from raw html using BeautifulSoup."""
+    metadata = {"source": url}
+
+    try:
+        from bs4 import BeautifulSoup
+    except ImportError:
+        return metadata
+    soup = BeautifulSoup(raw_html, "html.parser")
+    if title := soup.find("title"):
+        metadata["title"] = title.get_text()
+    if description := soup.find("meta", attrs={"name": "description"}):
+        metadata["description"] = description.get("content", None)
+    if html := soup.find("html"):
+        metadata["language"] = html.get("lang", None)
+    return metadata
+
+
+def recursive_web_loader(
+    *,
+    max_depth: int = 2,
+    extractor: Optional[Callable[[str], str]] = None,
+    metadata_extractor: Optional[Callable[[str, str], dict]] = None,
+) -> Pregel:
+    # assign default extractors
+    extractor = extractor or (lambda x: x)
+    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),
+        "client": Context(httpx_client, httpx_aclient),
+    }
+    # the main chain that gets executed recursively
+    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"]
+        )
+        # 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"]
+            ],
+        )
+    )
+    return Pregel(
+        channels=channels,
+        chains={
+            # use the base_url as the first url to visit
+            "input": Channel.subscribe_to("base_url") | Channel.write_to("next_urls"),
+            # add the main chain
+            "visitor": visitor,
+        },
+        # this will accept a string as input
+        input="base_url",
+        # and return a dict with documents and visited set
+        output=["documents", "visited"],
+        # debug logging
+        debug=True,
+    ).with_config({"recursion_limit": max_depth + 1})
+
+
+loader = recursive_web_loader(max_depth=3)
+
+documents = loader.invoke("https://docs.python.org/3.9/")
+
+print(len(documents["documents"]))

examples/runnable-pubsub.py

@@ -1,138 +0,0 @@
-from operator import itemgetter
-from pprint import pprint
-
-from langchain.chat_models.openai import ChatOpenAI
-from langchain.prompts import SystemMessagePromptTemplate
-from langchain.runnables.openai_functions import OpenAIFunctionsRouter
-from langchain.schema.output_parser import StrOutputParser
-
-from permchain.connection_inmemory import InMemoryPubSubConnection
-from permchain.pubsub import PubSub
-from permchain.topic import Topic
-
-drafter_prompt = (
-    SystemMessagePromptTemplate.from_template(
-        "You are an expert on turtles, who likes to write in pirate-speak. You have been tasked by your editor with drafting a 100-word article answering the following question."
-    )
-    + "Question:\n\n{question}"
-)
-
-reviser_prompt = (
-    SystemMessagePromptTemplate.from_template(
-        "You are an expert on turtles. You have been tasked by your editor with revising the following draft, which was written by a non-expert. You may follow the editor's notes or not, as you see fit."
-    )
-    + "Draft:\n\n{draft}"
-    + "Editor's notes:\n\n{notes}"
-)
-
-editor_prompt = (
-    SystemMessagePromptTemplate.from_template(
-        "You are an editor. You have been tasked with editing the following draft, which was written by a non-expert. Please accept the draft if it is good enough to publish, or send it for revision, along with your notes to guide the revision."
-    )
-    + "Draft:\n\n{draft}"
-)
-
-
-drafter_llm = ChatOpenAI(model="gpt-3.5-turbo")
-editor_llm = ChatOpenAI(model="gpt-4")
-
-# create topics
-editor_inbox = Topic("editor_inbox")
-reviser_inbox = Topic("reviser_inbox")
-
-
-# write a first draft
-drafter = (
-    Topic.IN.subscribe()
-    | {"draft": drafter_prompt | drafter_llm | StrOutputParser()}
-    | editor_inbox.publish()
-)
-
-# edit every draft, produce revision notes or accept
-editor = (
-    editor_inbox.subscribe()
-    | editor_prompt
-    | editor_llm.bind(
-        functions=[
-            {
-                "name": "revise",
-                "description": "Sends the draft for revision",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "notes": {
-                            "type": "string",
-                            "description": "The editor's notes to guide the revision.",
-                        },
-                    },
-                },
-            },
-            {
-                "name": "accept",
-                "description": "Accepts the draft",
-                "parameters": {
-                    "type": "object",
-                    "properties": {"ready": {"const": True}},
-                },
-            },
-        ]
-    )
-    | OpenAIFunctionsRouter(
-        {
-            "revise": (
-                {
-                    "notes": itemgetter("notes"),
-                    "draft": editor_inbox.current() | itemgetter("draft"),
-                    "question": Topic.IN.current() | itemgetter("question"),
-                }
-                | reviser_inbox.publish()
-            ),
-            "accept": editor_inbox.current() | Topic.OUT.publish(),
-        },
-    )
-)
-
-# every time revision notes are posted, revise latest draft
-reviser = (
-    reviser_inbox.subscribe()
-    | {"draft": reviser_prompt | drafter_llm | StrOutputParser()}
-    | editor_inbox.publish()
-)
-
-web_researcher = PubSub(
-    processes=(drafter, editor, reviser),
-    connection=InMemoryPubSubConnection(),
-)
-
-# for output in web_researcher.stream({"question": "What food do turtles eat?"}):
-#     print("got output", output)
-
-# print("---done with stream()---")
-
-# pprint(web_researcher.invoke({"question": "What food do turtles eat?"}))
-
-pprint(
-    web_researcher.batch(
-        [{"question": "What food do turtles eat?"}, {"question": "What is art?"}]
-    )
-)
-
-
-# agent = PubSub(
-#     Channel.IN | Channel("planner"),
-#     Channel("executor") | executor | Channel("planner"),
-#     Channel("planner")
-#     | planner
-#     | {"action": Channel("executor"), "finish": Channel.OUT},
-# )
-
-# graph = (
-#     drafter
-#     | editor
-#     | RouterRunnable(
-#         {
-#             "send_for_revision": reviser,
-#             "accept_draft": lambda x: x["draft"],
-#         }
-#     )
-# )

permchain/__init__.py

@@ -1,9 +1,4 @@
-from permchain.connection_inmemory import InMemoryPubSubConnection
-from permchain.pubsub import PubSub
-from permchain.topic import Topic
+from permchain.pregel import Channel, Pregel
+from permchain.pregel.read import PregelRead
 
-__all__ = [
-    "PubSub",
-    "Topic",
-    "InMemoryPubSubConnection",
-]
+__all__ = ["Channel", "Pregel", "PregelRead"]

permchain/channels/__init__.py

@@ -0,0 +1,15 @@
+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
+
+__all__ = [
+    "LastValue",
+    "Inbox",
+    "UniqueInbox",
+    "Archive",
+    "UniqueArchive",
+    "BinaryOperatorAggregate",
+    "Context",
+]

permchain/channels/archive.py

@@ -0,0 +1,94 @@
+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

@@ -0,0 +1,101 @@
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncGenerator,
+    Generator,
+    Generic,
+    Mapping,
+    Optional,
+    Sequence,
+    TypeVar,
+)
+
+from typing_extensions import Self
+
+Value = TypeVar("Value")
+Update = TypeVar("Update")
+
+
+class EmptyChannelError(Exception):
+    """Raised when attempting to get the value of a channel that hasn't been updated
+    for the first time yet."""
+
+    pass
+
+
+class InvalidUpdateError(Exception):
+    """Raised when attempting to update a channel with an invalid sequence of updates."""
+
+    pass
+
+
+class BaseChannel(Generic[Value, Update], ABC):
+    @property
+    @abstractmethod
+    def ValueType(self) -> Any:
+        """The type of the value stored in the channel."""
+
+    @property
+    @abstractmethod
+    def UpdateType(self) -> Any:
+        """The type of the update received by the channel."""
+
+    @contextmanager
+    @abstractmethod
+    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+        """Return a new identical channel, optionally initialized from a checkpoint."""
+
+    @asynccontextmanager
+    async def aempty(
+        self, checkpoint: Optional[str] = None
+    ) -> AsyncGenerator[Self, None]:
+        """Return a new identical channel, optionally initialized from a checkpoint."""
+        with self.empty(checkpoint) as value:
+            yield value
+
+    @abstractmethod
+    def update(self, values: Sequence[Update]) -> None:
+        """Update the channel's value with the given sequence of updates.
+        The order of the updates in the sequence is arbitrary.
+
+        Raises InvalidUpdateError if the sequence of updates is invalid."""
+
+    @abstractmethod
+    def get(self) -> Value:
+        """Return the current value of the channel.
+
+        Raises EmptyChannelError if the channel is empty (never updated yet)."""
+
+    @abstractmethod
+    def checkpoint(self) -> str | None:
+        """Return a string representation of the channel's current state,
+        or None if the channel doesn't support checkpoints.
+
+        Raises EmptyChannelError if the channel is empty (never updated yet)."""
+
+
+@contextmanager
+def ChannelsManager(
+    channels: Mapping[str, BaseChannel]
+) -> 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()}
+    try:
+        yield {k: v.__enter__() for k, v in empty.items()}
+    finally:
+        for v in empty.values():
+            v.__exit__(None, None, None)
+
+
+@asynccontextmanager
+async def AsyncChannelsManager(
+    channels: Mapping[str, BaseChannel]
+) -> 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()}
+    try:
+        yield {k: await v.__aenter__() for k, v in empty.items()}
+    finally:
+        for v in empty.values():
+            await v.__aexit__(None, None, None)

permchain/channels/binop.py

@@ -0,0 +1,65 @@
+import json
+from contextlib import contextmanager
+from typing import Callable, Generator, Generic, Optional, Sequence, Type
+
+from typing_extensions import Self
+
+from permchain.channels.base import BaseChannel, EmptyChannelError, Value
+
+
+class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
+    """Stores the result of applying a binary operator to the current value and each new value.
+
+    ```python
+    import operator
+
+    total = Channels.BinaryOperatorAggregate(int, operator.add)
+    ```
+    """
+
+    def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
+        self.typ = typ
+        self.operator = operator
+
+    @property
+    def ValueType(self) -> Type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @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, self.operator)
+        if checkpoint is not None:
+            empty.value = json.loads(checkpoint)
+        try:
+            yield empty
+        finally:
+            try:
+                del empty.value
+            except AttributeError:
+                pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        if not hasattr(self, "value"):
+            self.value = values[0]
+            values = values[1:]
+
+        for value in values:
+            self.value = self.operator(self.value, value)
+
+    def get(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> str:
+        try:
+            return json.dumps(self.value)
+        except AttributeError:
+            raise EmptyChannelError()

permchain/channels/context.py

@@ -0,0 +1,109 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncContextManager,
+    AsyncGenerator,
+    Callable,
+    ContextManager,
+    Generator,
+    Generic,
+    Optional,
+    Sequence,
+    Type,
+)
+
+from typing_extensions import Self
+
+from permchain.channels.base import (
+    BaseChannel,
+    EmptyChannelError,
+    InvalidUpdateError,
+    Value,
+)
+
+
+class Context(Generic[Value], BaseChannel[Value, None]):
+    """Exposes the value of a context manager, for the duration of an invocation.
+    Context manager is entered before the first step, and exited after the last step.
+    Optionally, provide an equivalent async context manager, which will be used
+    instead for async invocations.
+
+    ```python
+    import httpx
+
+    client = Channels.Context(httpx.Client, httpx.AsyncClient)
+    ```
+    """
+
+    value: Value
+
+    def __init__(
+        self,
+        ctx: Optional[Callable[[], ContextManager[Value]]] = None,
+        actx: Optional[Callable[[], AsyncContextManager[Value]]] = None,
+        typ: Optional[Type[Value]] = None,
+    ) -> None:
+        if ctx is None and actx is None:
+            raise ValueError("Must provide either sync or async context manager.")
+
+        self.typ = typ
+        self.ctx = ctx
+        self.actx = actx
+
+    @property
+    def ValueType(self) -> Any:
+        """The type of the value stored in the channel."""
+        return (
+            self.typ
+            or (self.ctx if hasattr(self.ctx, "__enter__") else None)
+            or (self.actx if hasattr(self.actx, "__aenter__") else None)
+            or None
+        )
+
+    @property
+    def UpdateType(self) -> Type[None]:
+        """The type of the update received by the channel."""
+        raise InvalidUpdateError()
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+        if self.ctx is None:
+            raise ValueError("Cannot enter sync context manager.")
+
+        empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ)
+        # ContextManager doesn't have a checkpoint
+        ctx = self.ctx()
+        empty.value = ctx.__enter__()
+        try:
+            yield empty
+        finally:
+            ctx.__exit__(None, None, None)
+
+    @asynccontextmanager
+    async def aempty(
+        self, checkpoint: Optional[str] = None
+    ) -> AsyncGenerator[Self, None]:
+        if self.actx is not None:
+            empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ)
+            # ContextManager doesn't have a checkpoint
+            actx = self.actx()
+            empty.value = await actx.__aenter__()
+            try:
+                yield empty
+            finally:
+                await actx.__aexit__(None, None, None)
+        else:
+            with self.empty() as empty:
+                yield empty
+
+    def update(self, values: Sequence[None]) -> None:
+        raise InvalidUpdateError()
+
+    def get(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> None:
+        return None

permchain/channels/inbox.py

@@ -0,0 +1,115 @@
+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

@@ -0,0 +1,60 @@
+import json
+from contextlib import contextmanager
+from typing import Generator, Generic, Optional, Sequence, Type
+
+from typing_extensions import Self
+
+from permchain.channels.base import (
+    BaseChannel,
+    EmptyChannelError,
+    InvalidUpdateError,
+    Value,
+)
+
+
+class LastValue(Generic[Value], BaseChannel[Value, Value]):
+    """Stores the last value received, can receive at most one value per step."""
+
+    def __init__(self, typ: Type[Value]) -> None:
+        self.typ = typ
+
+    @property
+    def ValueType(self) -> Type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @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.value = json.loads(checkpoint)
+        try:
+            yield empty
+        finally:
+            try:
+                del empty.value
+            except AttributeError:
+                pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        if len(values) != 1:
+            raise InvalidUpdateError()
+
+        self.value = values[-1]
+
+    def get(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> str:
+        try:
+            return json.dumps(self.value)
+        except AttributeError:
+            raise EmptyChannelError()

permchain/connection.py

@@ -1,65 +0,0 @@
-import asyncio
-from abc import ABC, abstractmethod
-from typing import Any, Callable, Iterator, TypedDict
-
-
-class PubSubMessage(TypedDict):
-    topic: str
-    value: Any
-    published_at: str
-    correlation_id: str
-
-
-PubSubListener = Callable[[PubSubMessage], None]
-
-
-class PubSubConnection(ABC):
-    def full_name(self, prefix: str, *parts: str) -> str:
-        """Return the full topic name for a given prefix and topic name."""
-        return ":".join(map(str, [prefix, *parts]))
-
-    @abstractmethod
-    def observe(self, prefix: str) -> Iterator[PubSubMessage]:
-        """Iterate over messages for all topics under this prefix,
-        without affecting listeners/iterators on each topic.
-        This method waits for new messages to arrive."""
-        ...
-
-    @abstractmethod
-    def iterate(
-        self, prefix: str, topic: str, *, wait: bool
-    ) -> Iterator[PubSubMessage]:
-        """Iterate over all currently queued messages for a topic, consuming them.
-        Optionally wait for new messages to arrive."""
-        ...
-
-    # TODO add aiterate() method
-
-    @abstractmethod
-    def listen(self, prefix: str, topic: str, listeners: list[PubSubListener]) -> None:
-        ...
-
-    async def alisten(
-        self, prefix: str, topic: str, listeners: list[PubSubListener]
-    ) -> None:
-        return await asyncio.get_event_loop().run_in_executor(
-            None, self.listen, prefix, topic, listeners
-        )
-
-    @abstractmethod
-    def send(self, prefix: str, topic: str, value: Any) -> None:
-        ...
-
-    async def asend(self, prefix: str, topic: str, value: Any) -> None:
-        return await asyncio.get_event_loop().run_in_executor(
-            None, self.send, prefix, topic, value
-        )
-
-    @abstractmethod
-    def disconnect(self, name: str) -> None:
-        ...
-
-    async def adisconnect(self, name: str) -> None:
-        return await asyncio.get_event_loop().run_in_executor(
-            None, self.disconnect, name
-        )

permchain/connection_inmemory.py

@@ -1,125 +0,0 @@
-import queue
-import threading
-from collections import defaultdict
-from datetime import datetime
-from typing import Any, Iterator, cast
-
-from permchain.connection import PubSubConnection, PubSubListener, PubSubMessage
-
-
-class IterableQueue(queue.SimpleQueue):
-    done_sentinel = object()
-
-    def put(
-        self, item: PubSubMessage, block: bool = True, timeout: float | None = None
-    ) -> None:
-        return super().put(item, block, timeout)
-
-    def get(
-        self, block: bool = True, timeout: float | None = None
-    ) -> PubSubMessage | object:
-        return super().get(block=block, timeout=timeout)
-
-    def __iter__(self) -> Iterator[PubSubMessage]:
-        return iter(self.get, self.done_sentinel)
-
-    def close(self) -> None:
-        self.put(self.done_sentinel)
-
-
-class InMemoryPubSubConnection(PubSubConnection):
-    clear_on_disconnect: bool
-    logs: defaultdict[str, IterableQueue]
-    topics: defaultdict[str, IterableQueue]
-    listeners: defaultdict[str, list[PubSubListener]]
-    lock: threading.RLock
-
-    def __init__(self, clear_on_disconnect: bool = True) -> None:
-        self.clear_on_disconnect = clear_on_disconnect
-        self.logs = defaultdict(IterableQueue)
-        self.topics = defaultdict(IterableQueue)
-        self.listeners = defaultdict(list)
-        self.lock = threading.RLock()
-
-    def observe(self, prefix: str) -> Iterator[PubSubMessage]:
-        return iter(self.logs[str(prefix)])
-
-    def iterate(
-        self, prefix: str, topic: str, *, wait: bool
-    ) -> Iterator[PubSubMessage]:
-        topic = self.full_name(prefix, topic)
-
-        # This connection doesn't support iterating over topics with listeners connected
-        with self.lock:
-            if self.listeners[topic]:
-                raise RuntimeError(
-                    f"Cannot iterate over topic {topic} while listeners are connected"
-                )
-
-        # If wait is False, add sentinel to queue to ensure the iterator terminates
-        if not wait:
-            self.topics[topic].close()
-
-        return iter(self.topics[topic])
-
-    def listen(self, prefix: str, topic: str, listeners: list[PubSubListener]) -> None:
-        full_name = self.full_name(prefix, topic)
-        self.disconnect(full_name)
-
-        with self.lock:
-            # Add the listeners for future messages
-            self.listeners[full_name].extend(listeners)
-
-            # Send any pending messages to the listeners
-            topic_queue = self.topics[full_name]
-            while not topic_queue.empty():
-                message = topic_queue.get()
-                if message is not topic_queue.done_sentinel:
-                    for listener in self.listeners[full_name]:
-                        listener(cast(PubSubMessage, message))
-
-    def send(self, prefix: str, topic: str, value: Any) -> None:
-        full_name = self.full_name(prefix, topic)
-        message = PubSubMessage(
-            value=value,
-            topic=topic,
-            correlation_id=str(prefix),
-            published_at=datetime.now().isoformat(),
-        )
-
-        # Add the message to the log
-        self.logs[str(prefix)].put(message)
-        with self.lock:
-            listeners = self.listeners[full_name]
-            if listeners:
-                # Send the message to listeners if any are connected
-                for listener in listeners:
-                    listener(message)
-            else:
-                # Otherwise add the message to the topic queue for later
-                self.topics[full_name].put(message)
-
-    def disconnect(self, name: str) -> None:
-        with self.lock:
-            if name in self.logs:
-                self.logs[name].close()
-                if self.clear_on_disconnect:
-                    del self.logs[name]
-
-            to_delete = []
-            for topic, q in self.topics.items():
-                if topic.startswith(name):
-                    q.close()
-                    if self.clear_on_disconnect:
-                        to_delete.append(topic)
-            # can't delete while iterating
-            for topic in to_delete:
-                del self.topics[topic]
-
-            to_delete = []
-            for topic in self.listeners:
-                if topic.startswith(name):
-                    to_delete.append(topic)
-            # can't delete while iterating
-            for topic in to_delete:
-                del self.listeners[topic]

permchain/constants.py

@@ -1,2 +0,0 @@
-CONFIG_GET_KEY = "pubsub_get"
-CONFIG_SEND_KEY = "pubsub_send"

permchain/pregel/__init__.py

@@ -0,0 +1,455 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+from collections import defaultdict, deque
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    cast,
+    overload,
+)
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
+from langchain.globals import get_debug
+from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator
+from langchain.schema.runnable import (
+    Runnable,
+    RunnablePassthrough,
+    RunnableSerializable,
+)
+from langchain.schema.runnable.base import RunnableLike, coerce_to_runnable
+from langchain.schema.runnable.config import (
+    RunnableConfig,
+    get_executor_for_config,
+    patch_config,
+)
+
+from permchain.channels.base import (
+    AsyncChannelsManager,
+    BaseChannel,
+    ChannelsManager,
+    EmptyChannelError,
+)
+from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
+from permchain.pregel.debug import print_checkpoint, print_step_start
+from permchain.pregel.io import map_input, map_output
+from permchain.pregel.log import logger
+from permchain.pregel.read import PregelBatch, PregelInvoke
+from permchain.pregel.validate import validate_chains_channels
+from permchain.pregel.write import PregelSink
+
+
+class Channel:
+    @overload
+    @classmethod
+    def subscribe_to(cls, channels: str, key: Optional[str] = None) -> PregelInvoke:
+        ...
+
+    @overload
+    @classmethod
+    def subscribe_to(cls, channels: Sequence[str], key: None = None) -> PregelInvoke:
+        ...
+
+    @classmethod
+    def subscribe_to(
+        cls, channels: str | Sequence[str], key: Optional[str] = None
+    ) -> PregelInvoke:
+        """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(
+            channels=cast(
+                Mapping[None, str] | Mapping[str, str],
+                {key: channels}
+                if isinstance(channels, str)
+                else {chan: chan for chan in channels},
+            )
+        )
+
+    @classmethod
+    def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> PregelBatch:
+        """Runs process.batch() with the content of inbox each time it is updated."""
+        return PregelBatch(channel=inbox, key=key)
+
+    @classmethod
+    def write_to(
+        cls,
+        *channels: str,
+        **kwargs: RunnableLike,
+    ) -> PregelSink:
+        """Writes to channels the result of the lambda, or None to skip writing."""
+        return PregelSink(
+            channels=(
+                [(c, RunnablePassthrough()) for c in channels]
+                + [(k, coerce_to_runnable(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, PregelInvoke | PregelBatch]
+
+    output: str | Sequence[str]
+
+    input: str | Sequence[str]
+
+    step_timeout: Optional[float] = None
+
+    debug: bool = Field(default_factory=get_debug)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @root_validator(skip_on_failure=True)
+    def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
+        validate_chains_channels(
+            values["chains"], values["channels"], values["input"], values["output"]
+        )
+        return values
+
+    @property
+    def InputType(self) -> Any:
+        if isinstance(self.input, str):
+            return self.channels[self.input].UpdateType
+
+    def get_input_schema(
+        self, config: Optional[RunnableConfig] = None
+    ) -> Type[BaseModel]:
+        if isinstance(self.input, str):
+            return super().get_input_schema(config)
+        else:
+            return create_model(  # type: ignore[call-overload]
+                "PregelInput",
+                **{
+                    k: (self.channels[k].UpdateType, None)
+                    for k in self.input or self.channels.keys()
+                },
+            )
+
+    @property
+    def OutputType(self) -> Any:
+        if isinstance(self.output, str):
+            return self.channels[self.output].ValueType
+
+    def get_output_schema(
+        self, config: Optional[RunnableConfig] = None
+    ) -> Type[BaseModel]:
+        if isinstance(self.output, str):
+            return super().get_output_schema(config)
+        else:
+            return create_model(  # type: ignore[call-overload]
+                "PregelOutput",
+                **{k: (self.channels[k].ValueType, None) for k in self.output},
+            )
+
+    def _transform(
+        self,
+        input: Iterator[dict[str, Any] | Any],
+        run_manager: CallbackManagerForChainRun,
+        config: RunnableConfig,
+    ) -> Iterator[dict[str, Any] | Any]:
+        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:
+            next_tasks = _apply_writes_and_prepare_next_tasks(
+                processes,
+                channels,
+                deque(w for c in input for w in map_input(self.input, c)),
+            )
+
+            def read(chan: str) -> Any:
+                try:
+                    return channels[chan].get()
+                except EmptyChannelError:
+                    return None
+
+            # Similarly to Bulk Synchronous Parallel / Pregel model
+            # computation proceeds in steps, while there are channel updates
+            # channel updates from step N are only visible in step N+1
+            # channels are guaranteed to be immutable for the duration of the step,
+            # with channel updates applied only at the transition between steps
+            for step in range(config["recursion_limit"]):
+                if self.debug:
+                    print_step_start(step, next_tasks)
+
+                # collect all writes to channels, without applying them yet
+                pending_writes = deque[tuple[str, 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,
+                            patch_config(
+                                config,
+                                callbacks=run_manager.get_child(f"pregel:step:{step}"),
+                                configurable={
+                                    # deque.extend is thread-safe
+                                    CONFIG_KEY_SEND: pending_writes.extend,
+                                    CONFIG_KEY_READ: read,
+                                },
+                            ),
+                        )
+                        for proc, input, _ in next_tasks
+                    ),
+                    return_when=concurrent.futures.FIRST_EXCEPTION,
+                    timeout=self.step_timeout,
+                )
+
+                # interrupt on failure or timeout
+                _interrupt_or_proceed(done, inflight, step)
+
+                # apply writes to channels, decide on next step
+                next_tasks = _apply_writes_and_prepare_next_tasks(
+                    processes, channels, pending_writes
+                )
+
+                if self.debug:
+                    print_checkpoint(step, channels)
+
+                # if any write to output channels in this step, yield current value
+                for output in map_output(self.output, pending_writes, channels):
+                    yield output
+
+                # TODO this is where we'd save checkpoint
+
+                # if no more tasks, we're done
+                if not next_tasks:
+                    break
+
+    async def _atransform(
+        self,
+        input: AsyncIterator[dict[str, Any] | Any],
+        run_manager: AsyncCallbackManagerForChainRun,
+        config: RunnableConfig,
+    ) -> AsyncIterator[dict[str, Any] | Any]:
+        processes = {**self.chains}
+        # TODO this is where we'd restore from checkpoint
+        async with AsyncChannelsManager(self.channels) as channels:
+            next_tasks = _apply_writes_and_prepare_next_tasks(
+                processes,
+                channels,
+                deque([w async for c in input for w in map_input(self.input, c)]),
+            )
+
+            def read(chan: str) -> Any:
+                try:
+                    return channels[chan].get()
+                except EmptyChannelError:
+                    return None
+
+            # Similarly to Bulk Synchronous Parallel / Pregel model
+            # computation proceeds in steps, while there are channel updates
+            # channel updates from step N are only visible in step N+1,
+            # channels are guaranteed to be immutable for the duration of the step,
+            # channel updates being applied only at the transition between steps
+            for step in range(config["recursion_limit"]):
+                if self.debug:
+                    print_step_start(step, next_tasks)
+
+                # collect all writes to channels, without applying them yet
+                pending_writes = deque[tuple[str, Any]]()
+
+                # execute tasks, and wait for one to fail or all to finish.
+                # each task is independent from all other concurrent tasks
+                done, inflight = await asyncio.wait(
+                    [
+                        asyncio.create_task(
+                            proc.ainvoke(
+                                input,
+                                patch_config(
+                                    config,
+                                    callbacks=run_manager.get_child(
+                                        f"pregel:step:{step}"
+                                    ),
+                                    configurable={
+                                        # deque.extend is thread-safe
+                                        CONFIG_KEY_SEND: pending_writes.extend,
+                                        CONFIG_KEY_READ: read,
+                                    },
+                                ),
+                            )
+                        )
+                        for proc, input, _ in next_tasks
+                    ],
+                    return_when=asyncio.FIRST_EXCEPTION,
+                    timeout=self.step_timeout,
+                )
+
+                # interrupt on failure or timeout
+                _interrupt_or_proceed(done, inflight, step)
+
+                # apply writes to channels, decide on next step
+                next_tasks = _apply_writes_and_prepare_next_tasks(
+                    processes, channels, pending_writes
+                )
+
+                if self.debug:
+                    print_checkpoint(step, channels)
+
+                # if any write to output channels in this step, yield current value
+                for output in map_output(self.output, pending_writes, channels):
+                    yield output
+
+                # TODO this is where we'd save checkpoint
+
+                # if no more tasks, we're done
+                if not next_tasks:
+                    break
+
+    def invoke(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any] | Any:
+        latest: dict[str, Any] | Any = None
+        for chunk in self.stream(input, config, **kwargs):
+            latest = chunk
+        return latest
+
+    def stream(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Iterator[dict[str, Any] | Any]:
+        return self.transform(iter([input]), config, **kwargs)
+
+    def transform(
+        self,
+        input: Iterator[dict[str, Any] | Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> Iterator[dict[str, Any] | Any]:
+        return self._transform_stream_with_config(
+            input, self._transform, config, **kwargs
+        )
+
+    async def ainvoke(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any] | Any:
+        latest: dict[str, Any] | Any = None
+        async for chunk in self.astream(input, config, **kwargs):
+            latest = chunk
+        return latest
+
+    async def astream(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[dict[str, Any] | Any]:
+        async def input_stream() -> AsyncIterator[dict[str, Any] | Any]:
+            yield input
+
+        async for chunk in self.atransform(input_stream(), config, **kwargs):
+            yield chunk
+
+    async def atransform(
+        self,
+        input: AsyncIterator[dict[str, Any] | Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
+    ) -> AsyncIterator[dict[str, Any] | Any]:
+        async for chunk in self._atransform_stream_with_config(
+            input, self._atransform, config, **kwargs
+        ):
+            yield chunk
+
+
+def _interrupt_or_proceed(
+    done: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
+    inflight: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]],
+    step: int,
+) -> None:
+    while done:
+        # if any task failed
+        if exc := done.pop().exception():
+            # cancel all pending tasks
+            while inflight:
+                inflight.pop().cancel()
+            # raise the exception
+            raise exc
+            # TODO this is where retry of an entire step would happen
+
+    if inflight:
+        # if we got here means we timed out
+        while inflight:
+            # cancel all pending tasks
+            inflight.pop().cancel()
+        # raise timeout error
+        raise TimeoutError(f"Timed out at step {step}")
+
+
+def _apply_writes_and_prepare_next_tasks(
+    processes: Mapping[str, PregelInvoke | PregelBatch],
+    channels: Mapping[str, BaseChannel],
+    pending_writes: Sequence[tuple[str, Any]],
+) -> list[tuple[Runnable, Any, str]]:
+    pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    # Group writes by channel
+    for chan, val in pending_writes:
+        pending_writes_by_channel[chan].append(val)
+
+    updated_channels: set[str] = set()
+    # Apply writes to channels
+    for chan, vals in pending_writes_by_channel.items():
+        if chan in channels:
+            channels[chan].update(vals)
+            updated_channels.add(chan)
+        else:
+            logger.warning(f"Skipping write for channel {chan} which has no readers")
+
+    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 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
+                try:
+                    val = {k: channels[chan].get() for k, chan in proc.channels.items()}
+                except EmptyChannelError:
+                    continue
+
+                # Processes that subscribe to a single keyless channel get
+                # the value directly, instead of a dict
+                if list(proc.channels.keys()) == [None]:
+                    tasks.append((proc, val[None], name))
+                else:
+                    tasks.append((proc, val, name))
+        elif isinstance(proc, PregelBatch):
+            # If the channel read by this process was updated
+            if proc.channel in updated_channels:
+                # Here we don't catch EmptyChannelError because the channel
+                # must be intialized if the previous `if` condition is true
+                val = channels[proc.channel].get()
+                if proc.key is not None:
+                    val = [{proc.key: v} for v in val]
+
+                tasks.append((proc, val, name))
+
+    return tasks

permchain/pregel/constants.py

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

permchain/pregel/debug.py

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

permchain/pregel/io.py

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

permchain/pregel/log.py

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

permchain/pregel/read.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+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.utils import ConfigurableFieldSpec
+
+from permchain.channels.base import BaseChannel
+from permchain.pregel.constants import CONFIG_KEY_READ
+
+
+class PregelRead(RunnableLambda):
+    channel: str
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id=CONFIG_KEY_READ,
+                name=CONFIG_KEY_READ,
+                description=None,
+                default=None,
+                annotation=Callable[[BaseChannel], Any],
+            ),
+        ]
+
+    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]
+        self.channel = channel
+
+    def _read(self, _: Any, config: RunnableConfig) -> Any:
+        try:
+            read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
+        except KeyError:
+            raise RuntimeError(
+                f"Runnable {self} is not configured with a read function"
+                "Make sure to call in the context of a Pregel process"
+            )
+        return read(self.channel)
+
+    async def _aread(self, _: Any, config: RunnableConfig) -> Any:
+        try:
+            read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ]
+        except KeyError:
+            raise RuntimeError(
+                f"Runnable {self} is not configured with a read function"
+                "Make sure to call in the context of a Pregel process"
+            )
+        return read(self.channel)
+
+
+class PregelInvoke(RunnableBinding):
+    channels: Mapping[None, str] | Mapping[str, str]
+
+    bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
+
+    kwargs: Mapping[str, Any] = Field(default_factory=dict)
+
+    def __init__(
+        self,
+        channels: Mapping[None, str] | Mapping[str, str],
+        *,
+        bound: Optional[Runnable[Any, Any]] = None,
+        kwargs: Optional[Mapping[str, Any]] = None,
+        config: Optional[RunnableConfig] = None,
+        **other_kwargs: Any,
+    ) -> None:
+        super().__init__(
+            channels=channels,
+            bound=bound or RunnablePassthrough(),
+            kwargs=kwargs or {},
+            config=config,
+            **other_kwargs,
+        )
+
+    def join(self, channels: Sequence[str]) -> PregelInvoke:
+        joiner = RunnablePassthrough.assign(
+            **{chan: PregelRead(chan) for chan in channels}
+        )
+        if isinstance(self.bound, RunnablePassthrough):
+            return PregelInvoke(channels=self.channels, bound=joiner)
+        else:
+            return PregelInvoke(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))
+        else:
+            # delegate to __or__ in self.bound
+            return PregelInvoke(channels=self.channels, bound=self.bound | other)
+
+    def __ror__(
+        self,
+        other: Runnable[Other, Any]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
+    ) -> Runnable:
+        raise NotImplementedError()
+
+
+class PregelBatch(RunnableEach):
+    channel: str
+
+    key: Optional[str]
+
+    bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough)
+
+    def join(self, channels: Sequence[str]) -> PregelBatch:
+        if self.key is None:
+            raise ValueError(
+                "Cannot join() additional channels without a key."
+                " Pass a key arg to Channel.subscribe_to_each()."
+            )
+
+        joiner = RunnablePassthrough.assign(
+            **{chan: PregelRead(chan) for chan in channels}
+        )
+        if isinstance(self.bound, RunnablePassthrough):
+            return PregelBatch(channel=self.channel, key=self.key, bound=joiner)
+        else:
+            return PregelBatch(
+                channel=self.channel, key=self.key, bound=self.bound | joiner
+            )
+
+    def __or__(  # type: ignore[override]
+        self,
+        other: Runnable[Any, Other]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
+    ) -> PregelBatch:
+        if isinstance(self.bound, RunnablePassthrough):
+            return PregelBatch(
+                channel=self.channel, key=self.key, bound=coerce_to_runnable(other)
+            )
+        else:
+            # delegate to __or__ in self.bound
+            return PregelBatch(
+                channel=self.channel, key=self.key, bound=self.bound | other
+            )
+
+    def __ror__(
+        self,
+        other: Runnable[Other, Any]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
+    ) -> Runnable:
+        raise NotImplementedError()

permchain/pregel/validate.py

@@ -0,0 +1,43 @@
+from typing import Mapping, Sequence
+
+from permchain.channels.base import BaseChannel
+from permchain.pregel.read import PregelBatch, PregelInvoke
+
+
+def validate_chains_channels(
+    chains: Mapping[str, PregelInvoke | PregelBatch],
+    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):
+            subscribed_channels.update(chain.channels.values())
+        elif isinstance(chain, PregelBatch):
+            subscribed_channels.add(chain.channel)
+        else:
+            raise TypeError(
+                f"Invalid chain type {type(chain)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()"
+            )
+
+    for chan in subscribed_channels:
+        if chan not in channels:
+            raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
+
+    if isinstance(input, str):
+        if input not in subscribed_channels:
+            raise ValueError(f"Input channel {input} is not subscribed to by any chain")
+    else:
+        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"
+            )
+
+    if isinstance(output, str):
+        if output not in channels:
+            raise ValueError(f"Output channel {output} is not initialized")
+    else:
+        for chan in output:
+            if chan not in channels:
+                raise ValueError(f"Output channel {chan} is not initialized")

permchain/pregel/write.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from langchain.schema.runnable import (
+    Runnable,
+    RunnableConfig,
+    RunnableLambda,
+)
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+from permchain.pregel.constants import CONFIG_KEY_SEND
+
+TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
+
+
+# TODO switch to RunnablePassthrough after updating langchain
+class PregelSink(RunnableLambda):
+    channels: Sequence[tuple[str, Runnable]]
+    """
+    Mapping of write channels to Runnables that return the value to be written,
+    or None to skip writing.
+    """
+
+    def __init__(
+        self,
+        *,
+        channels: Sequence[tuple[str, Runnable]],
+    ):
+        super().__init__(func=self._write, afunc=self._awrite)  # type: ignore[arg-type]
+        self.channels = channels
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id=CONFIG_KEY_SEND,
+                name=CONFIG_KEY_SEND,
+                description=None,
+                default=None,
+                annotation=TYPE_SEND,
+            ),
+        ]
+
+    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
+
+    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])
+
+        return input

permchain/pubsub.py

@@ -1,250 +0,0 @@
-from __future__ import annotations
-
-import threading
-from abc import ABC
-from collections import defaultdict
-from concurrent.futures import CancelledError, Future
-from functools import partial
-from typing import Any, Iterator, List, Optional, Sequence, Set, TypeVar
-
-from langchain.callbacks.manager import CallbackManagerForChainRun
-from langchain.schema.runnable import Runnable, RunnableConfig, patch_config
-from langchain.schema.runnable.config import get_executor_for_config
-
-from permchain.connection import PubSubConnection, PubSubMessage
-from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY
-from permchain.topic import (
-    INPUT_TOPIC,
-    OUTPUT_TOPIC,
-    RunnableReducer,
-    RunnableSubscriber,
-)
-
-T = TypeVar("T")
-T_in = TypeVar("T_in")
-T_out = TypeVar("T_out")
-
-Process = RunnableSubscriber[T_in] | RunnableReducer[T_in]
-
-
-class PubSub(Runnable[Any, Any], ABC):
-    processes: Sequence[Process]
-
-    connection: PubSubConnection
-
-    def __init__(
-        self,
-        *procs: Process | Sequence[Process],
-        processes: Sequence[Process] = (),
-        connection: PubSubConnection,
-    ) -> None:
-        super().__init__()
-
-        self.lock = threading.Lock()
-        self.inflight_namespaces = set()
-
-        self.connection = connection
-        self.processes = list(processes)
-        for proc in procs:
-            if isinstance(proc, Sequence):
-                self.processes.extend(proc)
-            else:
-                self.processes.append(proc)
-
-    def with_retry(self, **kwargs: Any) -> Runnable[Any, Any]:
-        return self.__class__(
-            processes=[p.with_retry(**kwargs) for p in self.processes],
-            connection=self.connection,
-        )
-
-    def _transform(
-        self,
-        input: Iterator[Any],
-        run_manager: CallbackManagerForChainRun,
-        config: RunnableConfig,
-    ) -> Iterator[Any]:
-        # Split processes into subscribers and reducers, and group by topic
-        subscribers: defaultdict[str, list[RunnableSubscriber[Any]]] = defaultdict(list)
-        reducers: defaultdict[str, list[RunnableReducer[Any]]] = defaultdict(list)
-        for process in self.processes:
-            if isinstance(process, RunnableReducer):
-                reducers[process.topic.name].append(process)
-            elif isinstance(process, RunnableSubscriber):
-                subscribers[process.topic.name].append(process)
-            else:
-                raise ValueError(f"Unknown process type: {process}")
-
-        # Consume input iterator into a single value
-        input_value = None
-        for chunk in input:
-            if input_value is None:
-                input_value = chunk
-            else:
-                input_value += chunk
-
-        with get_executor_for_config(config) as executor:
-            # Namespace topics for each run, default to run_id, ie. isolated
-            topic_prefix = str(config.get("correlation_id") or run_manager.run_id)
-
-            # Check if this correlation_id is currently inflight. If so, raise an error,
-            # as that would make the output iterator produce incorrect results.
-            with self.lock:
-                if topic_prefix in self.inflight_namespaces:
-                    raise RuntimeError(
-                        f"Cannot run {self} in namespace {topic_prefix} "
-                        "because it is currently in use"
-                    )
-                self.inflight_namespaces.add(topic_prefix)
-
-            # Track inflight futures
-            inflight: Set[Future] = set()
-            # Track exceptions
-            exceptions: List[BaseException] = []
-
-            def on_idle() -> None:
-                """Called when all subscribed topics are empty.
-                It first runs any topic reducers. Then, if all subscribed topics
-                still empty, it closes the computation.
-                """
-                if reducers:
-                    for topic_name, processes in reducers.items():
-                        # Collect all pending messages for each topic
-                        messages = list(
-                            self.connection.iterate(
-                                topic_prefix, topic_name, wait=False
-                            )
-                        )
-                        # Run each reducer once with the collected messages
-                        if messages:
-                            for process in processes:
-                                run_once(process, messages)
-
-                if not inflight:
-                    self.connection.disconnect(topic_prefix)
-
-            def check_if_idle(fut: Future) -> None:
-                """Cleanup after a process runs."""
-                inflight.discard(fut)
-
-                try:
-                    exc = fut.exception()
-                except CancelledError:
-                    exc = None
-                except Exception as e:
-                    exc = e
-                if exc is not None:
-                    exceptions.append(exc)
-
-                # Close output iterator if
-                # - all processes are done, or
-                # - an exception occurred
-                if not inflight or exc is not None:
-                    on_idle()
-
-            def run_once(
-                process: RunnableSubscriber[Any] | RunnableReducer[Any],
-                messages: PubSubMessage | list[PubSubMessage],
-            ) -> None:
-                """Run a process once."""
-                value = (
-                    [m["value"] for m in messages]
-                    if isinstance(messages, list)
-                    else messages["value"]
-                )
-
-                def get(topic_name: str) -> Any:
-                    if topic_name == INPUT_TOPIC:
-                        return input_value
-                    elif topic_name == process.topic.name:
-                        return value
-                    else:
-                        raise ValueError(
-                            f"Cannot get value for {topic_name} in this context"
-                        )
-
-                # Run process once in executor
-                try:
-                    fut = executor.submit(
-                        process.invoke,
-                        value,
-                        config={
-                            **patch_config(
-                                config,
-                                callbacks=run_manager.get_child(),
-                                run_name=f"Topic: {process.topic.name}",
-                            ),
-                            CONFIG_SEND_KEY: partial(
-                                self.connection.send, topic_prefix
-                            ),
-                            CONFIG_GET_KEY: get,
-                            # TODO below doesn't work for batch calls nested inside
-                            # another pubsub, eg. test_invoke_join_then_call_other_pubsub
-                            # as all messages in each batch would share same correlation_id
-                            # "correlation_id": self.connection.full_name(
-                            #     topic_prefix,
-                            #     process.topic.name,
-                            #     str(self.processes.index(process)),
-                            # ),
-                        },
-                    )
-
-                    # Add callback to cleanup
-                    inflight.add(fut)
-                    fut.add_done_callback(check_if_idle)
-                except RuntimeError:
-                    # If executor is now closed, just ignore this process
-                    # This could happen eg. if an OUT message was published durin
-                    # execution of run_once
-                    pass
-
-            # Listen on all subscribed topics
-            for topic_name, processes in subscribers.items():
-                self.connection.listen(
-                    topic_prefix,
-                    topic_name,
-                    [partial(run_once, process) for process in processes],
-                )
-
-            # Send input to input processes
-            self.connection.send(topic_prefix, INPUT_TOPIC, input_value)
-
-            try:
-                if inflight:
-                    # Yield output until all processes are done
-                    # This blocks the current thread, all other work needs to go
-                    # through the executor
-                    for chunk in self.connection.observe(topic_prefix):
-                        yield chunk
-                        if chunk["topic"] == OUTPUT_TOPIC:
-                            # All expected output has been received, close
-                            self.connection.disconnect(topic_prefix)
-                            break
-                else:
-                    on_idle()
-            finally:
-                # Cancel all inflight futures
-                while inflight:
-                    inflight.pop().cancel()
-
-                # Remove namespace from inflight set
-                with self.lock:
-                    self.inflight_namespaces.remove(topic_prefix)
-
-                # Raise exceptions if any
-                if exceptions:
-                    raise exceptions[0]
-
-    def stream(
-        self,
-        input: Any,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
-    ) -> Iterator[PubSubMessage]:
-        yield from self._transform_stream_with_config(
-            iter([input]), self._transform, config, **kwargs
-        )
-
-    def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
-        for chunk in self.stream(input, config):
-            if chunk["topic"] == OUTPUT_TOPIC:
-                return chunk["value"]

permchain/topic.py

@@ -1,168 +0,0 @@
-from __future__ import annotations
-
-from typing import (
-    Any,
-    Callable,
-    Generic,
-    Mapping,
-    Optional,
-    Sequence,
-    TypeVar,
-)
-
-from langchain.load.serializable import Serializable
-from langchain.pydantic_v1 import Field
-from langchain.schema.runnable import (
-    Runnable,
-    RunnableBinding,
-    RunnableConfig,
-    RunnablePassthrough,
-    RunnableSequence,
-)
-from langchain.schema.runnable.base import Other, coerce_to_runnable
-
-from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY
-
-T = TypeVar("T")
-
-
-INPUT_TOPIC = "__in__"
-OUTPUT_TOPIC = "__out__"
-
-
-class Topic(Serializable, Generic[T]):
-    name: str
-
-    def __init__(self, name: str):
-        super().__init__(name=name)
-
-    def subscribe(self) -> RunnableSubscriber[T]:
-        if self.name == OUTPUT_TOPIC:
-            raise ValueError("Cannot subscribe to output topic")
-
-        return RunnableSubscriber(topic=self)
-
-    def join(self) -> RunnableReducer[T]:
-        if self.name == OUTPUT_TOPIC:
-            raise ValueError("Cannot join on output topic")
-
-        return RunnableReducer(topic=self)
-
-    def current(self) -> RunnableCurrentValue[T]:
-        if self.name == OUTPUT_TOPIC:
-            raise ValueError("Cannot subscribe to output topic")
-
-        return RunnableCurrentValue(topic=self)
-
-    def publish(self) -> RunnablePublisher[T]:
-        if self.name == INPUT_TOPIC:
-            raise ValueError("Cannot publish to input topic")
-
-        return RunnablePublisher(topic=self)
-
-    def publish_each(self) -> RunnablePublisherEach[T]:
-        if self.name == INPUT_TOPIC:
-            raise ValueError("Cannot publish to input topic")
-
-        return RunnablePublisherEach(topic=self)
-
-    @classmethod
-    @property
-    def IN(cls) -> Topic:
-        return cls(INPUT_TOPIC)
-
-    @classmethod
-    @property
-    def OUT(cls) -> Topic:
-        return cls(OUTPUT_TOPIC)
-
-
-class RunnableConfigForPubSub(RunnableConfig):
-    send: Callable[[str, Any], None]
-    get: Callable[[str], Any]
-
-
-class RunnableSubscriber(RunnableBinding[T, Any]):
-    topic: Topic[T]
-
-    bound: Runnable[T, Any] = Field(default_factory=RunnablePassthrough)
-
-    kwargs: Mapping[str, Any] = Field(default_factory=dict)
-
-    def __or__(
-        self,
-        other: Runnable[Any, Other]
-        | Callable[[Any], Other]
-        | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
-    ) -> RunnableSubscriber[T, Other]:
-        if isinstance(self.bound, RunnablePassthrough):
-            return RunnableSubscriber(topic=self.topic, bound=coerce_to_runnable(other))
-        else:
-            return RunnableSubscriber(topic=self.topic, bound=self.bound | other)
-
-    def __ror__(
-        self,
-        other: Runnable[Other, Any]
-        | Callable[[Any], Other]
-        | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
-    ) -> RunnableSubscriber[Other, Any]:
-        raise NotImplementedError()
-
-
-class RunnableReducer(RunnableBinding[list[T], Any]):
-    topic: Topic[T]
-
-    bound: Runnable[list[T], Any] = Field(default_factory=RunnablePassthrough)
-
-    kwargs: Mapping[str, Any] = Field(default_factory=dict)
-
-    def __or__(
-        self,
-        other: Runnable[Any, Other]
-        | Callable[[Any], Other]
-        | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
-    ) -> RunnableSequence[list[T], Other]:
-        if isinstance(self.bound, RunnablePassthrough):
-            return RunnableReducer(topic=self.topic, bound=coerce_to_runnable(other))
-        else:
-            return RunnableReducer(topic=self.topic, bound=self.bound | other)
-
-    def __ror__(
-        self,
-        other: Runnable[Other, Any]
-        | Callable[[Any], Other]
-        | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]],
-    ) -> RunnableSequence[Other, Any]:
-        raise NotImplementedError()
-
-
-class RunnablePublisher(Serializable, Runnable[T, T]):
-    topic: Topic[T]
-
-    def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
-        send = config.get(CONFIG_SEND_KEY, None)
-        if send is not None:
-            send(self.topic.name, input)
-        return input
-
-
-class RunnablePublisherEach(RunnablePublisher[Sequence[T]]):
-    topic: Topic[T]
-
-    def invoke(
-        self, input: Sequence[T], config: Optional[RunnableConfigForPubSub] = None
-    ) -> Sequence[T]:
-        for item in input:
-            super().invoke(item, config)
-        return input
-
-
-class RunnableCurrentValue(Serializable, Runnable[Any, T]):
-    topic: Topic[T]
-
-    def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T:
-        get: Callable[[str], None] = config.get(CONFIG_GET_KEY, None)
-        if get is not None:
-            return get(self.topic.name)
-        else:
-            raise ValueError("Cannot get value in this context")

pyproject.toml

@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/permchain"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-langchain = ">=0.0.300"
+langchain = ">=0.0.313"
 
 
 [tool.poetry.group.test.dependencies]
@@ -29,7 +29,7 @@ ruff = "^0.0.249"
 black = {extras = ["jupyter"], version = "^23.7.0"}
 
 [tool.poetry.group.typing.dependencies]
-mypy = "^0.991"
+mypy = "^1.6.0"
 
 [tool.poetry.group.dev]
 optional = true
@@ -57,6 +57,7 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.pytest.ini_options]
+asyncio_mode = "auto"
 # --strict-markers will raise errors on unknown marks.
 # https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
 #
@@ -66,6 +67,6 @@ build-backend = "poetry.core.masonry.api"
 #
 # https://github.com/tophat/syrupy
 # --snapshot-warn-unused    Prints a warning on unused snapshots rather than fail the test suite.
-addopts = "--full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
+addopts = "-x --full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
 # Registering custom markers.
 # https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers

tests/test_channels.py

@@ -0,0 +1,204 @@
+import operator
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncGenerator, FrozenSet, 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
+
+
+def test_last_value() -> None:
+    with LastValue(int).empty() as channel:
+        assert channel.ValueType is int
+        assert channel.UpdateType is int
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5, 6])
+
+        channel.update([3])
+        assert channel.get() == 3
+        channel.update([4])
+        assert channel.get() == 4
+
+
+async def test_last_value_async() -> None:
+    async with LastValue(int).aempty() as channel:
+        assert channel.ValueType is int
+        assert channel.UpdateType is int
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5, 6])
+
+        channel.update([3])
+        assert channel.get() == 3
+        channel.update([4])
+        assert channel.get() == 4
+
+
+def test_inbox() -> None:
+    with Inbox(str).empty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, Sequence[str]]
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+
+        channel.update(["a", "b"])
+        assert channel.get() == ("a", "b")
+        channel.update([["c"], "d"])
+        assert channel.get() == ("c", "d")
+
+
+async def test_inbox_async() -> None:
+    async with Inbox(str).aempty() as channel:
+        assert channel.ValueType is Sequence[str]
+        assert channel.UpdateType is Union[str, Sequence[str]]
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+
+        channel.update(["a", "b"])
+        assert channel.get() == ("a", "b")
+        channel.update(["c"])
+        channel.update([["c"], "d"])
+        assert channel.get() == ("c", "d")
+
+
+def test_set() -> None:
+    with UniqueArchive(str).empty() as channel:
+        assert channel.ValueType is FrozenSet[str]
+        assert channel.UpdateType is 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"))
+
+
+async def test_set_async() -> None:
+    async with UniqueArchive(str).aempty() as channel:
+        assert channel.ValueType is FrozenSet[str]
+        assert channel.UpdateType is 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"))
+
+
+def test_binop() -> None:
+    with BinaryOperatorAggregate(int, operator.add).empty() as channel:
+        assert channel.ValueType is int
+        assert channel.UpdateType is int
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+
+        channel.update([1, 2, 3])
+        assert channel.get() == 6
+        channel.update([4])
+        assert channel.get() == 10
+
+
+async def test_binop_async() -> None:
+    async with BinaryOperatorAggregate(int, operator.add).aempty() as channel:
+        assert channel.ValueType is int
+        assert channel.UpdateType is int
+
+        with pytest.raises(EmptyChannelError):
+            channel.get()
+
+        channel.update([1, 2, 3])
+        assert channel.get() == 6
+        channel.update([4])
+        assert channel.get() == 10
+
+
+def test_ctx_manager(mocker: MockerFixture) -> None:
+    setup = mocker.Mock()
+    cleanup = mocker.Mock()
+
+    @contextmanager
+    def an_int() -> Generator[int, None, None]:
+        setup()
+        try:
+            yield 5
+        finally:
+            cleanup()
+
+    with Context(an_int, None, int).empty() as channel:
+        assert setup.call_count == 1
+        assert cleanup.call_count == 0
+
+        assert channel.ValueType is int
+        with pytest.raises(InvalidUpdateError):
+            assert channel.UpdateType is None
+
+        assert channel.get() == 5
+
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5])  # type: ignore
+
+    assert setup.call_count == 1
+    assert cleanup.call_count == 1
+
+
+def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
+    with Context(httpx.Client).empty() as channel:
+        assert channel.ValueType is httpx.Client
+        with pytest.raises(InvalidUpdateError):
+            assert channel.UpdateType is None
+
+        assert isinstance(channel.get(), httpx.Client)
+
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5])  # type: ignore
+
+
+async def test_ctx_manager_async(mocker: MockerFixture) -> None:
+    setup = mocker.Mock()
+    cleanup = mocker.Mock()
+
+    @contextmanager
+    def an_int_sync() -> Generator[int, None, None]:
+        try:
+            yield 5
+        finally:
+            pass
+
+    @asynccontextmanager
+    async def an_int() -> AsyncGenerator[int, None]:
+        setup()
+        try:
+            yield 5
+        finally:
+            cleanup()
+
+    async with Context(an_int_sync, an_int, int).aempty() as channel:
+        assert setup.call_count == 1
+        assert cleanup.call_count == 0
+
+        assert channel.ValueType is int
+        with pytest.raises(InvalidUpdateError):
+            assert channel.UpdateType is None
+
+        assert channel.get() == 5
+
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5])  # type: ignore
+
+    assert setup.call_count == 1
+    assert cleanup.call_count == 1

tests/test_invoke.py

@@ -1,498 +0,0 @@
-from typing import Iterator
-from uuid import uuid4
-
-import pytest
-from pytest_mock import MockerFixture
-
-from permchain.connection import PubSubMessage
-from permchain.connection_inmemory import InMemoryPubSubConnection
-from permchain.pubsub import PubSub
-from permchain.topic import RunnableSubscriber, Topic
-
-
-def clean_log(
-    logs: Iterator[PubSubMessage], correlation_id: bool | None = None
-) -> list[PubSubMessage]:
-    if correlation_id is False:
-        return [{**m, "published_at": None, "correlation_id": None} for m in logs]
-    else:
-        return [{**m, "published_at": None} for m in logs]
-
-
-def test_invoke_single_process_in_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(chain, connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-    # Then invoke pubsub
-    assert pubsub.invoke(2) == 3
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_two_processes_in_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
-    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(chain_one, chain_two, connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-    # Then invoke pubsub
-    assert pubsub.invoke(2) == 4
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-@pytest.mark.skip("TODO")
-def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
-    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection(clear_on_disconnect=False)
-    pubsub_one = PubSub(chain_one, connection=conn)
-    pubsub_two = PubSub(chain_two, connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-    # Then invoke both pubsubs, as a group
-    # The second picks up where the first left off
-    correlation_id = uuid4()
-
-    # invoke() step 1
-    assert clean_log(pubsub_one.stream(2, {"correlation_id": correlation_id})) == [
-        {
-            "value": 2,
-            "topic": "__in__",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 3,
-            "topic": "one",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-    ]
-
-    # IN, one
-    assert len(conn.topics) == 2
-    topic_one_full_name = conn.full_name(correlation_id, topic_one.name)
-    # the actual message publishd by chain_one, and a sentinel "end" value
-    assert conn.topics[topic_one_full_name].qsize() == 2
-
-    # invoke() step 2
-    # this picks up where the first left off, and produces same result as
-    # `test_invoke_two_processes_in_out`
-    assert clean_log(pubsub_two.stream(None, {"correlation_id": correlation_id})) == [
-        {
-            "value": None,
-            "topic": "__in__",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 4,
-            "topic": "__out__",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-    ]
-    # listeners are still cleared, even though state is preserved
-    assert conn.listeners == {}
-    # IN, OUT, one
-    assert len(conn.topics) == 3
-
-    for topic_name, queue in conn.topics.items():
-        if topic_name.endswith("IN"):
-            # Contains two sentinel "end" values, and None
-            # passed in as input to chain_two, which doesn't subscribe to it
-            assert queue.qsize() == 3
-        if topic_name.endswith("OUT"):
-            # Empty because this was consumed by invoke()
-            assert queue.qsize() == 0
-        if topic_name.endswith("one"):
-            # Contains 2 sentinel "end" values
-            assert queue.qsize() == 2
-
-
-def test_invoke_many_processes_in_out(mocker: MockerFixture):
-    test_size = 100
-
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topics: list[Topic] = [Topic("zero")]
-    chains: list[RunnableSubscriber] = [
-        Topic.IN.subscribe() | add_one | topics[0].publish()
-    ]
-    for i in range(test_size - 2):
-        topics.append(Topic(str(i)))
-        chains.append(topics[-2].subscribe() | add_one | topics[-1].publish())
-    chains.append(topics[-1].subscribe() | add_one | Topic.OUT.publish())
-
-    # Chains can be invoked directly for testing
-    for chain in chains:
-        assert chain.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=chains, connection=conn)
-
-    for _ in range(10):
-        # Using in-memory conn internals to make assertions about pubsub
-        # If we start with 0 listeners
-        assert conn.listeners == {}
-        # Then invoke pubsub
-        assert pubsub.invoke(2) == 2 + test_size
-        # After invoke returns the listeners were cleaned up
-        assert conn.listeners == {}
-
-
-def test_batch_many_processes_in_out(mocker: MockerFixture):
-    test_size = 100
-
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topics: list[Topic] = [Topic("zero")]
-    chains: list[RunnableSubscriber] = [
-        Topic.IN.subscribe() | add_one | topics[0].publish()
-    ]
-    for i in range(test_size - 2):
-        topics.append(Topic(str(i)))
-        chains.append(topics[-2].subscribe() | add_one | topics[-1].publish())
-    chains.append(topics[-1].subscribe() | add_one | Topic.OUT.publish())
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=chains, connection=conn)
-
-    for _ in range(10):
-        # Using in-memory conn internals to make assertions about pubsub
-        # If we start with 0 listeners
-        assert conn.listeners == {}
-        # Then invoke pubsub
-        assert pubsub.batch([2, 1, 3, 4, 5]) == [
-            2 + test_size,
-            1 + test_size,
-            3 + test_size,
-            4 + test_size,
-            5 + test_size,
-        ]
-        # After invoke returns the listeners were cleaned up
-        assert conn.listeners == {}
-
-
-def test_invoke_two_processes_two_in_two_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    chain_one = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
-    chain_two = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # We get only one of the two return values, as computation is closed
-    # as soon as we publish to OUT for the first time
-    assert pubsub.invoke(2) == 3
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
-    topic_one = Topic("one")
-    topic_two = Topic("two")
-    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
-    chain_two = topic_one.subscribe() | add_one | topic_two.publish()
-    chain_three = Topic.IN.subscribe() | add_one | topic_two.publish()
-    chain_four = topic_two.join() | add_10_each | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_four.invoke([2, 3]) == [12, 13]
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub((chain_one, chain_two, chain_three, chain_four), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # We get a single array result as chain_four waits for all publishers to finish
-    # before operating on all elements published to topic_two as an array
-    assert pubsub.invoke(2) == [13, 14]
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_join_then_subscribe(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
-
-    topic_one = Topic("one")
-    topic_two = Topic("two")
-
-    chain_one = Topic.IN.subscribe() | add_10_each | topic_one.publish_each()
-    chain_two = topic_one.join() | sum | topic_two.publish()
-    chain_three = topic_two.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_two.invoke([2, 3]) == 5
-    assert chain_three.invoke(5) == 6
-
-    correlation_id = uuid4()
-    conn = InMemoryPubSubConnection(clear_on_disconnect=False)
-    pubsub = PubSub((chain_one, chain_two, chain_three), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # We get a single array result as chain_four waits for all publishers to finish
-    # before operating on all elements published to topic_two as an array
-    assert clean_log(pubsub.stream([2, 3], {"correlation_id": correlation_id})) == [
-        {
-            "value": [2, 3],
-            "topic": "__in__",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 12,
-            "topic": "one",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 13,
-            "topic": "one",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 25,
-            "topic": "two",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-        {
-            "value": 26,
-            "topic": "__out__",
-            "correlation_id": str(correlation_id),
-            "published_at": None,
-        },
-    ]
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
-    conn = InMemoryPubSubConnection(clear_on_disconnect=False)
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-
-    inner_pubsub = PubSub(
-        (Topic.IN.subscribe() | add_one | Topic.OUT.publish(),), connection=conn
-    )
-
-    add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
-
-    topic_one = Topic("one")
-    topic_two = Topic("two")
-
-    chain_one = Topic.IN.subscribe() | add_10_each | topic_one.publish_each()
-    chain_two = topic_one.join() | inner_pubsub.map() | sorted | topic_two.publish()
-    chain_three = topic_two.subscribe() | sum | Topic.OUT.publish()
-
-    pubsub = PubSub((chain_one, chain_two, chain_three), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    for _ in range(10):
-        assert clean_log(pubsub.stream([2, 3]), correlation_id=False) == [
-            {
-                "value": [2, 3],
-                "topic": "__in__",
-                "correlation_id": None,
-                "published_at": None,
-            },
-            {
-                "value": 12,
-                "topic": "one",
-                "correlation_id": None,
-                "published_at": None,
-            },
-            {
-                "value": 13,
-                "topic": "one",
-                "correlation_id": None,
-                "published_at": None,
-            },
-            {
-                "value": [13, 14],
-                "topic": "two",
-                "correlation_id": None,
-                "published_at": None,
-            },
-            {
-                "value": 27,
-                "topic": "__out__",
-                "correlation_id": None,
-                "published_at": None,
-            },
-        ]
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    # Topic.publish() is passthrough so we can publish to multiple topics in sequence
-    chain_one = (
-        Topic.IN.subscribe() | add_one | Topic.OUT.publish() | topic_one.publish()
-    )
-    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # pubsub stopped executing after publishing to OUT, so only one value is returned
-    assert clean_log(pubsub.stream(2), correlation_id=False) == [
-        {
-            "value": 2,
-            "topic": "__in__",
-            "correlation_id": None,
-            "published_at": None,
-        },
-        {
-            "value": 3,
-            "topic": "__out__",
-            "correlation_id": None,
-            "published_at": None,
-        },
-    ]
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_two_processes_no_out(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
-    chain_two = topic_one.subscribe() | add_one
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # It finishes executing (once no more messages being published)
-    # but returns nothing, as nothing was published to OUT topic
-    assert pubsub.invoke(2) is None
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-def test_invoke_two_processes_no_in(mocker: MockerFixture):
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    chain_one = topic_one.subscribe() | add_one | Topic.OUT.publish()
-    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-
-    # Then invoke pubsub
-    # It returns without any output as there is nothing to run
-    assert pubsub.invoke(2) is None
-
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
-
-
-@pytest.mark.skip("TODO")
-def test_invoke_two_processes_simple_cycle(mocker: MockerFixture) -> None:
-    add_one = mocker.Mock(side_effect=lambda x: x + 1)
-    topic_one = Topic("one")
-    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
-    chain_two = topic_one.subscribe() | add_one | topic_one.publish()
-
-    # Chains can be invoked directly for testing
-    assert chain_one.invoke(2) == 3
-    assert chain_two.invoke(2) == 3
-
-    conn = InMemoryPubSubConnection()
-    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
-
-    # Using in-memory conn internals to make assertions about pubsub
-    # If we start with 0 listeners
-    assert conn.listeners == {}
-    # Then invoke pubsub
-    with pytest.raises(RecursionError):
-        pubsub.invoke(2)
-    # After invoke returns the listeners were cleaned up
-    for key in conn.listeners:
-        assert not conn.listeners[key]

tests/test_pregel.py

@@ -0,0 +1,459 @@
+import time
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Generator
+
+import pytest
+from langchain.schema.runnable import RunnablePassthrough
+from pytest_mock import MockerFixture
+
+from permchain import Channel, Pregel
+from permchain.channels.base import InvalidUpdateError
+from permchain.channels.context import Context
+from permchain.channels.inbox import Inbox
+from permchain.channels.last_value import LastValue
+
+
+def test_invoke_single_process_in_out(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,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
+    assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
+    assert app.invoke(2) == 3
+
+
+def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
+
+    app = Pregel(
+        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.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {"output": {"title": "Output", "type": "integer"}},
+    }
+    assert app.invoke(2) == {"output": 3}
+
+
+def test_invoke_single_process_in_dict_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")
+
+    app = Pregel(
+        chains={
+            "one": chain,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+        },
+        input=["input"],
+        output=["output"],
+    )
+
+    assert app.input_schema.schema() == {
+        "title": "PregelInput",
+        "type": "object",
+        "properties": {"input": {"title": "Input", "type": "integer"}},
+    }
+    assert app.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {"output": {"title": "Output", "type": "integer"}},
+    }
+    assert app.invoke({"input": 2}) == {"output": 3}
+
+
+def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    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",
+    )
+
+    assert app.invoke(2) == 4
+
+
+def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one, "chain_two": chain_two},
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+        },
+        input=["input", "inbox"],
+        output="output",
+    )
+
+    assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4]  # [12 + 1, 2 + 1 + 1]
+
+
+def test_batch_two_processes_in_out() -> None:
+    def add_one_with_delay(inp: int) -> int:
+        time.sleep(inp / 10)
+        return inp + 1
+
+    chain_one = (
+        Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
+    )
+    chain_two = (
+        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",
+    )
+
+    assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
+
+
+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")
+
+    for _ in range(10):
+        assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
+
+    with ThreadPoolExecutor() as executor:
+        assert [
+            *executor.map(app.invoke, [2] * 10, [{"recursion_limit": test_size}] * 10)
+        ] == [2 + test_size] * 10
+
+
+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")
+
+    for _ in range(10):
+        assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
+            2 + test_size,
+            1 + test_size,
+            3 + test_size,
+            4 + test_size,
+            5 + test_size,
+        ]
+
+    with ThreadPoolExecutor() as executor:
+        assert [
+            *executor.map(
+                app.batch, [[2, 1, 3, 4, 5]] * 10, [{"recursion_limit": test_size}] * 10
+            )
+        ] == [
+            [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
+        ] * 10
+
+
+def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = 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",
+    )
+
+    with pytest.raises(InvalidUpdateError):
+        # LastValue channels can only be updated once per iteration
+        app.invoke(2)
+
+
+def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = 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": Inbox(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    # An Inbox channel accumulates updates into a sequence
+    assert app.invoke(2) == (3, 3)
+
+
+def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
+
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_four = (
+        Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
+    )
+
+    app = Pregel(
+        chains={
+            "chain_one": chain_one,
+            "chain_three": chain_three,
+            "chain_four": chain_four,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    # Then invoke app
+    # We get a single array result as chain_four waits for all publishers to finish
+    # before operating on all elements published to topic_two as an array
+    for _ in range(100):
+        assert app.invoke(2) == [13, 13]
+
+    with ThreadPoolExecutor() as executor:
+        assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
+
+
+def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
+
+    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 = (
+        Channel.subscribe_to("input")
+        | add_10_each
+        | Channel.write_to("inbox_one").map()
+    )
+    chain_two = (
+        Channel.subscribe_to("inbox_one")
+        | inner_app.map()
+        | sorted
+        | Channel.write_to("outbox_one")
+    )
+    chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
+
+    app = Pregel(
+        chains={
+            "chain_one": chain_one,
+            "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",
+    )
+
+    for _ in range(10):
+        assert app.invoke([2, 3]) == 27
+
+    with ThreadPoolExecutor() as executor:
+        assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10
+
+
+def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = (
+        Channel.subscribe_to("input")
+        | add_one
+        | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
+    )
+    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",
+    )
+
+    assert [c for c in app.stream(2)] == [3, 4]
+
+
+def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    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",
+    )
+
+    # It finishes executing (once no more messages being published)
+    # but returns nothing, as nothing was published to OUT topic
+    assert app.invoke(2) is None
+
+
+def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
+    chain_two = Channel.subscribe_to("between") | add_one
+
+    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",
+        )
+
+
+def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
+    setup = mocker.Mock()
+    cleanup = mocker.Mock()
+
+    @contextmanager
+    def an_int() -> Generator[int, None, None]:
+        setup()
+        try:
+            yield 5
+        finally:
+            cleanup()
+
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one, "chain_two": chain_two},
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+            "ctx": Context(an_int, typ=int),
+        },
+        input="input",
+        output=["inbox", "output"],
+    )
+
+    assert setup.call_count == 0
+    assert cleanup.call_count == 0
+    for i, chunk in enumerate(app.stream(2)):
+        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,)}
+        elif i == 1:
+            assert chunk == {"output": 4}
+        else:
+            assert False, "Expected only two chunks"
+    assert cleanup.call_count == 1, "Expected cleanup to be called once"

tests/test_pregel_async.py

@@ -0,0 +1,473 @@
+import asyncio
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, AsyncIterator, Generator
+
+import pytest
+from langchain.schema.runnable import RunnablePassthrough
+from pytest_mock import MockerFixture
+
+from permchain import Channel, Pregel
+from permchain.channels.base import InvalidUpdateError
+from permchain.channels.context import Context
+from permchain.channels.inbox import Inbox
+from permchain.channels.last_value import LastValue
+
+
+async def test_invoke_single_process_in_out(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,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    assert await app.ainvoke(2) == 3
+
+
+async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
+
+    app = Pregel(
+        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.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {"output": {"title": "Output", "type": "integer"}},
+    }
+    assert await app.ainvoke(2) == {"output": 3}
+
+
+async def test_invoke_single_process_in_dict_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")
+
+    app = Pregel(
+        chains={
+            "one": chain,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+        },
+        input=["input"],
+        output=["output"],
+    )
+
+    assert app.input_schema.schema() == {
+        "title": "PregelInput",
+        "type": "object",
+        "properties": {"input": {"title": "Input", "type": "integer"}},
+    }
+    assert app.output_schema.schema() == {
+        "title": "PregelOutput",
+        "type": "object",
+        "properties": {"output": {"title": "Output", "type": "integer"}},
+    }
+    assert await app.ainvoke({"input": 2}) == {"output": 3}
+
+
+async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    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",
+    )
+
+    assert await app.ainvoke(2) == 4
+
+
+async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    pubsub = Pregel(
+        chains={"chain_one": chain_one, "chain_two": chain_two},
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+        },
+        input=["input", "inbox"],
+        output="output",
+    )
+
+    # [12 + 1, 2 + 1 + 1]
+    assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [13, 4]
+
+
+async def test_batch_two_processes_in_out() -> None:
+    async def add_one_with_delay(inp: int) -> int:
+        await asyncio.sleep(inp / 10)
+        return inp + 1
+
+    chain_one = (
+        Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
+    )
+    chain_two = (
+        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",
+    )
+
+    assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
+
+
+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")
+
+    # No state is left over from previous invocations
+    for _ in range(10):
+        assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size
+
+    # Concurrent invocations do not interfere with each other
+    assert await asyncio.gather(
+        *(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10))
+    ) == [2 + test_size for _ in range(10)]
+
+
+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")
+
+    # No state is left over from previous invocations
+    for _ in range(10):
+        # Then invoke pubsub
+        assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
+            2 + test_size,
+            1 + test_size,
+            3 + test_size,
+            4 + test_size,
+            5 + test_size,
+        ]
+
+    # 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)
+        )
+    ) == [
+        [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
+        for _ in range(10)
+    ]
+
+
+async def test_invoke_two_processes_two_in_two_out_invalid(
+    mocker: MockerFixture,
+) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = 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",
+    )
+
+    with pytest.raises(InvalidUpdateError):
+        # LastValue channels can only be updated once per iteration
+        await app.ainvoke(2)
+
+
+async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = 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": Inbox(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    # An Inbox channel accumulates updates into a sequence
+    assert await app.ainvoke(2) == (3, 3)
+
+
+async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
+
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_four = (
+        Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
+    )
+
+    app = Pregel(
+        chains={
+            "chain_one": chain_one,
+            "chain_three": chain_three,
+            "chain_four": chain_four,
+        },
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+        },
+        input="input",
+        output="output",
+    )
+
+    # Then invoke app
+    # We get a single array result as chain_four waits for all publishers to finish
+    # before operating on all elements published to topic_two as an array
+    for _ in range(100):
+        assert await app.ainvoke(2) == [13, 13]
+
+    assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [
+        [13, 13] for _ in range(100)
+    ]
+
+
+async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
+
+    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 = (
+        Channel.subscribe_to("input")
+        | add_10_each
+        | Channel.write_to("inbox_one").map()
+    )
+    chain_two = (
+        Channel.subscribe_to("inbox_one")
+        | inner_app.map()
+        | sorted
+        | Channel.write_to("outbox_one")
+    )
+    chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
+
+    app = Pregel(
+        chains={
+            "chain_one": chain_one,
+            "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",
+    )
+
+    # Then invoke pubsub
+    for _ in range(10):
+        assert await app.ainvoke([2, 3]) == 27
+
+    assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [
+        27 for _ in range(10)
+    ]
+
+
+async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+
+    chain_one = (
+        Channel.subscribe_to("input")
+        | add_one
+        | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
+    )
+    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",
+    )
+
+    # Then invoke pubsub
+    assert [c async for c in app.astream(2)] == [3, 4]
+
+
+async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    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",
+    )
+
+    # Then invoke pubsub
+    # It finishes executing (once no more messages being published)
+    # but returns nothing, as nothing was published to OUT topic
+    assert await app.ainvoke(2) is None
+
+
+async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
+    setup_sync = mocker.Mock()
+    cleanup_sync = mocker.Mock()
+    setup_async = mocker.Mock()
+    cleanup_async = mocker.Mock()
+
+    @contextmanager
+    def an_int() -> Generator[int, None, None]:
+        setup_sync()
+        try:
+            yield 5
+        finally:
+            cleanup_sync()
+
+    @asynccontextmanager
+    async def an_int_async() -> AsyncGenerator[int, None]:
+        setup_async()
+        try:
+            yield 5
+        finally:
+            cleanup_async()
+
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    chain_two = (
+        Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one, "chain_two": chain_two},
+        channels={
+            "input": LastValue(int),
+            "output": LastValue(int),
+            "inbox": Inbox(int),
+            "ctx": Context(an_int, an_int_async, typ=int),
+        },
+        input="input",
+        output=["inbox", "output"],
+    )
+
+    async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]:
+        i = 0
+        async for chunk in aiter:
+            yield i, chunk
+            i += 1
+
+    assert setup_sync.call_count == 0
+    assert cleanup_sync.call_count == 0
+    assert setup_async.call_count == 0
+    assert cleanup_async.call_count == 0
+    async for i, chunk in aenumerate(app.astream(2)):
+        assert setup_sync.call_count == 0, "Sync context manager should not be used"
+        assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
+        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,)}
+        elif i == 1:
+            assert chunk == {"output": 4}
+        else:
+            assert False, "Expected only two chunks"
+    assert setup_sync.call_count == 0
+    assert cleanup_sync.call_count == 0
+    assert setup_async.call_count == 1, "Expected setup to be called once"
+    assert cleanup_async.call_count == 1, "Expected cleanup to be called once"