From dbe8c1362f47af09980f0bebdbd97a30a187e5ee Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 15:42:51 +0100 Subject: [PATCH] Move subscribe/write cls methods to Channel --- README.md | 9 +- examples/combine_docs.ipynb | 25 +++-- examples/draft-revise-loop.py | 21 ++-- examples/readme.py | 9 +- examples/recursive-web-loader.py | 19 ++-- permchain/__init__.py | 5 +- permchain/pregel/__init__.py | 98 ++++++++-------- permchain/pregel/read.py | 2 +- permchain/pregel/validate.py | 2 +- tests/test_channels.py | 28 +++-- tests/test_pregel.py | 185 +++++++++++++++++-------------- tests/test_pregel_async.py | 175 +++++++++++++++-------------- 12 files changed, 307 insertions(+), 271 deletions(-) diff --git a/README.md b/README.md index e7cc038c1..aad68045e 100644 --- a/README.md +++ b/README.md @@ -47,17 +47,18 @@ Repeat until no chains are planned for execution, or a maximum number of steps i ## Example ```python -from permchain import Channels, Pregel +from permchain import Channel, Pregel +from permchain.channels import LastValue grow_value = ( - Pregel.subscribe_to("value") + Channel.subscribe_to("value") | (lambda x: x + x) - | Pregel.write_to(value=lambda x: x if len(x) < 10 else None) + | Channel.write_to(value=lambda x: x if len(x) < 10 else None) ) app = Pregel( chains={"grow_value": grow_value}, - channels={"value": Channels.LastValue(str)}, + channels={"value": LastValue(str)}, input="value", output="value", ) diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index ee83cbd21..6fefc2516 100644 --- a/examples/combine_docs.ipynb +++ b/examples/combine_docs.ipynb @@ -25,7 +25,8 @@ "from langchain.schema.document import Document\n", "from langchain.schema import format_document\n", "\n", - "from permchain import Channels, Pregel, PregelRead" + "from permchain import Channel, Pregel, PregelRead\n", + "from permchain.channels import LastValue, Inbox" ] }, { @@ -208,12 +209,12 @@ "source": [ "channels = {\n", " # input\n", - " \"question\": Channels.LastValue(str),\n", - " \"docs\": Channels.Inbox(Document),\n", + " \"question\": LastValue(str),\n", + " \"docs\": Inbox(Document),\n", " # intermediate\n", - " \"docs_to_finalize\": Channels.Inbox(Document),\n", + " \"docs_to_finalize\": Inbox(Document),\n", " # output\n", - " \"answer\": Channels.LastValue(str),\n", + " \"answer\": LastValue(str),\n", "}" ] }, @@ -227,14 +228,14 @@ "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 Pregel.write_to(\"docs\")\n", + " return Channel.write_to(\"docs\")\n", " else:\n", " # send to the finalizer if we're ready to produce final answer\n", - " return Pregel.write_to(\"docs_to_finalize\")\n", + " return Channel.write_to(\"docs_to_finalize\")\n", "\n", "\n", "collapse = (\n", - " Pregel.subscribe_to(\"docs\")\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", @@ -246,9 +247,9 @@ "\n", "# Convert final set of docs to an answer\n", "finalize = (\n", - " Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n", + " Channel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n", " | stuff_chain\n", - " | Pregel.write_to(\"answer\")\n", + " | Channel.write_to(\"answer\")\n", ")" ] }, @@ -314,7 +315,7 @@ "\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 used to work at Kensho.',\n", + "\u001b[0m{'answer': 'Harrison worked at Kensho.',\n", " 'docs': (...),\n", " 'docs_to_finalize': (...),\n", " 'question': 'where did harrison work'}\n" @@ -323,7 +324,7 @@ { "data": { "text/plain": [ - "'Harrison used to work at Kensho.'" + "'Harrison worked at Kensho.'" ] }, "execution_count": 13, diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index 7384f64f3..050c064a8 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -5,7 +5,8 @@ from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser from langchain.prompts import SystemMessagePromptTemplate from langchain.schema.output_parser import StrOutputParser -from permchain import Channels, Pregel +from permchain import Channel, Pregel +from permchain.channels import LastValue # prompts @@ -75,23 +76,23 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser() # application channels = { - "question": Channels.LastValue(str), - "draft": Channels.LastValue(str), - "notes": Channels.LastValue(str), + "question": LastValue(str), + "draft": LastValue(str), + "notes": LastValue(str), } drafter = ( # subscribe to question channel as a dict with a single key, "question" - Pregel.subscribe_to(["question"]) + Channel.subscribe_to(["question"]) | drafter_chain - | Pregel.write_to("draft") + | Channel.write_to("draft") ) editor = ( # subscribe to draft channel as a dict with a single key, "draft" - Pregel.subscribe_to(["draft"]) + Channel.subscribe_to(["draft"]) | editor_chain - | Pregel.write_to( + | 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" @@ -102,9 +103,9 @@ editor = ( reviser = ( # subscribe to new values of "notes" channel, # and join them with the input value (question) and "draft" - Pregel.subscribe_to(["notes"]).join(["question", "draft"]) + Channel.subscribe_to(["notes"]).join(["question", "draft"]) | reviser_chain - | Pregel.write_to("draft") + | Channel.write_to("draft") ) draft_revise_loop = Pregel( diff --git a/examples/readme.py b/examples/readme.py index 647396b17..9e87a74c3 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -1,14 +1,15 @@ -from permchain import Channels, Pregel +from permchain import Channel, Pregel +from permchain.channels import LastValue grow_value = ( - Pregel.subscribe_to("value") + Channel.subscribe_to("value") | (lambda x: x + x) - | Pregel.write_to(value=lambda x: x if len(x) < 10 else None) + | Channel.write_to(value=lambda x: x if len(x) < 10 else None) ) app = Pregel( chains={"grow_value": grow_value}, - channels={"value": Channels.LastValue(str)}, + channels={"value": LastValue(str)}, input="value", output="value", ) diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index c16881380..241745123 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -6,7 +6,8 @@ from langchain.schema import Document from langchain.schema.runnable import RunnableLambda, RunnablePassthrough from langchain.utils.html import extract_sub_links -from permchain import Channels, Pregel +from permchain import Channel, Pregel +from permchain.channels import Archive, Context, LastValue, UniqueArchive, UniqueInbox # Load url with sync httpx client @@ -84,23 +85,23 @@ def recursive_web_loader( metadata_extractor = metadata_extractor or _metadata_extractor # define the channels channels = { - "base_url": Channels.LastValue(str), - "next_urls": Channels.UniqueInbox(str), - "documents": Channels.Archive(Document), - "visited": Channels.UniqueArchive(str), - "client": Channels.Context(httpx_client, httpx_aclient), + "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 - Pregel.subscribe_to_each("next_urls", key="url").join( + 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)) - | Pregel.write_to( + | Channel.write_to( # send this url to the visited set visited=lambda x: x["url"], # send a new document to the documents stream @@ -123,7 +124,7 @@ def recursive_web_loader( channels=channels, chains={ # use the base_url as the first url to visit - "input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"), + "input": Channel.subscribe_to("base_url") | Channel.write_to("next_urls"), # add the main chain "visitor": visitor, }, diff --git a/permchain/__init__.py b/permchain/__init__.py index 732952882..57ff1145c 100644 --- a/permchain/__init__.py +++ b/permchain/__init__.py @@ -1,5 +1,4 @@ -import permchain.channels as Channels -from permchain.pregel import Pregel +from permchain.pregel import Channel, Pregel from permchain.pregel.read import PregelRead -__all__ = ["Channels", "Pregel", "PregelRead"] +__all__ = ["Channel", "Pregel", "PregelRead"] diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 11ff4f16e..593993410 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -48,6 +48,56 @@ 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] @@ -106,54 +156,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): **{k: (self.channels[k].ValueType, None) for k in self.output}, ) - @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()] - ) - ) - def _transform( self, input: Iterator[dict[str, Any] | Any], diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py index 709b4e6b8..19884243a 100644 --- a/permchain/pregel/read.py +++ b/permchain/pregel/read.py @@ -123,7 +123,7 @@ class PregelBatch(RunnableEach): if self.key is None: raise ValueError( "Cannot join() additional channels without a key." - " Pass a key arg to Pregel.subscribe_to_each()." + " Pass a key arg to Channel.subscribe_to_each()." ) joiner = RunnablePassthrough.assign( diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index bd269845c..80bdb3b01 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -18,7 +18,7 @@ def validate_chains_channels( subscribed_channels.add(chain.channel) else: raise TypeError( - f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()" + f"Invalid chain type {type(chain)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()" ) for chan in subscribed_channels: diff --git a/tests/test_channels.py b/tests/test_channels.py index 4b45a67c6..55067afc6 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -6,12 +6,16 @@ import httpx import pytest from pytest_mock import MockerFixture -import permchain.channels as Channels +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 Channels.LastValue(int).empty() as channel: + with LastValue(int).empty() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -27,7 +31,7 @@ def test_last_value() -> None: async def test_last_value_async() -> None: - async with Channels.LastValue(int).aempty() as channel: + async with LastValue(int).aempty() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -43,7 +47,7 @@ async def test_last_value_async() -> None: def test_inbox() -> None: - with Channels.Inbox(str).empty() as channel: + with Inbox(str).empty() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, Sequence[str]] @@ -57,7 +61,7 @@ def test_inbox() -> None: async def test_inbox_async() -> None: - async with Channels.Inbox(str).aempty() as channel: + async with Inbox(str).aempty() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, Sequence[str]] @@ -72,7 +76,7 @@ async def test_inbox_async() -> None: def test_set() -> None: - with Channels.UniqueArchive(str).empty() as channel: + with UniqueArchive(str).empty() as channel: assert channel.ValueType is FrozenSet[str] assert channel.UpdateType is str @@ -84,7 +88,7 @@ def test_set() -> None: async def test_set_async() -> None: - async with Channels.UniqueArchive(str).aempty() as channel: + async with UniqueArchive(str).aempty() as channel: assert channel.ValueType is FrozenSet[str] assert channel.UpdateType is str @@ -96,7 +100,7 @@ async def test_set_async() -> None: def test_binop() -> None: - with Channels.BinaryOperatorAggregate(int, operator.add).empty() as channel: + with BinaryOperatorAggregate(int, operator.add).empty() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -110,7 +114,7 @@ def test_binop() -> None: async def test_binop_async() -> None: - async with Channels.BinaryOperatorAggregate(int, operator.add).aempty() as channel: + async with BinaryOperatorAggregate(int, operator.add).aempty() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -135,7 +139,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None: finally: cleanup() - with Channels.Context(an_int, None, int).empty() as channel: + with Context(an_int, None, int).empty() as channel: assert setup.call_count == 1 assert cleanup.call_count == 0 @@ -153,7 +157,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None: def test_ctx_manager_ctx(mocker: MockerFixture) -> None: - with Channels.Context(httpx.Client).empty() as channel: + with Context(httpx.Client).empty() as channel: assert channel.ValueType is httpx.Client with pytest.raises(InvalidUpdateError): assert channel.UpdateType is None @@ -183,7 +187,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None: finally: cleanup() - async with Channels.Context(an_int_sync, an_int, int).aempty() as channel: + async with Context(an_int_sync, an_int, int).aempty() as channel: assert setup.call_count == 1 assert cleanup.call_count == 0 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6f544c671..daf31b232 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -7,21 +7,24 @@ import pytest from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture -from permchain import Channels, Pregel +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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", @@ -34,15 +37,15 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output=["output"], @@ -59,15 +62,15 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input=["input"], output=["output"], @@ -88,15 +91,17 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input="input", output="output", @@ -107,15 +112,17 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input=["input", "inbox"], output="output", @@ -130,18 +137,18 @@ def test_batch_two_processes_in_out() -> None: return inp + 1 chain_one = ( - Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one") + Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") ) chain_two = ( - Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "one": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "one": LastValue(int), }, input="input", output="output", @@ -155,17 +162,17 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chans = { - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "-1": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "-1": LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} + chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = Channels.LastValue(int) + chans[str(i)] = LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) + Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") + chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -183,17 +190,17 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chans = { - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "-1": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "-1": LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} + chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = Channels.LastValue(int) + chans[str(i)] = LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) + Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") + chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -219,14 +226,14 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: 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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", @@ -240,14 +247,14 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N 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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.Inbox(int), + "input": LastValue(int), + "output": Inbox(int), }, input="input", output="output", @@ -261,9 +268,11 @@ 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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output") + 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={ @@ -272,9 +281,9 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None "chain_four": chain_four, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input="input", output="output", @@ -296,26 +305,28 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: inner_app = Pregel( chains={ - "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", ) chain_one = ( - Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map() + Channel.subscribe_to("input") + | add_10_each + | Channel.write_to("inbox_one").map() ) chain_two = ( - Pregel.subscribe_to("inbox_one") + Channel.subscribe_to("inbox_one") | inner_app.map() | sorted - | Pregel.write_to("outbox_one") + | Channel.write_to("outbox_one") ) - chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output") + chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") app = Pregel( chains={ @@ -324,10 +335,10 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: "chain_three": chain_three, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox_one": Channels.Inbox(int), - "outbox_one": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox_one": Inbox(int), + "outbox_one": LastValue(int), }, input="input", output="output", @@ -344,18 +355,18 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain_one = ( - Pregel.subscribe_to("input") + Channel.subscribe_to("input") | add_one - | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) + | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "between": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "between": LastValue(int), }, input="input", output="output", @@ -366,15 +377,15 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between") - chain_two = Pregel.subscribe_to("between") | add_one + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "between": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "between": LastValue(int), }, input="input", output="output", @@ -388,16 +399,16 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") - chain_two = Pregel.subscribe_to("between") | add_one + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "between": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "between": LastValue(int), }, input="input", output="output", @@ -417,16 +428,18 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: cleanup() add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), - "ctx": Channels.Context(an_int, typ=int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), + "ctx": Context(an_int, typ=int), }, input="input", output=["inbox", "output"], diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 99487adfd..cdf827612 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -6,21 +6,24 @@ import pytest from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture -from permchain import Channels, Pregel +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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", @@ -31,15 +34,15 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output=["output"], @@ -56,15 +59,15 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( chains={ "one": chain, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input=["input"], output=["output"], @@ -85,15 +88,17 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input="input", output="output", @@ -104,15 +109,17 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input=["input", "inbox"], output="output", @@ -128,18 +135,18 @@ async def test_batch_two_processes_in_out() -> None: return inp + 1 chain_one = ( - Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one") + Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") ) chain_two = ( - Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "one": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "one": LastValue(int), }, input="input", output="output", @@ -153,17 +160,17 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chans = { - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "-1": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "-1": LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} + chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = Channels.LastValue(int) + chans[str(i)] = LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) + Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") + chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -182,17 +189,17 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chans = { - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "-1": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "-1": LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} + chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = Channels.LastValue(int) + chans[str(i)] = LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) + Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") + chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -224,14 +231,14 @@ async def test_invoke_two_processes_two_in_two_out_invalid( ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", @@ -245,14 +252,14 @@ async def test_invoke_two_processes_two_in_two_out_invalid( 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 = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.Inbox(int), + "input": LastValue(int), + "output": Inbox(int), }, input="input", output="output", @@ -266,9 +273,11 @@ async 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)) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output") + 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={ @@ -277,9 +286,9 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - "chain_four": chain_four, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), }, input="input", output="output", @@ -302,26 +311,28 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None inner_app = Pregel( chains={ - "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), }, input="input", output="output", ) chain_one = ( - Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map() + Channel.subscribe_to("input") + | add_10_each + | Channel.write_to("inbox_one").map() ) chain_two = ( - Pregel.subscribe_to("inbox_one") + Channel.subscribe_to("inbox_one") | inner_app.map() | sorted - | Pregel.write_to("outbox_one") + | Channel.write_to("outbox_one") ) - chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output") + chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") app = Pregel( chains={ @@ -330,10 +341,10 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None "chain_three": chain_three, }, channels={ - "input": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox_one": Channels.Inbox(int), - "outbox_one": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "inbox_one": Inbox(int), + "outbox_one": LastValue(int), }, input="input", output="output", @@ -352,18 +363,18 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non add_one = mocker.Mock(side_effect=lambda x: x + 1) chain_one = ( - Pregel.subscribe_to("input") + Channel.subscribe_to("input") | add_one - | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) + | Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "between": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "between": LastValue(int), }, input="input", output="output", @@ -375,15 +386,15 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between") - chain_two = Pregel.subscribe_to("between") | add_one + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "between": Channels.LastValue(int), + "input": LastValue(int), + "output": LastValue(int), + "between": LastValue(int), }, input="input", output="output", @@ -418,16 +429,18 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: cleanup_async() add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") + 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": Channels.LastValue(int), - "output": Channels.LastValue(int), - "inbox": Channels.Inbox(int), - "ctx": Channels.Context(an_int, an_int_async, typ=int), + "input": LastValue(int), + "output": LastValue(int), + "inbox": Inbox(int), + "ctx": Context(an_int, an_int_async, typ=int), }, input="input", output=["inbox", "output"],