From d0df4103cd7a3a103b73fd225a6aa3acf647648c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 14:12:44 +0100 Subject: [PATCH] Add doc strings for all channels, small update to names/semantics --- examples/combine_docs.ipynb | 26 +++--- examples/recursive-web-loader.py | 4 +- permchain/channels/__init__.py | 6 +- permchain/channels/{stream.py => archive.py} | 97 +++++++++++--------- permchain/channels/base.py | 21 ++++- permchain/channels/binop.py | 5 +- permchain/channels/context.py | 12 +++ permchain/channels/inbox.py | 74 +++++++-------- permchain/channels/last_value.py | 7 +- tests/test_channels.py | 4 +- 10 files changed, 149 insertions(+), 107 deletions(-) rename permchain/channels/{stream.py => archive.py} (72%) diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index cf4415f73..ffd417a0f 100644 --- a/examples/combine_docs.ipynb +++ b/examples/combine_docs.ipynb @@ -25,7 +25,7 @@ "from langchain.schema.document import Document\n", "from langchain.schema import format_document\n", "\n", - "from permchain import Pregel, PregelRead, channels\n" + "from permchain import Pregel, PregelRead, channels" ] }, { @@ -45,7 +45,7 @@ "metadata": {}, "outputs": [], "source": [ - "from langchain.schema.runnable import RunnableLambda\n" + "from langchain.schema.runnable import RunnableLambda" ] }, { @@ -59,7 +59,7 @@ "\n", "_combine_documents = RunnableLambda(\n", " lambda x: format_document(x, DEFAULT_DOCUMENT_PROMPT)\n", - ").map() | (lambda x: \"\\n\\n\".join(x))\n" + ").map() | (lambda x: \"\\n\\n\".join(x))" ] }, { @@ -72,7 +72,7 @@ "docs = [\n", " Document(page_content=\"Harrison used to work at Kensho\"),\n", " Document(page_content=\"Ankush worked at Facebook\"),\n", - "]\n" + "]" ] }, { @@ -98,7 +98,7 @@ " )\n", " | ChatOpenAI()\n", " | StrOutputParser()\n", - ")\n" + ")" ] }, { @@ -119,7 +119,7 @@ } ], "source": [ - "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})\n" + "stuff_chain.invoke({\"question\": \"where did harrison work\", \"docs\": docs})" ] }, { @@ -139,7 +139,7 @@ "metadata": {}, "outputs": [], "source": [ - "many_docs = docs * 5\n" + "many_docs = docs * 5" ] }, { @@ -164,7 +164,7 @@ " new_result_doc_list.append(_sub_result_docs[:-1])\n", " _sub_result_docs = _sub_result_docs[-1:]\n", " new_result_doc_list.append(_sub_result_docs)\n", - " return new_result_doc_list\n" + " return new_result_doc_list" ] }, { @@ -196,7 +196,7 @@ "source": [ "# Just to show what its like split\n", "split_docs = _split_list_of_docs(many_docs)\n", - "split_docs\n" + "split_docs" ] }, { @@ -214,7 +214,7 @@ " \"docs_to_finalize\": channels.Inbox(Document),\n", " # output\n", " \"answer\": channels.LastValue(str),\n", - "}\n" + "}" ] }, { @@ -249,7 +249,7 @@ " Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n", " | stuff_chain\n", " | Pregel.write_to(\"answer\")\n", - ")\n" + ")" ] }, { @@ -268,7 +268,7 @@ " input=[\"question\", \"docs\"],\n", " output=\"answer\",\n", " debug=True,\n", - ")\n" + ")" ] }, { @@ -332,7 +332,7 @@ } ], "source": [ - "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})\n" + "reduce_chain.invoke({\"question\": \"where did harrison work\", \"docs\": many_docs})" ] }, { diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index 3ca220fbc..a29e66e84 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -86,8 +86,8 @@ def recursive_web_loader( channels = { "base_url": Channels.LastValue(str), "next_urls": Channels.UniqueInbox(str), - "documents": Channels.Stream(Document), - "visited": Channels.Set(str), + "documents": Channels.Archive(Document), + "visited": Channels.UniqueArchive(str), "client": Channels.ContextManager(httpx_client, httpx_aclient), } # the main chain that gets executed recursively diff --git a/permchain/channels/__init__.py b/permchain/channels/__init__.py index a92933e03..4574034d5 100644 --- a/permchain/channels/__init__.py +++ b/permchain/channels/__init__.py @@ -1,15 +1,15 @@ +from permchain.channels.archive import Archive, UniqueArchive from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.context import ContextManager from permchain.channels.inbox import Inbox, UniqueInbox from permchain.channels.last_value import LastValue -from permchain.channels.stream import Set, Stream __all__ = [ "LastValue", "Inbox", "UniqueInbox", + "Archive", + "UniqueArchive", "BinaryOperatorAggregate", - "Set", - "Stream", "ContextManager", ] diff --git a/permchain/channels/stream.py b/permchain/channels/archive.py similarity index 72% rename from permchain/channels/stream.py rename to permchain/channels/archive.py index d51b02b8d..05015fbc4 100644 --- a/permchain/channels/stream.py +++ b/permchain/channels/archive.py @@ -5,50 +5,11 @@ from typing import Any, FrozenSet, Generator, Generic, Optional, Sequence, Type from typing_extensions import Self from permchain.channels.base import Channel, EmptyChannelError, Value +from permchain.channels.inbox import flatten -class Set(Generic[Value], Channel[FrozenSet[Value], Value]): - """Stores all unique values received.""" - - def __init__(self, typ: Type[Value]) -> None: - self.typ = typ - self.set = set[Value]() - - @property - def ValueType(self) -> Type[FrozenSet[Value]]: - """The type of the value stored in the channel.""" - return FrozenSet[self.typ] # type: ignore[name-defined] - - @property - def UpdateType(self) -> Type[Value]: - """The type of the update received by the channel.""" - return self.typ - - @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) - if checkpoint is not None: - empty.set = set(json.loads(checkpoint)) - try: - yield empty - finally: - pass - - def update(self, values: Sequence[Value]) -> None: - self.set.update(values) - - def get(self) -> FrozenSet[Value]: - try: - return frozenset(self.set) - except AttributeError: - raise EmptyChannelError() - - def checkpoint(self) -> str: - return json.dumps(list(self.set)) - - -class Stream(Generic[Value], Channel[Sequence[Value], Value]): - """Stores all unique values received.""" +class Archive(Generic[Value], Channel[Sequence[Value], Value | list[Value]]): + """Stores all unique values received, persists across steps.""" def __init__(self, typ: Type[Value]) -> None: self.typ = typ @@ -74,8 +35,8 @@ class Stream(Generic[Value], Channel[Sequence[Value], Value]): finally: pass - def update(self, values: Sequence[Value]) -> None: - self.set.extend(values) + def update(self, values: Sequence[Value | list[Value]]) -> None: + self.set.extend(flatten(values)) def get(self) -> Sequence[Value]: try: @@ -84,4 +45,50 @@ class Stream(Generic[Value], Channel[Sequence[Value], Value]): raise EmptyChannelError() def checkpoint(self) -> str: - return json.dumps(self.set) + try: + return json.dumps(self.set) + except AttributeError: + raise EmptyChannelError() + + +class UniqueArchive(Generic[Value], Channel[FrozenSet[Value], Value]): + """Stores all unique values received, persists across steps.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + self.set = set[Value]() + + @property + def ValueType(self) -> Type[FrozenSet[Value]]: + """The type of the value stored in the channel.""" + return FrozenSet[self.typ] # type: ignore[name-defined] + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.set = set(json.loads(checkpoint)) + try: + yield empty + finally: + pass + + def update(self, values: Sequence[Value | list[Value]]) -> None: + self.set.update(flatten(values)) + + def get(self) -> FrozenSet[Value]: + try: + return frozenset(self.set) + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> str: + try: + return json.dumps(list(self.set)) + except AttributeError: + raise EmptyChannelError() diff --git a/permchain/channels/base.py b/permchain/channels/base.py index ee6748445..789d7e2af 100644 --- a/permchain/channels/base.py +++ b/permchain/channels/base.py @@ -18,10 +18,15 @@ Update = TypeVar("Update") class EmptyChannelError(Exception): + """Raised when attempting to get the value of a channel that hasn't been updated + for the first time yet.""" + pass class InvalidUpdateError(Exception): + """Raised when attempting to update a channel with an invalid sequence of updates.""" + pass @@ -51,21 +56,30 @@ class Channel(Generic[Value, Update], ABC): @abstractmethod def update(self, values: Sequence[Update]) -> None: - ... + """Update the channel's value with the given sequence of updates. + The order of the updates in the sequence is arbitrary. + + Raises InvalidUpdateError if the sequence of updates is invalid.""" @abstractmethod def get(self) -> Value: - ... + """Return the current value of the channel. + + Raises EmptyChannelError if the channel is empty (never updated yet).""" @abstractmethod def checkpoint(self) -> str | None: - ... + """Return a string representation of the channel's current state, + or None if the channel doesn't support checkpoints. + + Raises EmptyChannelError if the channel is empty (never updated yet).""" @contextmanager def ChannelsManager( channels: Mapping[str, Channel] ) -> Generator[Mapping[str, Channel], None, None]: + """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" empty = {k: v.empty() for k, v in channels.items()} try: yield {k: v.__enter__() for k, v in empty.items()} @@ -78,6 +92,7 @@ def ChannelsManager( async def AsyncChannelsManager( channels: Mapping[str, Channel] ) -> AsyncGenerator[Mapping[str, Channel], None]: + """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" empty = {k: v.aempty() for k, v in channels.items()} try: yield {k: await v.__aenter__() for k, v in empty.items()} diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py index 2e703f048..9861c85e5 100644 --- a/permchain/channels/binop.py +++ b/permchain/channels/binop.py @@ -59,4 +59,7 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]): raise EmptyChannelError() def checkpoint(self) -> str: - return json.dumps(self.value) + try: + return json.dumps(self.value) + except AttributeError: + raise EmptyChannelError() diff --git a/permchain/channels/context.py b/permchain/channels/context.py index 1248dbc78..87664aacb 100644 --- a/permchain/channels/context.py +++ b/permchain/channels/context.py @@ -23,6 +23,18 @@ from permchain.channels.base import ( class ContextManager(Generic[Value], Channel[Value, None]): + """Exposes the value of a context manager, for the duration of an invocation. + Context manager is entered before the first step, and exited after the last step. + Optionally, provide an equivalent async context manager, which will be used + instead for async invocations. + + ```python + import httpx + + client = ContextManager(httpx.Client, httpx.AsyncClient) + ``` + """ + value: Value def __init__( diff --git a/permchain/channels/inbox.py b/permchain/channels/inbox.py index afe9c90df..f49ae8887 100644 --- a/permchain/channels/inbox.py +++ b/permchain/channels/inbox.py @@ -1,17 +1,31 @@ import json from contextlib import contextmanager -from typing import Any, Generator, Generic, Optional, Sequence, Type, Union, cast +from typing import ( + Any, + FrozenSet, + Generator, + Generic, + Iterator, + Optional, + Sequence, + Type, + Union, +) from typing_extensions import Self -from permchain.channels.base import ( - Channel, - EmptyChannelError, - Value, -) +from permchain.channels.base import Channel, EmptyChannelError, Value -class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): +def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]: + for value in values: + if isinstance(value, list): + yield from value + else: + yield value + + +class Inbox(Generic[Value], Channel[Sequence[Value], Value | list[Value]]): """Stores all values received, resets in each step.""" def __init__(self, typ: Type[Value]) -> None: @@ -40,16 +54,8 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): except AttributeError: pass - def update(self, values: Sequence[Value | Sequence[Value]]) -> None: - self.queue = tuple( - cast(Value, v) - for value in values - for v in ( - (value,) - if isinstance(value, self.typ) - else cast(Sequence[Value], value) - ) - ) + def update(self, values: Sequence[Value | list[Value]]) -> None: + self.queue = tuple(flatten(values)) def get(self) -> Sequence[Value]: try: @@ -58,19 +64,22 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): raise EmptyChannelError() def checkpoint(self) -> str: - return json.dumps(self.queue) + try: + return json.dumps(self.queue) + except AttributeError: + raise EmptyChannelError() -class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]): +class UniqueInbox(Generic[Value], Channel[FrozenSet[Value], Value | list[Value]]): """Stores all unique values received, resets in each step.""" def __init__(self, typ: Type[Value]) -> None: self.typ = typ @property - def ValueType(self) -> Type[Sequence[Value]]: + def ValueType(self) -> Type[FrozenSet[Value]]: """The type of the value stored in the channel.""" - return Sequence[self.typ] # type: ignore[name-defined] + return FrozenSet[self.typ] # type: ignore[name-defined] @property def UpdateType(self) -> Any: @@ -81,7 +90,7 @@ class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Valu def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: empty = self.__class__(self.typ) if checkpoint is not None: - empty.queue = tuple(json.loads(checkpoint)) + empty.queue = frozenset(json.loads(checkpoint)) try: yield empty finally: @@ -90,24 +99,17 @@ class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Valu except AttributeError: pass - def update(self, values: Sequence[Value | Sequence[Value]]) -> None: - self.queue = tuple( - set( - cast(Value, v) - for value in values - for v in ( - (value,) - if isinstance(value, self.typ) - else cast(Sequence[Value], value) - ) - ) - ) + def update(self, values: Sequence[Value | list[Value]]) -> None: + self.queue = frozenset(flatten(values)) - def get(self) -> Sequence[Value]: + def get(self) -> FrozenSet[Value]: try: return self.queue except AttributeError: raise EmptyChannelError() def checkpoint(self) -> str: - return json.dumps(self.queue) + try: + return json.dumps(self.queue) + except AttributeError: + raise EmptyChannelError() diff --git a/permchain/channels/last_value.py b/permchain/channels/last_value.py index 266e44ce0..d60b064d6 100644 --- a/permchain/channels/last_value.py +++ b/permchain/channels/last_value.py @@ -13,7 +13,7 @@ from permchain.channels.base import ( class LastValue(Generic[Value], Channel[Value, Value]): - """Stores the last value received.""" + """Stores the last value received, can receive at most one value per step.""" def __init__(self, typ: Type[Value]) -> None: self.typ = typ @@ -54,4 +54,7 @@ class LastValue(Generic[Value], Channel[Value, Value]): raise EmptyChannelError() def checkpoint(self) -> str: - return json.dumps(self.value) + try: + return json.dumps(self.value) + except AttributeError: + raise EmptyChannelError() diff --git a/tests/test_channels.py b/tests/test_channels.py index 4e27f1a0b..640520fde 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -72,7 +72,7 @@ async def test_inbox_async() -> None: def test_set() -> None: - with Channels.Set(str).empty() as channel: + with Channels.UniqueArchive(str).empty() as channel: assert channel.ValueType is FrozenSet[str] assert channel.UpdateType is str @@ -84,7 +84,7 @@ def test_set() -> None: async def test_set_async() -> None: - async with Channels.Set(str).aempty() as channel: + async with Channels.UniqueArchive(str).aempty() as channel: assert channel.ValueType is FrozenSet[str] assert channel.UpdateType is str