From f0e31a8ffa41734391b1c408fb51ce5a1138cf46 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 22 Oct 2023 19:07:10 +0100 Subject: [PATCH] Split into smaller files --- README.md | 1 + permchain/channels/base.py | 25 ++ permchain/{pregel.py => pregel/__init__.py} | 314 +------------------- permchain/pregel/constants.py | 3 + permchain/pregel/read.py | 152 ++++++++++ permchain/pregel/validate.py | 44 +++ permchain/pregel/write.py | 82 +++++ 7 files changed, 322 insertions(+), 299 deletions(-) rename permchain/{pregel.py => pregel/__init__.py} (64%) create mode 100644 permchain/pregel/constants.py create mode 100644 permchain/pregel/read.py create mode 100644 permchain/pregel/validate.py create mode 100644 permchain/pregel/write.py diff --git a/README.md b/README.md index 278e04042..464592a85 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ Check `examples` for more examples. - [x] Test different input and output types (str, str sequence) - [x] Add tests for Stream, UniqueInbox - [ ] Add tests for subscribe_to_each().join() +- [ ] Add optional debug logging - [ ] Implement checkpointing - [ ] Save checkpoints at end of each step - [ ] Load checkpoint at start of invocation diff --git a/permchain/channels/base.py b/permchain/channels/base.py index a7e62cd43..ee6748445 100644 --- a/permchain/channels/base.py +++ b/permchain/channels/base.py @@ -5,6 +5,7 @@ from typing import ( AsyncGenerator, Generator, Generic, + Mapping, Optional, Sequence, TypeVar, @@ -59,3 +60,27 @@ class Channel(Generic[Value, Update], ABC): @abstractmethod def checkpoint(self) -> str | None: ... + + +@contextmanager +def ChannelsManager( + channels: Mapping[str, Channel] +) -> Generator[Mapping[str, Channel], None, None]: + empty = {k: v.empty() for k, v in channels.items()} + try: + yield {k: v.__enter__() for k, v in empty.items()} + finally: + for v in empty.values(): + v.__exit__(None, None, None) + + +@asynccontextmanager +async def AsyncChannelsManager( + channels: Mapping[str, Channel] +) -> AsyncGenerator[Mapping[str, Channel], None]: + empty = {k: v.aempty() for k, v in channels.items()} + try: + yield {k: await v.__aenter__() for k, v in empty.items()} + finally: + for v in empty.values(): + await v.__aexit__(None, None, None) diff --git a/permchain/pregel.py b/permchain/pregel/__init__.py similarity index 64% rename from permchain/pregel.py rename to permchain/pregel/__init__.py index 88384955c..02e04318e 100644 --- a/permchain/pregel.py +++ b/permchain/pregel/__init__.py @@ -4,283 +4,39 @@ import asyncio import concurrent.futures import logging from collections import defaultdict, deque -from contextlib import asynccontextmanager, contextmanager -from typing import ( - Any, - AsyncGenerator, - AsyncIterator, - Callable, - Generator, - Iterator, - Mapping, - Optional, - Sequence, - Type, - cast, -) +from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast from langchain.callbacks.manager import ( AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, ) -from langchain.pydantic_v1 import BaseModel, Field, create_model +from langchain.pydantic_v1 import BaseModel, create_model from langchain.schema.runnable import ( Runnable, - RunnableBinding, - RunnableLambda, RunnablePassthrough, RunnableSerializable, ) -from langchain.schema.runnable.base import ( - Other, - RunnableEach, - RunnableLike, - coerce_to_runnable, -) +from langchain.schema.runnable.base import RunnableLike, coerce_to_runnable from langchain.schema.runnable.config import ( RunnableConfig, get_executor_for_config, patch_config, ) -from langchain.schema.runnable.utils import ConfigurableFieldSpec -from permchain.channels.base import Channel, EmptyChannelError +from permchain.channels.base import ( + AsyncChannelsManager, + Channel, + ChannelsManager, + EmptyChannelError, +) +from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP +from permchain.pregel.read import PregelBatch, PregelInvoke +from permchain.pregel.validate import validate_chains_channels +from permchain.pregel.write import PregelSink logger = logging.getLogger(__name__) -CONFIG_KEY_STEP = "__pregel_step" -CONFIG_KEY_SEND = "__pregel_send" -CONFIG_KEY_READ = "__pregel_read" - -TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] - - -class PregelRead(RunnableLambda): - channel: str - - @property - def config_specs(self) -> Sequence[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_READ, - name=CONFIG_KEY_READ, - description=None, - default=None, - annotation=Callable[[Channel], Any], - ), - ] - - def __init__(self, channel: str) -> None: - super().__init__(func=self._read, afunc=self._aread) # type: ignore[arg-type] - self.channel = channel - - def _read(self, _: Any, config: RunnableConfig) -> Any: - try: - read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] - except KeyError: - raise RuntimeError( - f"Runnable {self} is not configured with a read function" - "Make sure to call in the context of a Pregel process" - ) - return read(self.channel) - - async def _aread(self, _: Any, config: RunnableConfig) -> Any: - try: - read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] - except KeyError: - raise RuntimeError( - f"Runnable {self} is not configured with a read function" - "Make sure to call in the context of a Pregel process" - ) - return read(self.channel) - - -class PregelInvoke(RunnableBinding): - channels: Mapping[None, str] | Mapping[str, str] - - bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) - - kwargs: Mapping[str, Any] = Field(default_factory=dict) - - def __init__( - self, - channels: Mapping[None, str] | Mapping[str, str], - *, - bound: Optional[Runnable[Any, Any]] = None, - kwargs: Optional[Mapping[str, Any]] = None, - config: Optional[RunnableConfig] = None, - **other_kwargs: Any, - ) -> None: - super().__init__( - channels=channels, - bound=bound or RunnablePassthrough(), - kwargs=kwargs or {}, - config=config, - **other_kwargs, - ) - - def __or__( - self, - other: Runnable[Any, Other] - | Callable[[Any], Other] - | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], - ) -> PregelInvoke: - if isinstance(self.bound, RunnablePassthrough): - return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other)) - else: - # delegate to __or__ in self.bound - return PregelInvoke(channels=self.channels, bound=self.bound | other) - - def __ror__( - self, - other: Runnable[Other, Any] - | Callable[[Any], Other] - | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], - ) -> Runnable: - raise NotImplementedError() - - -class PregelBatch(RunnableEach): - channel: str - - key: Optional[str] - - bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) - - def join(self, channels: Sequence[str]) -> PregelBatch: - if self.key is None: - raise ValueError( - "Cannot join() additional channels without a key." - " Pass a key arg to Pregel.subscribe_to_each()." - ) - - joiner = RunnablePassthrough.assign( - **{chan: PregelRead(chan) for chan in channels} - ) - if isinstance(self.bound, RunnablePassthrough): - return PregelBatch(channel=self.channel, key=self.key, bound=joiner) - else: - return PregelBatch( - channel=self.channel, key=self.key, bound=self.bound | joiner - ) - - def __or__( # type: ignore[override] - self, - other: Runnable[Any, Other] - | Callable[[Any], Other] - | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], - ) -> PregelBatch: - if isinstance(self.bound, RunnablePassthrough): - return PregelBatch( - channel=self.channel, key=self.key, bound=coerce_to_runnable(other) - ) - else: - # delegate to __or__ in self.bound - return PregelBatch( - channel=self.channel, key=self.key, bound=self.bound | other - ) - - def __ror__( - self, - other: Runnable[Other, Any] - | Callable[[Any], Other] - | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], - ) -> Runnable: - raise NotImplementedError() - - -class PregelSink(RunnableLambda): - channels: Sequence[tuple[str, Runnable]] - """ - Mapping of write channels to Runnables that return the value to be written, - or None to skip writing. - """ - - max_steps: Optional[int] - - def __init__( - self, - *, - channels: Sequence[tuple[str, Runnable]], - max_steps: Optional[int] = None, - ): - super().__init__(func=self._write, afunc=self._awrite) # type: ignore[arg-type] - self.channels = channels - self.max_steps = max_steps - - @property - def config_specs(self) -> Sequence[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_STEP, - name=CONFIG_KEY_STEP, - description=None, - default=None, - annotation=int, - ), - ConfigurableFieldSpec( - id=CONFIG_KEY_SEND, - name=CONFIG_KEY_SEND, - description=None, - default=None, - annotation=TYPE_SEND, - ), - ] - - def _write(self, input: Any, config: RunnableConfig) -> None: - step: int = config["configurable"][CONFIG_KEY_STEP] - - if self.max_steps is not None and step >= self.max_steps: - return - - write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - - values = [(chan, r.invoke(input, config)) for chan, r in self.channels] - - write([(chan, val) for chan, val in values if val is not None]) - - return input - - async def _awrite(self, input: Any, config: RunnableConfig) -> None: - step: int = config["configurable"][CONFIG_KEY_STEP] - - if self.max_steps is not None and step >= self.max_steps: - return - - write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - - values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels] - - write([(chan, val) for chan, val in values if val is not None]) - - return input - - -@contextmanager -def ChannelsManager( - channels: Mapping[str, Channel] -) -> Generator[Mapping[str, Channel], None, None]: - empty = {k: v.empty() for k, v in channels.items()} - try: - yield {k: v.__enter__() for k, v in empty.items()} - finally: - for v in empty.values(): - v.__exit__(None, None, None) - - -@asynccontextmanager -async def AsyncChannelsManager( - channels: Mapping[str, Channel] -) -> AsyncGenerator[Mapping[str, Channel], None]: - empty = {k: v.aempty() for k, v in channels.items()} - try: - yield {k: await v.__aenter__() for k, v in empty.items()} - finally: - for v in empty.values(): - await v.__aexit__(None, None, None) - - class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): channels: Mapping[str, Channel] @@ -428,7 +184,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]() - # execute tasks, and wait for one to fail or all to finish + # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks done, inflight = concurrent.futures.wait( ( @@ -528,7 +284,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]() - # execute tasks, and wait for one to fail or all to finish + # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks done, inflight = await asyncio.wait( [ @@ -704,43 +460,3 @@ def _apply_writes_and_prepare_next_tasks( tasks.append((proc, val)) return tasks - - -def validate_chains_channels( - chains: Sequence[PregelInvoke | PregelBatch], - channels: Mapping[str, Channel], - input: str | Sequence[str], - output: str | Sequence[str], -) -> None: - subscribed_channels = set[str]() - for chain in chains: - if isinstance(chain, PregelInvoke): - subscribed_channels.update(chain.channels.values()) - elif isinstance(chain, PregelBatch): - subscribed_channels.add(chain.channel) - else: - raise TypeError( - f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()" - ) - - for chan in subscribed_channels: - if chan not in channels: - raise ValueError(f"Channel {chan} is subscribed to, but not initialized") - - if isinstance(input, str): - if input not in subscribed_channels: - raise ValueError(f"Input channel {input} is not subscribed to by any chain") - else: - for chan in input: - if chan not in subscribed_channels: - raise ValueError( - f"Input channel {chan} is not subscribed to by any chain" - ) - - if isinstance(output, str): - if output not in channels: - raise ValueError(f"Output channel {output} is not initialized") - else: - for chan in output: - if chan not in channels: - raise ValueError(f"Output channel {chan} is not initialized") diff --git a/permchain/pregel/constants.py b/permchain/pregel/constants.py new file mode 100644 index 000000000..eda362e78 --- /dev/null +++ b/permchain/pregel/constants.py @@ -0,0 +1,3 @@ +CONFIG_KEY_STEP = "__pregel_step" +CONFIG_KEY_SEND = "__pregel_send" +CONFIG_KEY_READ = "__pregel_read" diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py new file mode 100644 index 000000000..d6d5eadbd --- /dev/null +++ b/permchain/pregel/read.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from typing import Any, Callable, Mapping, Optional, Sequence + +from langchain.pydantic_v1 import Field +from langchain.schema.runnable import ( + Runnable, + RunnableBinding, + RunnableConfig, + RunnableLambda, + RunnablePassthrough, +) +from langchain.schema.runnable.base import Other, RunnableEach, coerce_to_runnable +from langchain.schema.runnable.utils import ConfigurableFieldSpec + +from permchain.channels.base import Channel +from permchain.pregel.constants import CONFIG_KEY_READ + + +class PregelRead(RunnableLambda): + channel: str + + @property + def config_specs(self) -> Sequence[ConfigurableFieldSpec]: + return [ + ConfigurableFieldSpec( + id=CONFIG_KEY_READ, + name=CONFIG_KEY_READ, + description=None, + default=None, + annotation=Callable[[Channel], Any], + ), + ] + + def __init__(self, channel: str) -> None: + # TODO remove type ignore after updating langchain + super().__init__(func=self._read, afunc=self._aread) # type: ignore[arg-type] + self.channel = channel + + def _read(self, _: Any, config: RunnableConfig) -> Any: + try: + read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] + except KeyError: + raise RuntimeError( + f"Runnable {self} is not configured with a read function" + "Make sure to call in the context of a Pregel process" + ) + return read(self.channel) + + async def _aread(self, _: Any, config: RunnableConfig) -> Any: + try: + read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] + except KeyError: + raise RuntimeError( + f"Runnable {self} is not configured with a read function" + "Make sure to call in the context of a Pregel process" + ) + return read(self.channel) + + +class PregelInvoke(RunnableBinding): + channels: Mapping[None, str] | Mapping[str, str] + + bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) + + kwargs: Mapping[str, Any] = Field(default_factory=dict) + + def __init__( + self, + channels: Mapping[None, str] | Mapping[str, str], + *, + bound: Optional[Runnable[Any, Any]] = None, + kwargs: Optional[Mapping[str, Any]] = None, + config: Optional[RunnableConfig] = None, + **other_kwargs: Any, + ) -> None: + super().__init__( + channels=channels, + bound=bound or RunnablePassthrough(), + kwargs=kwargs or {}, + config=config, + **other_kwargs, + ) + + def __or__( + self, + other: Runnable[Any, Other] + | Callable[[Any], Other] + | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], + ) -> PregelInvoke: + if isinstance(self.bound, RunnablePassthrough): + return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other)) + else: + # delegate to __or__ in self.bound + return PregelInvoke(channels=self.channels, bound=self.bound | other) + + def __ror__( + self, + other: Runnable[Other, Any] + | Callable[[Any], Other] + | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], + ) -> Runnable: + raise NotImplementedError() + + +class PregelBatch(RunnableEach): + channel: str + + key: Optional[str] + + bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) + + def join(self, channels: Sequence[str]) -> PregelBatch: + if self.key is None: + raise ValueError( + "Cannot join() additional channels without a key." + " Pass a key arg to Pregel.subscribe_to_each()." + ) + + joiner = RunnablePassthrough.assign( + **{chan: PregelRead(chan) for chan in channels} + ) + if isinstance(self.bound, RunnablePassthrough): + return PregelBatch(channel=self.channel, key=self.key, bound=joiner) + else: + return PregelBatch( + channel=self.channel, key=self.key, bound=self.bound | joiner + ) + + def __or__( # type: ignore[override] + self, + other: Runnable[Any, Other] + | Callable[[Any], Other] + | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], + ) -> PregelBatch: + if isinstance(self.bound, RunnablePassthrough): + return PregelBatch( + channel=self.channel, key=self.key, bound=coerce_to_runnable(other) + ) + else: + # delegate to __or__ in self.bound + return PregelBatch( + channel=self.channel, key=self.key, bound=self.bound | other + ) + + def __ror__( + self, + other: Runnable[Other, Any] + | Callable[[Any], Other] + | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], + ) -> Runnable: + raise NotImplementedError() diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py new file mode 100644 index 000000000..dd5d06249 --- /dev/null +++ b/permchain/pregel/validate.py @@ -0,0 +1,44 @@ +from typing import Mapping, Sequence + +from permchain.channels.base import Channel +from permchain.pregel.read import PregelBatch, PregelInvoke + + +def validate_chains_channels( + chains: Sequence[PregelInvoke | PregelBatch], + channels: Mapping[str, Channel], + input: str | Sequence[str], + output: str | Sequence[str], +) -> None: + subscribed_channels = set[str]() + for chain in chains: + if isinstance(chain, PregelInvoke): + subscribed_channels.update(chain.channels.values()) + elif isinstance(chain, PregelBatch): + subscribed_channels.add(chain.channel) + else: + raise TypeError( + f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()" + ) + + for chan in subscribed_channels: + if chan not in channels: + raise ValueError(f"Channel {chan} is subscribed to, but not initialized") + + if isinstance(input, str): + if input not in subscribed_channels: + raise ValueError(f"Input channel {input} is not subscribed to by any chain") + else: + for chan in input: + if chan not in subscribed_channels: + raise ValueError( + f"Input channel {chan} is not subscribed to by any chain" + ) + + if isinstance(output, str): + if output not in channels: + raise ValueError(f"Output channel {output} is not initialized") + else: + for chan in output: + if chan not in channels: + raise ValueError(f"Output channel {chan} is not initialized") diff --git a/permchain/pregel/write.py b/permchain/pregel/write.py new file mode 100644 index 000000000..e04d93f2b --- /dev/null +++ b/permchain/pregel/write.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from typing import Any, Callable, Optional, Sequence + +from langchain.schema.runnable import ( + Runnable, + RunnableConfig, + RunnableLambda, +) +from langchain.schema.runnable.utils import ConfigurableFieldSpec + +from permchain.pregel.constants import CONFIG_KEY_SEND, CONFIG_KEY_STEP + +TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] + + +# TODO switch to RunnablePassthrough after updating langchain +class PregelSink(RunnableLambda): + channels: Sequence[tuple[str, Runnable]] + """ + Mapping of write channels to Runnables that return the value to be written, + or None to skip writing. + """ + + max_steps: Optional[int] + + def __init__( + self, + *, + channels: Sequence[tuple[str, Runnable]], + max_steps: Optional[int] = None, + ): + super().__init__(func=self._write, afunc=self._awrite) # type: ignore[arg-type] + self.channels = channels + self.max_steps = max_steps + + @property + def config_specs(self) -> Sequence[ConfigurableFieldSpec]: + return [ + ConfigurableFieldSpec( + id=CONFIG_KEY_STEP, + name=CONFIG_KEY_STEP, + description=None, + default=None, + annotation=int, + ), + ConfigurableFieldSpec( + id=CONFIG_KEY_SEND, + name=CONFIG_KEY_SEND, + description=None, + default=None, + annotation=TYPE_SEND, + ), + ] + + def _write(self, input: Any, config: RunnableConfig) -> None: + step: int = config["configurable"][CONFIG_KEY_STEP] + + if self.max_steps is not None and step >= self.max_steps: + return + + write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] + + values = [(chan, r.invoke(input, config)) for chan, r in self.channels] + + write([(chan, val) for chan, val in values if val is not None]) + + return input + + async def _awrite(self, input: Any, config: RunnableConfig) -> None: + step: int = config["configurable"][CONFIG_KEY_STEP] + + if self.max_steps is not None and step >= self.max_steps: + return + + write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] + + values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels] + + write([(chan, val) for chan, val in values if val is not None]) + + return input