diff --git a/README.md b/README.md index 3b9ceb24a..0dd4c3e04 100644 --- a/README.md +++ b/README.md @@ -22,13 +22,15 @@ Some of the use cases are: Channels are used to communicate between chains. Each channel has a value type, an update type, and an update function – which takes a sequence of updates and modifies the stored value. Channels can be used to send data from one chain to another, or to send data from a chain to itself in a future step. PermChain provides a number of built-in channels: -- `LastValue`: stores the last value sent to the channel, useful for input values, and single-value outputs -- `Inbox`: stores an ephemeral sequence of values sent to the channel, useful for sending data from one chain to another -- `UniqueInbox`: same as Inbox, but deduplicates values sent to the channel -- `Archive`: stores a persistent sequence of values sent to the channel, useful for accumulating data over multiple steps -- `UniqueArchive`: same as Archive, but deduplicates values sent to the channel -- `BinaryOperatorAggregate`: stores a persistent value, updated by applying a binary operator to the current value and each update sent to the channel, useful for computing aggregates over multiple steps. eg. `total = BinaryOperatorAggregate(int, operator.add)` +#### Basic channels: LastValue and Topic + +- `LastValue`: The default channel, stores the last value sent to the channel, useful for input and output values, or for sending data from one step to the next +- `Topic`: A configurable PubSub Topic, useful for sending multiple values between chains, or for accumulating output. Can be configured to deduplicate values, and/or to accummulate values over the course of multiple steps. + +#### Advanced channels: Context and BinaryOperatorAggregate + - `Context`: exposes the value of a context manager, managing its lifecycle. Useful for accessing external resources that require setup and/or teardown. eg. `client = Context(httpx.Client)` +- `BinaryOperatorAggregate`: stores a persistent value, updated by applying a binary operator to the current value and each update sent to the channel, useful for computing aggregates over multiple steps. eg. `total = BinaryOperatorAggregate(int, operator.add)` ### Chains @@ -48,7 +50,6 @@ Repeat until no chains are planned for execution, or a maximum number of steps i ```python from permchain import Channel, Pregel -from permchain.channels import LastValue grow_value = ( Channel.subscribe_to("value") @@ -58,13 +59,11 @@ grow_value = ( app = Pregel( chains={"grow_value": grow_value}, - channels={"value": LastValue(str)}, input="value", output="value", ) assert app.invoke("a") == "aaaaaaaa" - ``` Check `examples` for more examples. diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index 6fefc2516..7ad00124e 100644 --- a/examples/combine_docs.ipynb +++ b/examples/combine_docs.ipynb @@ -25,8 +25,8 @@ "from langchain.schema.document import Document\n", "from langchain.schema import format_document\n", "\n", - "from permchain import Channel, Pregel, PregelRead\n", - "from permchain.channels import LastValue, Inbox" + "from permchain import Channel, Pregel\n", + "from permchain.channels import LastValue, Topic" ] }, { @@ -209,18 +209,15 @@ "source": [ "channels = {\n", " # input\n", - " \"question\": LastValue(str),\n", - " \"docs\": Inbox(Document),\n", + " \"docs\": Topic(Document),\n", " # intermediate\n", - " \"docs_to_finalize\": Inbox(Document),\n", - " # output\n", - " \"answer\": LastValue(str),\n", + " \"docs_to_finalize\": Topic(Document),\n", "}" ] }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 23, "id": "67370694-86f4-4b64-9d4f-38b2e306abeb", "metadata": {}, "outputs": [], @@ -234,12 +231,16 @@ " return Channel.write_to(\"docs_to_finalize\")\n", "\n", "\n", + "def split_docs_with_question(input: dict[str, str | list[Document]]) -> list[dict[str, str | list[Document]]]:\n", + " return [\n", + " {\"docs\": docs, \"question\": input[\"question\"]}\n", + " for docs in _split_list_of_docs(input[\"docs\"])\n", + " ]\n", + "\n", + "\n", "collapse = (\n", - " Channel.subscribe_to(\"docs\")\n", - " | _split_list_of_docs\n", - " | {\"docs_list\": RunnablePassthrough(), \"question\": PregelRead(\"question\")}\n", - " # {docs: list[list[Doc]], question: str} -> list[{docs: list[Doc], question: str}]\n", - " | (lambda x: [{\"docs\": docs, \"question\": x[\"question\"]} for docs in x[\"docs_list\"]])\n", + " Channel.subscribe_to([\"docs\", \"question\"])\n", + " | split_docs_with_question\n", " | stuff_chain.map() # Collapse each list of docs to a single string\n", " | (lambda x: [Document(page_content=s) for s in x]) # A new (smaller) list of docs\n", " | decide\n", @@ -255,7 +256,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 24, "id": "3019e7d2-ab7f-4868-b43c-ad898d824a26", "metadata": {}, "outputs": [], @@ -274,7 +275,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 25, "id": "69fcb829-3dae-432a-8db3-11bbb179a7d2", "metadata": {}, "outputs": [ @@ -283,51 +284,52 @@ "output_type": "stream", "text": [ "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 0 with 1 task. Next tasks:\n", - "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho'),\n", - " Document(page_content='Ankush worked at Facebook'),\n", - " Document(page_content='Harrison used to work at Kensho'),\n", - " Document(page_content='Ankush worked at Facebook'),\n", - " Document(page_content='Harrison used to work at Kensho'),\n", - " Document(page_content='Ankush worked at Facebook'),\n", - " Document(page_content='Harrison used to work at Kensho'),\n", - " Document(page_content='Ankush worked at Facebook'),\n", - " Document(page_content='Harrison used to work at Kensho'),\n", - " Document(page_content='Ankush worked at Facebook')))\n", + "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho'),\n", + " Document(page_content='Ankush worked at Facebook'),\n", + " Document(page_content='Harrison used to work at Kensho'),\n", + " Document(page_content='Ankush worked at Facebook'),\n", + " Document(page_content='Harrison used to work at Kensho'),\n", + " Document(page_content='Ankush worked at Facebook'),\n", + " Document(page_content='Harrison used to work at Kensho'),\n", + " Document(page_content='Ankush worked at Facebook'),\n", + " Document(page_content='Harrison used to work at Kensho'),\n", + " Document(page_content='Ankush worked at Facebook')],\n", + " 'question': 'where did harrison work'})\n", "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 0. Channel values:\n", - "\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n", + "\u001b[0m{'docs': [...], 'docs_to_finalize': [], 'question': 'where did harrison work'}\n", "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 1 with 1 task. Next tasks:\n", - "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.')))\n", + "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.')],\n", + " 'question': 'where did harrison work'})\n", "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 1. Channel values:\n", - "\u001b[0m{'docs': (...), 'question': 'where did harrison work'}\n", + "\u001b[0m{'docs': [...], 'docs_to_finalize': [], 'question': 'where did harrison work'}\n", "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 2 with 1 task. Next tasks:\n", - "\u001b[0m- collapse((Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.')))\n", + "\u001b[0m- collapse({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.')],\n", + " 'question': 'where did harrison work'})\n", "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 2. Channel values:\n", - "\u001b[0m{'docs': (...),\n", - " 'docs_to_finalize': (...),\n", - " 'question': 'where did harrison work'}\n", + "\u001b[0m{'docs': [], 'docs_to_finalize': [...], 'question': 'where did harrison work'}\n", "\u001b[36;1m\u001b[1;3m[pregel/step]\u001b[0m \u001b[1mStarting step 3 with 1 task. Next tasks:\n", - "\u001b[0m- finalize({'docs': (Document(page_content='Harrison used to work at Kensho.'),\n", - " Document(page_content='Harrison used to work at Kensho.'))})\n", + "\u001b[0m- finalize({'docs': [Document(page_content='Harrison used to work at Kensho.'),\n", + " Document(page_content='Harrison used to work at Kensho.')]})\n", "\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 3. Channel values:\n", - "\u001b[0m{'answer': 'Harrison worked at Kensho.',\n", - " 'docs': (...),\n", - " 'docs_to_finalize': (...),\n", + "\u001b[0m{'answer': 'Harrison used to work at Kensho.',\n", + " 'docs': [],\n", + " 'docs_to_finalize': [],\n", " 'question': 'where did harrison work'}\n" ] }, { "data": { "text/plain": [ - "'Harrison worked at Kensho.'" + "'Harrison used to work at Kensho.'" ] }, - "execution_count": 13, + "execution_count": 25, "metadata": {}, "output_type": "execute_result" } diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index 489bb3970..2619f828a 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -75,12 +75,6 @@ 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") @@ -105,7 +99,6 @@ reviser = ( ) draft_revise_loop = Pregel( - channels=channels, chains={ "drafter": drafter, "editor": editor, @@ -113,27 +106,12 @@ draft_revise_loop = Pregel( }, # input will be a dict with a single key, "question" input=["question"], - # output will be a dict with keys "draft" and "notes" - output=["draft", "notes"], + # output will be the value of "draft" + output="draft", # debug logging debug=True, ) # run -for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}): - print(draft) - print("---") - - -async def main(): - async for draft in draft_revise_loop.astream( - {"question": "What food do turtles eat?"} - ): - print(draft) - print("---") - - -# import asyncio - -# asyncio.run(main()) +print(draft_revise_loop.invoke({"question": "What food do turtles eat?"})) diff --git a/examples/readme.py b/examples/readme.py index 9e87a74c3..0d86505d4 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -1,5 +1,4 @@ from permchain import Channel, Pregel -from permchain.channels import LastValue grow_value = ( Channel.subscribe_to("value") @@ -9,7 +8,6 @@ grow_value = ( app = Pregel( chains={"grow_value": grow_value}, - channels={"value": LastValue(str)}, input="value", output="value", ) diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index 241745123..5d4b8fae1 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -7,7 +7,8 @@ from langchain.schema.runnable import RunnableLambda, RunnablePassthrough from langchain.utils.html import extract_sub_links from permchain import Channel, Pregel -from permchain.channels import Archive, Context, LastValue, UniqueArchive, UniqueInbox +from permchain.channels.context import Context +from permchain.channels.topic import Topic # Load url with sync httpx client @@ -85,39 +86,28 @@ def recursive_web_loader( metadata_extractor = metadata_extractor or _metadata_extractor # define the channels channels = { - "base_url": LastValue(str), - "next_urls": UniqueInbox(str), - "documents": Archive(Document), - "visited": UniqueArchive(str), + "next_urls": Topic(str, unique=True), + "documents": Topic(Document, accumulate=True), "client": Context(httpx_client, httpx_aclient), } # the main chain that gets executed recursively + # while there are urls in next_urls visitor = ( - # while there are urls in next_urls # run the chain below for each url in next_urls - # adding the current values of visited set, base_url and httpx client - Channel.subscribe_to_each("next_urls", key="url").join( - ["visited", "client", "base_url"] - ) + # adding the current values of base_url and httpx client + Channel.subscribe_to_each("next_urls", key="url").join(["client", "base_url"]) # load the url (with sync and async implementations) | RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async)) | Channel.write_to( - # send this url to the visited set - visited=lambda x: x["url"], # send a new document to the documents stream documents=lambda x: Document( page_content=extractor(x["body"]), metadata=metadata_extractor(x["body"], x["url"]), ), - # send the next urls to the next_urls set - # only if not visited already - next_urls=lambda x: [ - url - for url in extract_sub_links( - x["body"], x["url"], base_url=x["base_url"] - ) - if url not in x["visited"] and url != x["url"] - ], + # send the next urls to the next_urls topic + next_urls=lambda x: extract_sub_links( + x["body"], x["url"], base_url=x["base_url"] + ), ) ) return Pregel( diff --git a/permchain/channels/__init__.py b/permchain/channels/__init__.py index 90d415010..cc4de04bb 100644 --- a/permchain/channels/__init__.py +++ b/permchain/channels/__init__.py @@ -1,15 +1,11 @@ -from permchain.channels.archive import Archive, UniqueArchive from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.context import Context -from permchain.channels.inbox import Inbox, UniqueInbox from permchain.channels.last_value import LastValue +from permchain.channels.topic import Topic __all__ = [ "LastValue", - "Inbox", - "UniqueInbox", - "Archive", - "UniqueArchive", - "BinaryOperatorAggregate", + "Topic", "Context", + "BinaryOperatorAggregate", ] diff --git a/permchain/channels/archive.py b/permchain/channels/archive.py deleted file mode 100644 index a491e9971..000000000 --- a/permchain/channels/archive.py +++ /dev/null @@ -1,94 +0,0 @@ -import json -from contextlib import contextmanager -from typing import Any, FrozenSet, Generator, Generic, Optional, Sequence, Type - -from typing_extensions import Self - -from permchain.channels.base import BaseChannel, EmptyChannelError, Value -from permchain.channels.inbox import flatten - - -class Archive(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): - """Stores all unique values received, persists across steps.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - self.set = list[Value]() - - @property - def ValueType(self) -> Any: - """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.set = json.loads(checkpoint) - try: - yield empty - finally: - pass - - def update(self, values: Sequence[Value | list[Value]]) -> None: - self.set.extend(flatten(values)) - - def get(self) -> Sequence[Value]: - try: - return tuple(self.set) - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - try: - return json.dumps(self.set) - except AttributeError: - raise EmptyChannelError() - - -class UniqueArchive(Generic[Value], BaseChannel[FrozenSet[Value], Value]): - """Stores all unique values received, persists across steps.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - self.set = set[Value]() - - @property - def ValueType(self) -> Type[FrozenSet[Value]]: - """The type of the value stored in the channel.""" - return FrozenSet[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.set = set(json.loads(checkpoint)) - try: - yield empty - finally: - pass - - def update(self, values: Sequence[Value | list[Value]]) -> None: - self.set.update(flatten(values)) - - def get(self) -> FrozenSet[Value]: - try: - return frozenset(self.set) - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - try: - return json.dumps(list(self.set)) - except AttributeError: - raise EmptyChannelError() diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py index 8f125673b..981322b46 100644 --- a/permchain/channels/binop.py +++ b/permchain/channels/binop.py @@ -45,6 +45,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]): pass def update(self, values: Sequence[Value]) -> None: + if not values: + return if not hasattr(self, "value"): self.value = values[0] values = values[1:] diff --git a/permchain/channels/context.py b/permchain/channels/context.py index 99d10c19e..183fdafd6 100644 --- a/permchain/channels/context.py +++ b/permchain/channels/context.py @@ -97,7 +97,8 @@ class Context(Generic[Value], BaseChannel[Value, None]): yield empty def update(self, values: Sequence[None]) -> None: - raise InvalidUpdateError() + if values: + raise InvalidUpdateError() def get(self) -> Value: try: diff --git a/permchain/channels/inbox.py b/permchain/channels/inbox.py deleted file mode 100644 index 60972a79b..000000000 --- a/permchain/channels/inbox.py +++ /dev/null @@ -1,115 +0,0 @@ -import json -from contextlib import contextmanager -from typing import ( - Any, - FrozenSet, - Generator, - Generic, - Iterator, - Optional, - Sequence, - Type, - Union, -) - -from typing_extensions import Self - -from permchain.channels.base import BaseChannel, EmptyChannelError, Value - - -def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]: - for value in values: - if isinstance(value, list): - yield from value - else: - yield value - - -class Inbox(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): - """Stores all values received, resets in each step.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - - @property - def ValueType(self) -> Type[Sequence[Value]]: - """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Any: - """The type of the update received by the channel.""" - return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.queue = tuple(json.loads(checkpoint)) - try: - yield empty - finally: - try: - del empty.queue - except AttributeError: - pass - - def update(self, values: Sequence[Value | list[Value]]) -> None: - self.queue = tuple(flatten(values)) - - def get(self) -> Sequence[Value]: - try: - return self.queue - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - try: - return json.dumps(self.queue) - except AttributeError: - raise EmptyChannelError() - - -class UniqueInbox(Generic[Value], BaseChannel[FrozenSet[Value], Value | list[Value]]): - """Stores all unique values received, resets in each step.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - - @property - def ValueType(self) -> Type[FrozenSet[Value]]: - """The type of the value stored in the channel.""" - return FrozenSet[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Any: - """The type of the update received by the channel.""" - return Union[self.typ, Sequence[self.typ]] # type: ignore[name-defined] - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.queue = frozenset(json.loads(checkpoint)) - try: - yield empty - finally: - try: - del empty.queue - except AttributeError: - pass - - def update(self, values: Sequence[Value | list[Value]]) -> None: - self.queue = frozenset(flatten(values)) - - def get(self) -> FrozenSet[Value]: - try: - return self.queue - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - try: - return json.dumps(self.queue) - except AttributeError: - raise EmptyChannelError() diff --git a/permchain/channels/last_value.py b/permchain/channels/last_value.py index 953d16dad..6dc60e138 100644 --- a/permchain/channels/last_value.py +++ b/permchain/channels/last_value.py @@ -42,6 +42,8 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]): pass def update(self, values: Sequence[Value]) -> None: + if len(values) == 0: + return if len(values) != 1: raise InvalidUpdateError() diff --git a/permchain/channels/topic.py b/permchain/channels/topic.py new file mode 100644 index 000000000..71b821efe --- /dev/null +++ b/permchain/channels/topic.py @@ -0,0 +1,76 @@ +import json +from contextlib import contextmanager +from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, Union + +from typing_extensions import Self + +from permchain.channels.base import BaseChannel, Value + + +def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]: + for value in values: + if isinstance(value, list): + yield from value + else: + yield value + + +class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): + """A configurable PubSub Topic. + + Args: + typ: The type of the value stored in the channel. + unique: Whether to discard duplicate values. + accumulate: Whether to accummulate values across steps. If False, the channel will be emptied after each step. + """ + + def __init__( + self, typ: Type[Value], unique: bool = False, accumulate: bool = False + ) -> None: + # attrs + self.typ = typ + self.unique = unique + self.accumulate = accumulate + # state + self.seen = set[Value]() + self.values = list[Value]() + + @property + def ValueType(self) -> Any: + """The type of the value stored in the channel.""" + return Sequence[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return Union[self.typ, list[self.typ]] # type: ignore[name-defined] + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ, self.unique, self.accumulate) + if checkpoint is not None: + parsed = json.loads(checkpoint) + empty.seen = set(parsed["seen"]) + empty.values = list(parsed["values"]) + try: + yield empty + finally: + pass + + def update(self, values: Sequence[Value | list[Value]]) -> None: + if not self.accumulate: + self.values = list[Value]() + if flat_values := flatten(values): + if self.unique: + for value in flat_values: + if value not in self.seen: + self.seen.add(value) + self.values.append(value) + else: + self.values.extend(flat_values) + + def get(self) -> Sequence[Value]: + return list(self.values) + + def checkpoint(self) -> str: + return json.dumps({"seen": list(self.seen), "values": self.values}) diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 7132ae62a..c76ad01a7 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -99,13 +99,13 @@ class Channel: class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): - channels: Mapping[str, BaseChannel] - chains: Mapping[str, ChannelInvoke | ChannelBatch] - output: str | Sequence[str] + channels: Mapping[str, BaseChannel] = Field(default_factory=dict) - input: str | Sequence[str] + output: str | Sequence[str] = "output" + + input: str | Sequence[str] = "input" step_timeout: Optional[float] = None @@ -421,6 +421,10 @@ def _apply_writes_and_prepare_next_tasks( updated_channels.add(chan) else: logger.warning(f"Skipping write for channel {chan} which has no readers") + # Channels that weren't updated in this step are notified of a new step + for chan in channels: + if chan not in updated_channels: + channels[chan].update([]) tasks: list[tuple[Runnable, Any, str]] = [] # Check if any processes should be run in next step diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index eb42c972d..56aa31961 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -1,6 +1,7 @@ -from typing import Mapping, Sequence +from typing import Any, Mapping, Sequence from permchain.channels.base import BaseChannel +from permchain.channels.last_value import LastValue from permchain.pregel.read import ChannelBatch, ChannelInvoke @@ -23,7 +24,7 @@ def validate_chains_channels( for chan in subscribed_channels: if chan not in channels: - raise ValueError(f"Channel {chan} is subscribed to, but not initialized") + channels[chan] = LastValue(Any) if isinstance(input, str): if input not in subscribed_channels: @@ -36,8 +37,8 @@ def validate_chains_channels( if isinstance(output, str): if output not in channels: - raise ValueError(f"Output channel {output} is not initialized") + channels[output] = LastValue(Any) else: for chan in output: if chan not in channels: - raise ValueError(f"Output channel {chan} is not initialized") + channels[chan] = LastValue(Any) diff --git a/poetry.lock b/poetry.lock index afc2b69e1..93fbf093e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -840,73 +840,68 @@ files = [ [[package]] name = "greenlet" -version = "3.0.0" +version = "3.0.1" description = "Lightweight in-process concurrent programming" optional = false python-versions = ">=3.7" files = [ - {file = "greenlet-3.0.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e09dea87cc91aea5500262993cbd484b41edf8af74f976719dd83fe724644cd6"}, - {file = "greenlet-3.0.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f47932c434a3c8d3c86d865443fadc1fbf574e9b11d6650b656e602b1797908a"}, - {file = "greenlet-3.0.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bdfaeecf8cc705d35d8e6de324bf58427d7eafb55f67050d8f28053a3d57118c"}, - {file = "greenlet-3.0.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6a68d670c8f89ff65c82b936275369e532772eebc027c3be68c6b87ad05ca695"}, - {file = "greenlet-3.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38ad562a104cd41e9d4644f46ea37167b93190c6d5e4048fcc4b80d34ecb278f"}, - {file = "greenlet-3.0.0-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02a807b2a58d5cdebb07050efe3d7deaf915468d112dfcf5e426d0564aa3aa4a"}, - {file = "greenlet-3.0.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:b1660a15a446206c8545edc292ab5c48b91ff732f91b3d3b30d9a915d5ec4779"}, - {file = "greenlet-3.0.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:813720bd57e193391dfe26f4871186cf460848b83df7e23e6bef698a7624b4c9"}, - {file = "greenlet-3.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:aa15a2ec737cb609ed48902b45c5e4ff6044feb5dcdfcf6fa8482379190330d7"}, - {file = "greenlet-3.0.0-cp310-universal2-macosx_11_0_x86_64.whl", hash = "sha256:7709fd7bb02b31908dc8fd35bfd0a29fc24681d5cc9ac1d64ad07f8d2b7db62f"}, - {file = "greenlet-3.0.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:211ef8d174601b80e01436f4e6905aca341b15a566f35a10dd8d1e93f5dbb3b7"}, - {file = "greenlet-3.0.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6512592cc49b2c6d9b19fbaa0312124cd4c4c8a90d28473f86f92685cc5fef8e"}, - {file = "greenlet-3.0.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:871b0a8835f9e9d461b7fdaa1b57e3492dd45398e87324c047469ce2fc9f516c"}, - {file = "greenlet-3.0.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b505fcfc26f4148551826a96f7317e02c400665fa0883fe505d4fcaab1dabfdd"}, - {file = "greenlet-3.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:123910c58234a8d40eaab595bc56a5ae49bdd90122dde5bdc012c20595a94c14"}, - {file = "greenlet-3.0.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:96d9ea57292f636ec851a9bb961a5cc0f9976900e16e5d5647f19aa36ba6366b"}, - {file = "greenlet-3.0.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0b72b802496cccbd9b31acea72b6f87e7771ccfd7f7927437d592e5c92ed703c"}, - {file = "greenlet-3.0.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:527cd90ba3d8d7ae7dceb06fda619895768a46a1b4e423bdb24c1969823b8362"}, - {file = "greenlet-3.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:37f60b3a42d8b5499be910d1267b24355c495064f271cfe74bf28b17b099133c"}, - {file = "greenlet-3.0.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:1482fba7fbed96ea7842b5a7fc11d61727e8be75a077e603e8ab49d24e234383"}, - {file = "greenlet-3.0.0-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:be557119bf467d37a8099d91fbf11b2de5eb1fd5fc5b91598407574848dc910f"}, - {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:73b2f1922a39d5d59cc0e597987300df3396b148a9bd10b76a058a2f2772fc04"}, - {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d1e22c22f7826096ad503e9bb681b05b8c1f5a8138469b255eb91f26a76634f2"}, - {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1d363666acc21d2c204dd8705c0e0457d7b2ee7a76cb16ffc099d6799744ac99"}, - {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:334ef6ed8337bd0b58bb0ae4f7f2dcc84c9f116e474bb4ec250a8bb9bd797a66"}, - {file = "greenlet-3.0.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6672fdde0fd1a60b44fb1751a7779c6db487e42b0cc65e7caa6aa686874e79fb"}, - {file = "greenlet-3.0.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:952256c2bc5b4ee8df8dfc54fc4de330970bf5d79253c863fb5e6761f00dda35"}, - {file = "greenlet-3.0.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:269d06fa0f9624455ce08ae0179430eea61085e3cf6457f05982b37fd2cefe17"}, - {file = "greenlet-3.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:9adbd8ecf097e34ada8efde9b6fec4dd2a903b1e98037adf72d12993a1c80b51"}, - {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c6b5ce7f40f0e2f8b88c28e6691ca6806814157ff05e794cdd161be928550f4c"}, - {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ecf94aa539e97a8411b5ea52fc6ccd8371be9550c4041011a091eb8b3ca1d810"}, - {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:80dcd3c938cbcac986c5c92779db8e8ce51a89a849c135172c88ecbdc8c056b7"}, - {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e52a712c38e5fb4fd68e00dc3caf00b60cb65634d50e32281a9d6431b33b4af1"}, - {file = "greenlet-3.0.0-cp37-cp37m-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5539f6da3418c3dc002739cb2bb8d169056aa66e0c83f6bacae0cd3ac26b423"}, - {file = "greenlet-3.0.0-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:343675e0da2f3c69d3fb1e894ba0a1acf58f481f3b9372ce1eb465ef93cf6fed"}, - {file = "greenlet-3.0.0-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:abe1ef3d780de56defd0c77c5ba95e152f4e4c4e12d7e11dd8447d338b85a625"}, - {file = "greenlet-3.0.0-cp37-cp37m-win32.whl", hash = "sha256:e693e759e172fa1c2c90d35dea4acbdd1d609b6936115d3739148d5e4cd11947"}, - {file = "greenlet-3.0.0-cp37-cp37m-win_amd64.whl", hash = "sha256:bdd696947cd695924aecb3870660b7545a19851f93b9d327ef8236bfc49be705"}, - {file = "greenlet-3.0.0-cp37-universal2-macosx_11_0_x86_64.whl", hash = "sha256:cc3e2679ea13b4de79bdc44b25a0c4fcd5e94e21b8f290791744ac42d34a0353"}, - {file = "greenlet-3.0.0-cp38-cp38-macosx_11_0_universal2.whl", hash = "sha256:63acdc34c9cde42a6534518e32ce55c30f932b473c62c235a466469a710bfbf9"}, - {file = "greenlet-3.0.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4a1a6244ff96343e9994e37e5b4839f09a0207d35ef6134dce5c20d260d0302c"}, - {file = "greenlet-3.0.0-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b822fab253ac0f330ee807e7485769e3ac85d5eef827ca224feaaefa462dc0d0"}, - {file = "greenlet-3.0.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8060b32d8586e912a7b7dac2d15b28dbbd63a174ab32f5bc6d107a1c4143f40b"}, - {file = "greenlet-3.0.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:621fcb346141ae08cb95424ebfc5b014361621b8132c48e538e34c3c93ac7365"}, - {file = "greenlet-3.0.0-cp38-cp38-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6bb36985f606a7c49916eff74ab99399cdfd09241c375d5a820bb855dfb4af9f"}, - {file = "greenlet-3.0.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:10b5582744abd9858947d163843d323d0b67be9432db50f8bf83031032bc218d"}, - {file = "greenlet-3.0.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:f351479a6914fd81a55c8e68963609f792d9b067fb8a60a042c585a621e0de4f"}, - {file = "greenlet-3.0.0-cp38-cp38-win32.whl", hash = "sha256:9de687479faec7db5b198cc365bc34addd256b0028956501f4d4d5e9ca2e240a"}, - {file = "greenlet-3.0.0-cp38-cp38-win_amd64.whl", hash = "sha256:3fd2b18432e7298fcbec3d39e1a0aa91ae9ea1c93356ec089421fabc3651572b"}, - {file = "greenlet-3.0.0-cp38-universal2-macosx_11_0_x86_64.whl", hash = "sha256:3c0d36f5adc6e6100aedbc976d7428a9f7194ea79911aa4bf471f44ee13a9464"}, - {file = "greenlet-3.0.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:4cd83fb8d8e17633ad534d9ac93719ef8937568d730ef07ac3a98cb520fd93e4"}, - {file = "greenlet-3.0.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6a5b2d4cdaf1c71057ff823a19d850ed5c6c2d3686cb71f73ae4d6382aaa7a06"}, - {file = "greenlet-3.0.0-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2e7dcdfad252f2ca83c685b0fa9fba00e4d8f243b73839229d56ee3d9d219314"}, - {file = "greenlet-3.0.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c94e4e924d09b5a3e37b853fe5924a95eac058cb6f6fb437ebb588b7eda79870"}, - {file = "greenlet-3.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad6fb737e46b8bd63156b8f59ba6cdef46fe2b7db0c5804388a2d0519b8ddb99"}, - {file = "greenlet-3.0.0-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d55db1db455c59b46f794346efce896e754b8942817f46a1bada2d29446e305a"}, - {file = "greenlet-3.0.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:56867a3b3cf26dc8a0beecdb4459c59f4c47cdd5424618c08515f682e1d46692"}, - {file = "greenlet-3.0.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:9a812224a5fb17a538207e8cf8e86f517df2080c8ee0f8c1ed2bdaccd18f38f4"}, - {file = "greenlet-3.0.0-cp39-cp39-win32.whl", hash = "sha256:0d3f83ffb18dc57243e0151331e3c383b05e5b6c5029ac29f754745c800f8ed9"}, - {file = "greenlet-3.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:831d6f35037cf18ca5e80a737a27d822d87cd922521d18ed3dbc8a6967be50ce"}, - {file = "greenlet-3.0.0-cp39-universal2-macosx_11_0_x86_64.whl", hash = "sha256:a048293392d4e058298710a54dfaefcefdf49d287cd33fb1f7d63d55426e4355"}, - {file = "greenlet-3.0.0.tar.gz", hash = "sha256:19834e3f91f485442adc1ee440171ec5d9a4840a1f7bd5ed97833544719ce10b"}, + {file = "greenlet-3.0.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:f89e21afe925fcfa655965ca8ea10f24773a1791400989ff32f467badfe4a064"}, + {file = "greenlet-3.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:28e89e232c7593d33cac35425b58950789962011cc274aa43ef8865f2e11f46d"}, + {file = "greenlet-3.0.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b8ba29306c5de7717b5761b9ea74f9c72b9e2b834e24aa984da99cbfc70157fd"}, + {file = "greenlet-3.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:19bbdf1cce0346ef7341705d71e2ecf6f41a35c311137f29b8a2dc2341374565"}, + {file = "greenlet-3.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:599daf06ea59bfedbec564b1692b0166a0045f32b6f0933b0dd4df59a854caf2"}, + {file = "greenlet-3.0.1-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b641161c302efbb860ae6b081f406839a8b7d5573f20a455539823802c655f63"}, + {file = "greenlet-3.0.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d57e20ba591727da0c230ab2c3f200ac9d6d333860d85348816e1dca4cc4792e"}, + {file = "greenlet-3.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:5805e71e5b570d490938d55552f5a9e10f477c19400c38bf1d5190d760691846"}, + {file = "greenlet-3.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:52e93b28db27ae7d208748f45d2db8a7b6a380e0d703f099c949d0f0d80b70e9"}, + {file = "greenlet-3.0.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f7bfb769f7efa0eefcd039dd19d843a4fbfbac52f1878b1da2ed5793ec9b1a65"}, + {file = "greenlet-3.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:91e6c7db42638dc45cf2e13c73be16bf83179f7859b07cfc139518941320be96"}, + {file = "greenlet-3.0.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1757936efea16e3f03db20efd0cd50a1c86b06734f9f7338a90c4ba85ec2ad5a"}, + {file = "greenlet-3.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:19075157a10055759066854a973b3d1325d964d498a805bb68a1f9af4aaef8ec"}, + {file = "greenlet-3.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e9d21aaa84557d64209af04ff48e0ad5e28c5cca67ce43444e939579d085da72"}, + {file = "greenlet-3.0.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2847e5d7beedb8d614186962c3d774d40d3374d580d2cbdab7f184580a39d234"}, + {file = "greenlet-3.0.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:97e7ac860d64e2dcba5c5944cfc8fa9ea185cd84061c623536154d5a89237884"}, + {file = "greenlet-3.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:b2c02d2ad98116e914d4f3155ffc905fd0c025d901ead3f6ed07385e19122c94"}, + {file = "greenlet-3.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:22f79120a24aeeae2b4471c711dcf4f8c736a2bb2fabad2a67ac9a55ea72523c"}, + {file = "greenlet-3.0.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:100f78a29707ca1525ea47388cec8a049405147719f47ebf3895e7509c6446aa"}, + {file = "greenlet-3.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:60d5772e8195f4e9ebf74046a9121bbb90090f6550f81d8956a05387ba139353"}, + {file = "greenlet-3.0.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:daa7197b43c707462f06d2c693ffdbb5991cbb8b80b5b984007de431493a319c"}, + {file = "greenlet-3.0.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea6b8aa9e08eea388c5f7a276fabb1d4b6b9d6e4ceb12cc477c3d352001768a9"}, + {file = "greenlet-3.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8d11ebbd679e927593978aa44c10fc2092bc454b7d13fdc958d3e9d508aba7d0"}, + {file = "greenlet-3.0.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dbd4c177afb8a8d9ba348d925b0b67246147af806f0b104af4d24f144d461cd5"}, + {file = "greenlet-3.0.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:20107edf7c2c3644c67c12205dc60b1bb11d26b2610b276f97d666110d1b511d"}, + {file = "greenlet-3.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8bef097455dea90ffe855286926ae02d8faa335ed8e4067326257cb571fc1445"}, + {file = "greenlet-3.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:b2d3337dcfaa99698aa2377c81c9ca72fcd89c07e7eb62ece3f23a3fe89b2ce4"}, + {file = "greenlet-3.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:80ac992f25d10aaebe1ee15df45ca0d7571d0f70b645c08ec68733fb7a020206"}, + {file = "greenlet-3.0.1-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:337322096d92808f76ad26061a8f5fccb22b0809bea39212cd6c406f6a7060d2"}, + {file = "greenlet-3.0.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b9934adbd0f6e476f0ecff3c94626529f344f57b38c9a541f87098710b18af0a"}, + {file = "greenlet-3.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dc4d815b794fd8868c4d67602692c21bf5293a75e4b607bb92a11e821e2b859a"}, + {file = "greenlet-3.0.1-cp37-cp37m-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:41bdeeb552d814bcd7fb52172b304898a35818107cc8778b5101423c9017b3de"}, + {file = "greenlet-3.0.1-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:6e6061bf1e9565c29002e3c601cf68569c450be7fc3f7336671af7ddb4657166"}, + {file = "greenlet-3.0.1-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:fa24255ae3c0ab67e613556375a4341af04a084bd58764731972bcbc8baeba36"}, + {file = "greenlet-3.0.1-cp37-cp37m-win32.whl", hash = "sha256:b489c36d1327868d207002391f662a1d163bdc8daf10ab2e5f6e41b9b96de3b1"}, + {file = "greenlet-3.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:f33f3258aae89da191c6ebaa3bc517c6c4cbc9b9f689e5d8452f7aedbb913fa8"}, + {file = "greenlet-3.0.1-cp38-cp38-macosx_11_0_universal2.whl", hash = "sha256:d2905ce1df400360463c772b55d8e2518d0e488a87cdea13dd2c71dcb2a1fa16"}, + {file = "greenlet-3.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a02d259510b3630f330c86557331a3b0e0c79dac3d166e449a39363beaae174"}, + {file = "greenlet-3.0.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:55d62807f1c5a1682075c62436702aaba941daa316e9161e4b6ccebbbf38bda3"}, + {file = "greenlet-3.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3fcc780ae8edbb1d050d920ab44790201f027d59fdbd21362340a85c79066a74"}, + {file = "greenlet-3.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4eddd98afc726f8aee1948858aed9e6feeb1758889dfd869072d4465973f6bfd"}, + {file = "greenlet-3.0.1-cp38-cp38-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eabe7090db68c981fca689299c2d116400b553f4b713266b130cfc9e2aa9c5a9"}, + {file = "greenlet-3.0.1-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:f2f6d303f3dee132b322a14cd8765287b8f86cdc10d2cb6a6fae234ea488888e"}, + {file = "greenlet-3.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:d923ff276f1c1f9680d32832f8d6c040fe9306cbfb5d161b0911e9634be9ef0a"}, + {file = "greenlet-3.0.1-cp38-cp38-win32.whl", hash = "sha256:0b6f9f8ca7093fd4433472fd99b5650f8a26dcd8ba410e14094c1e44cd3ceddd"}, + {file = "greenlet-3.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:990066bff27c4fcf3b69382b86f4c99b3652bab2a7e685d968cd4d0cfc6f67c6"}, + {file = "greenlet-3.0.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:ce85c43ae54845272f6f9cd8320d034d7a946e9773c693b27d620edec825e376"}, + {file = "greenlet-3.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89ee2e967bd7ff85d84a2de09df10e021c9b38c7d91dead95b406ed6350c6997"}, + {file = "greenlet-3.0.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:87c8ceb0cf8a5a51b8008b643844b7f4a8264a2c13fcbcd8a8316161725383fe"}, + {file = "greenlet-3.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d6a8c9d4f8692917a3dc7eb25a6fb337bff86909febe2f793ec1928cd97bedfc"}, + {file = "greenlet-3.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9fbc5b8f3dfe24784cee8ce0be3da2d8a79e46a276593db6868382d9c50d97b1"}, + {file = "greenlet-3.0.1-cp39-cp39-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:85d2b77e7c9382f004b41d9c72c85537fac834fb141b0296942d52bf03fe4a3d"}, + {file = "greenlet-3.0.1-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:696d8e7d82398e810f2b3622b24e87906763b6ebfd90e361e88eb85b0e554dc8"}, + {file = "greenlet-3.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:329c5a2e5a0ee942f2992c5e3ff40be03e75f745f48847f118a3cfece7a28546"}, + {file = "greenlet-3.0.1-cp39-cp39-win32.whl", hash = "sha256:cf868e08690cb89360eebc73ba4be7fb461cfbc6168dd88e2fbbe6f31812cd57"}, + {file = "greenlet-3.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:ac4a39d1abae48184d420aa8e5e63efd1b75c8444dd95daa3e03f6c6310e9619"}, + {file = "greenlet-3.0.1.tar.gz", hash = "sha256:816bd9488a94cba78d93e1abb58000e8266fa9cc2aa9ccdd6eb0696acb24005b"}, ] [package.extras] @@ -2063,27 +2058,6 @@ files = [ docs = ["furo (>=2023.7.26)", "proselint (>=0.13)", "sphinx (>=7.1.1)", "sphinx-autodoc-typehints (>=1.24)"] test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.4)", "pytest-cov (>=4.1)", "pytest-mock (>=3.11.1)"] -[[package]] -name = "playwright" -version = "1.39.0" -description = "A high-level API to automate web browsers" -optional = false -python-versions = ">=3.8" -files = [ - {file = "playwright-1.39.0-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:384e195a6d09343f319031cf552e9cd601ede78fe9c082b9fa197537c5cbfe7a"}, - {file = "playwright-1.39.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:d2c3634411828d9273196ed6f69f2fa7645c89732b3c982dcf09ab03ed4c5d2b"}, - {file = "playwright-1.39.0-py3-none-macosx_11_0_universal2.whl", hash = "sha256:d2fd90f370599cf9a2c6a041bd79a5eeec62baf0e943c7c5c2079b29be476d2a"}, - {file = "playwright-1.39.0-py3-none-manylinux1_x86_64.whl", hash = "sha256:699a8e707ca5f3567aa28223ee1be7e42d2bf25eda7d3d86babda71e36e5f16f"}, - {file = "playwright-1.39.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:654bb3ae0dc3c69ffddc0c38c127c3b8e93032d8cf3928e2c4f21890cb39514b"}, - {file = "playwright-1.39.0-py3-none-win32.whl", hash = "sha256:40ed7f2546c64f1bb3d22b2295b4d43ed5a2f0b7ea7599d93a72f723a1883e1e"}, - {file = "playwright-1.39.0-py3-none-win_amd64.whl", hash = "sha256:a420d814e21b05e1156747e2a9fae6c3cca2b46bb4a0226fb26ee65538ce09c9"}, -] - -[package.dependencies] -greenlet = "3.0.0" -pyee = "11.0.1" -typing-extensions = {version = "*", markers = "python_version <= \"3.8\""} - [[package]] name = "pluggy" version = "1.3.0" @@ -2328,23 +2302,6 @@ files = [ [package.dependencies] typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0" -[[package]] -name = "pyee" -version = "11.0.1" -description = "A rough port of Node.js's EventEmitter to Python with a few tricks of its own" -optional = false -python-versions = ">=3.8" -files = [ - {file = "pyee-11.0.1-py3-none-any.whl", hash = "sha256:9bcc9647822234f42c228d88de63d0f9ffa881e87a87f9d36ddf5211f6ac977d"}, - {file = "pyee-11.0.1.tar.gz", hash = "sha256:a642c51e3885a33ead087286e35212783a4e9b8d6514a10a5db4e57ac57b2b29"}, -] - -[package.dependencies] -typing-extensions = "*" - -[package.extras] -dev = ["black", "flake8", "flake8-black", "isort", "jupyter-console", "mkdocs", "mkdocs-include-markdown-plugin", "mkdocstrings[python]", "pytest", "pytest-asyncio", "pytest-trio", "toml", "tox", "trio", "trio", "trio-typing", "twine", "twisted", "validate-pyproject[all]"] - [[package]] name = "pygments" version = "2.16.1" @@ -2963,19 +2920,19 @@ win32 = ["pywin32"] [[package]] name = "setuptools" -version = "67.8.0" +version = "68.2.2" description = "Easily download, build, install, upgrade, and uninstall Python packages" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "setuptools-67.8.0-py3-none-any.whl", hash = "sha256:5df61bf30bb10c6f756eb19e7c9f3b473051f48db77fddbe06ff2ca307df9a6f"}, - {file = "setuptools-67.8.0.tar.gz", hash = "sha256:62642358adc77ffa87233bc4d2354c4b2682d214048f500964dbe760ccedf102"}, + {file = "setuptools-68.2.2-py3-none-any.whl", hash = "sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a"}, + {file = "setuptools-68.2.2.tar.gz", hash = "sha256:4ac1475276d2f1c48684874089fefcd83bd7162ddaafb81fac866ba0db282a87"}, ] [package.extras] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-hoverxref (<2)", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (==0.8.3)", "sphinx-reredirects", "sphinxcontrib-towncrier"] -testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pip-run (>=8.8)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-ruff", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"] -testing-integration = ["build[virtualenv]", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"] +docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-hoverxref (<2)", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier"] +testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-ruff", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"] +testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "packaging (>=23.1)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"] [[package]] name = "six" @@ -3524,4 +3481,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "0abda52d59fdb79b9a0d9ddb838fc6952fc92715bd7fe075da17bc25ec2489f0" +content-hash = "ce9a7fe6e1972d14c6fd8aec807d4146098f6e4bc0efcbb0e8f10d922ff4901f" diff --git a/pyproject.toml b/pyproject.toml index 820f0f0d2..0e58b144b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,8 +35,6 @@ optional = true [tool.poetry.group.dev.dependencies] jupyter = "^1.0.0" -playwright = "^1.28.0" -setuptools = "^67.6.1" openai = "^0.27.8" [tool.ruff] diff --git a/tests/test_channels.py b/tests/test_channels.py index 55067afc6..2b562da62 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -1,17 +1,16 @@ import operator from contextlib import asynccontextmanager, contextmanager -from typing import AsyncGenerator, FrozenSet, Generator, Sequence, Union +from typing import AsyncGenerator, Generator, Sequence, Union import httpx import pytest from pytest_mock import MockerFixture -from permchain.channels.archive import UniqueArchive from permchain.channels.base import EmptyChannelError, InvalidUpdateError from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.context import Context -from permchain.channels.inbox import Inbox from permchain.channels.last_value import LastValue +from permchain.channels.topic import Topic def test_last_value() -> None: @@ -46,57 +45,108 @@ async def test_last_value_async() -> None: assert channel.get() == 4 -def test_inbox() -> None: - with Inbox(str).empty() as channel: +def test_topic() -> None: + with Topic(str).empty() as channel: assert channel.ValueType is Sequence[str] - assert channel.UpdateType is Union[str, Sequence[str]] - - with pytest.raises(EmptyChannelError): - channel.get() + assert channel.UpdateType is Union[str, list[str]] channel.update(["a", "b"]) - assert channel.get() == ("a", "b") - channel.update([["c"], "d"]) - assert channel.get() == ("c", "d") + assert channel.get() == ["a", "b"] + channel.update([["c", "d"], "d"]) + assert channel.get() == ["c", "d", "d"] + channel.update([]) + assert channel.get() == [] -async def test_inbox_async() -> None: - async with Inbox(str).aempty() as channel: +async def test_topic_async() -> None: + async with Topic(str).aempty() as channel: assert channel.ValueType is Sequence[str] - assert channel.UpdateType is Union[str, Sequence[str]] - - with pytest.raises(EmptyChannelError): - channel.get() + assert channel.UpdateType is Union[str, list[str]] channel.update(["a", "b"]) - assert channel.get() == ("a", "b") - channel.update(["c"]) - channel.update([["c"], "d"]) - assert channel.get() == ("c", "d") + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["b", "c", "d", "d"] + channel.update([]) + assert channel.get() == [] -def test_set() -> None: - with UniqueArchive(str).empty() as channel: - assert channel.ValueType is FrozenSet[str] - assert channel.UpdateType is str +def test_topic_unique() -> None: + with Topic(str, unique=True).empty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] - assert channel.get() == frozenset() channel.update(["a", "b"]) - assert channel.get() == frozenset(("a", "b")) - channel.update(["b", "c"]) - assert channel.get() == frozenset(("a", "b", "c")) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" + channel.update([]) + assert channel.get() == [] -async def test_set_async() -> None: - async with UniqueArchive(str).aempty() as channel: - assert channel.ValueType is FrozenSet[str] - assert channel.UpdateType is str +async def test_topic_unique_async() -> None: + async with Topic(str, unique=True).aempty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] - assert channel.get() == frozenset() channel.update(["a", "b"]) - assert channel.get() == frozenset(("a", "b")) - channel.update(["b", "c"]) - assert channel.get() == frozenset(("a", "b", "c")) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" + channel.update([]) + assert channel.get() == [] + + +def test_topic_accumulate() -> None: + with Topic(str, accumulate=True).empty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] + + channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + channel.update([]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + + +async def test_topic_accumulate_async() -> None: + async with Topic(str, accumulate=True).aempty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] + + channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + channel.update([]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + + +def test_topic_unique_accumulate() -> None: + with Topic(str, unique=True, accumulate=True).empty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] + + channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["a", "b", "c", "d"] + channel.update([]) + assert channel.get() == ["a", "b", "c", "d"] + + +async def test_topic_unique_accumulate_async() -> None: + async with Topic(str, unique=True, accumulate=True).aempty() as channel: + assert channel.ValueType is Sequence[str] + assert channel.UpdateType is Union[str, list[str]] + + channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["a", "b", "c", "d"] + channel.update([]) + assert channel.get() == ["a", "b", "c", "d"] def test_binop() -> None: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 5946b5461..619221270 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -10,8 +10,8 @@ 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 +from permchain.channels.topic import Topic def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -35,6 +35,17 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 3 +def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + + app = Pregel(chains={"one": chain}) + + assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.output_schema.schema() == {"title": "PregelOutput"} + assert app.invoke(2) == 3 + + def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") @@ -43,19 +54,14 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: chains={ "one": chain, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", output=["output"], ) - assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == { "title": "PregelOutput", "type": "object", - "properties": {"output": {"title": "Output", "type": "integer"}}, + "properties": {"output": {"title": "Output"}}, } assert app.invoke(2) == {"output": 3} @@ -68,10 +74,6 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: chains={ "one": chain, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, input=["input"], output=["output"], ) @@ -79,12 +81,12 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: assert app.input_schema.schema() == { "title": "PregelInput", "type": "object", - "properties": {"input": {"title": "Input", "type": "integer"}}, + "properties": {"input": {"title": "Input"}}, } assert app.output_schema.schema() == { "title": "PregelOutput", "type": "object", - "properties": {"output": {"title": "Output", "type": "integer"}}, + "properties": {"output": {"title": "Output"}}, } assert app.invoke({"input": 2}) == {"output": 3} @@ -98,13 +100,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, - input="input", - output="output", + channels={"inbox": Topic(int)}, ) assert app.invoke(2) == 4 @@ -119,13 +115,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, + channels={"inbox": Topic(int)}, input=["input", "inbox"], - output="output", ) assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4] # [12 + 1, 2 + 1 + 1] @@ -143,16 +134,7 @@ def test_batch_two_processes_in_out() -> None: Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") ) - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "one": LastValue(int), - }, - input="input", - output="output", - ) + app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -161,20 +143,14 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chans = { - "input": LastValue(int), - "output": LastValue(int), - "-1": LastValue(int), - } chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = LastValue(int) chains[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains) for _ in range(10): assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size @@ -189,20 +165,14 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chans = { - "input": LastValue(int), - "output": LastValue(int), - "-1": LastValue(int), - } chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = LastValue(int) chains[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains) for _ in range(3): assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ @@ -229,15 +199,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", - output="output", - ) + app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -252,16 +214,11 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": Inbox(int), - }, - input="input", - output="output", + channels={"output": Topic(int)}, ) # An Inbox channel accumulates updates into a sequence - assert app.invoke(2) == (3, 3) + assert app.invoke(2) == [3, 3] def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -280,13 +237,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None "chain_three": chain_three, "chain_four": chain_four, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, - input="input", - output="output", + channels={"inbox": Topic(int)}, ) # Then invoke app @@ -306,13 +257,7 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: inner_app = Pregel( chains={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", - output="output", + } ) chain_one = ( @@ -334,14 +279,7 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: "chain_two": chain_two, "chain_three": chain_three, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox_one": Inbox(int), - "outbox_one": LastValue(int), - }, - input="input", - output="output", + channels={"inbox_one": Topic(int)}, ) for _ in range(10): @@ -363,13 +301,6 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "between": LastValue(int), - }, - input="input", - output="output", ) assert [c for c in app.stream(2)] == [3, 4] @@ -382,13 +313,6 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "between": LastValue(int), - }, - input="input", - output="output", ) # It finishes executing (once no more messages being published) @@ -405,13 +329,6 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: with pytest.raises(ValueError): Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "between": LastValue(int), - }, - input="input", - output="output", ) @@ -436,12 +353,9 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), + "inbox": Topic(int), "ctx": Context(an_int, typ=int), }, - input="input", output=["inbox", "output"], ) @@ -451,7 +365,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert setup.call_count == 1, "Expected setup to be called once" assert cleanup.call_count == 0, "Expected cleanup to not be called yet" if i == 0: - assert chunk == {"inbox": (3,)} + assert chunk == {"inbox": [3]} elif i == 1: assert chunk == {"output": 4} else: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3ecc8634e..b191b7438 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -9,8 +9,8 @@ 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 +from permchain.channels.topic import Topic async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -29,6 +29,21 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: output="output", ) + assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"} + assert await app.ainvoke(2) == 3 + + +async def test_invoke_single_process_in_out_implicit_channels( + mocker: MockerFixture +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + + app = Pregel(chains={"one": chain}) + + assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.output_schema.schema() == {"title": "PregelOutput"} assert await app.ainvoke(2) == 3 @@ -40,19 +55,14 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: chains={ "one": chain, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", output=["output"], ) - assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + assert app.input_schema.schema() == {"title": "PregelInput"} assert app.output_schema.schema() == { "title": "PregelOutput", "type": "object", - "properties": {"output": {"title": "Output", "type": "integer"}}, + "properties": {"output": {"title": "Output"}}, } assert await app.ainvoke(2) == {"output": 3} @@ -65,10 +75,6 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> chains={ "one": chain, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, input=["input"], output=["output"], ) @@ -76,12 +82,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> assert app.input_schema.schema() == { "title": "PregelInput", "type": "object", - "properties": {"input": {"title": "Input", "type": "integer"}}, + "properties": {"input": {"title": "Input"}}, } assert app.output_schema.schema() == { "title": "PregelOutput", "type": "object", - "properties": {"output": {"title": "Output", "type": "integer"}}, + "properties": {"output": {"title": "Output"}}, } assert await app.ainvoke({"input": 2}) == {"output": 3} @@ -95,13 +101,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, - input="input", - output="output", + channels={"inbox": Topic(int)}, ) assert await app.ainvoke(2) == 4 @@ -116,13 +116,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: pubsub = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, + channels={"inbox": Topic(int)}, input=["input", "inbox"], - output="output", ) # [12 + 1, 2 + 1 + 1] @@ -143,13 +138,7 @@ async def test_batch_two_processes_in_out() -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "one": LastValue(int), - }, - input="input", - output="output", + channels={"one": LastValue(int)}, ) assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -159,20 +148,14 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chans = { - "input": LastValue(int), - "output": LastValue(int), - "-1": LastValue(int), - } chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = LastValue(int) chains[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains) # No state is left over from previous invocations for _ in range(10): @@ -188,20 +171,14 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - chans = { - "input": LastValue(int), - "output": LastValue(int), - "-1": LastValue(int), - } chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} for i in range(test_size - 2): - chans[str(i)] = LastValue(int) chains[str(i)] = ( Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) ) chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(chains=chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains) # No state is left over from previous invocations for _ in range(3): @@ -231,15 +208,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid( chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", - output="output", - ) + app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -254,16 +223,11 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": Inbox(int), - }, - input="input", - output="output", + channels={"output": Topic(int)}, ) - # An Inbox channel accumulates updates into a sequence - assert await app.ainvoke(2) == (3, 3) + # An Topic channel accumulates updates into a sequence + assert await app.ainvoke(2) == [3, 3] async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -282,13 +246,7 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - "chain_three": chain_three, "chain_four": chain_four, }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), - }, - input="input", - output="output", + channels={"inbox": Topic(int)}, ) # Then invoke app @@ -309,13 +267,7 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None inner_app = Pregel( chains={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - }, - channels={ - "input": LastValue(int), - "output": LastValue(int), - }, - input="input", - output="output", + } ) chain_one = ( @@ -338,13 +290,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None "chain_three": chain_three, }, channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox_one": Inbox(int), + "inbox_one": Topic(int), "outbox_one": LastValue(int), }, - input="input", - output="output", ) # Then invoke pubsub @@ -366,16 +314,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non ) chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "between": LastValue(int), - }, - input="input", - output="output", - ) + app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) # Then invoke pubsub assert [c async for c in app.astream(2)] == [3, 4] @@ -386,20 +325,10 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") chain_two = Channel.subscribe_to("between") | add_one - app = Pregel( - chains={"chain_one": chain_one, "chain_two": chain_two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "between": LastValue(int), - }, - input="input", - output="output", - ) + app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two}) - # Then invoke pubsub # It finishes executing (once no more messages being published) - # but returns nothing, as nothing was published to OUT topic + # but returns nothing, as nothing was published to "output" topic assert await app.ainvoke(2) is None @@ -434,12 +363,9 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Inbox(int), + "inbox": Topic(int), "ctx": Context(an_int, an_int_async, typ=int), }, - input="input", output=["inbox", "output"], ) @@ -459,7 +385,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert setup_async.call_count == 1, "Expected setup to be called once" assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet" if i == 0: - assert chunk == {"inbox": (3,)} + assert chunk == {"inbox": [3]} elif i == 1: assert chunk == {"output": 4} else: