diff --git a/examples/combine_docs.ipynb b/examples/combine_docs.ipynb index cf4415f73..4657ed08c 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 Channels, Pregel, PregelRead\n" ] }, { @@ -206,14 +206,14 @@ "metadata": {}, "outputs": [], "source": [ - "chans = {\n", + "channels = {\n", " # input\n", - " \"question\": channels.LastValue(str),\n", - " \"docs\": channels.Inbox(Document),\n", + " \"question\": Channels.LastValue(str),\n", + " \"docs\": Channels.Inbox(Document),\n", " # intermediate\n", - " \"docs_to_finalize\": channels.Inbox(Document),\n", + " \"docs_to_finalize\": Channels.Inbox(Document),\n", " # output\n", - " \"answer\": channels.LastValue(str),\n", + " \"answer\": Channels.LastValue(str),\n", "}\n" ] }, @@ -264,7 +264,7 @@ " \"collapse\": collapse,\n", " \"finalize\": finalize,\n", " },\n", - " channels=chans,\n", + " channels=channels,\n", " input=[\"question\", \"docs\"],\n", " output=\"answer\",\n", " debug=True,\n", 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