diff --git a/README.md b/README.md index 594e76deb..cf4a8750f 100644 --- a/README.md +++ b/README.md @@ -9,26 +9,32 @@ ```python from permchain import Pregel, channels -value = channels.LastValue[str]("value") - grow_value = ( - Pregel.subscribe_to(value) + Pregel.subscribe_to("value") | (lambda x: x + x) - | Pregel.send_to({value: lambda x: x if len(x) < 10 else None}) + | Pregel.send_to(value=lambda x: x if len(x) < 10 else None) ) -pubsub = Pregel(grow_value, input=value, output=value) +pubsub = Pregel( + grow_value, + channels={"value": channels.LastValue[str]()}, + input="value", + output="value", +) assert pubsub.invoke("a") == "aaaaaaaa" + ``` Check `examples` for more examples. ## Near-term Roadmap -- [ ] Iterate on API - - [ ] do we want api to receive output from multiple channels in invoke() - - [ ] do we want api to send input to multiple channels in invoke() +- [x] Iterate on API + - [x] do we want api to receive output from multiple channels in invoke() + - [x] do we want api to send input to multiple channels in invoke() + - [ ] Finish updating tests to new API +- [ ] Implement input_schema and output_schema in Pregel - [ ] Implement checkpointing - [ ] Save checkpoints at end of each step - [ ] Load checkpoint at start of invocation diff --git a/examples/readme.py b/examples/readme.py index e90794a38..3bb54429b 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -1,14 +1,16 @@ -import operator from permchain import Pregel, channels -value = channels.LastValue[str]("value") - grow_value = ( - Pregel.subscribe_to(value) + Pregel.subscribe_to("value") | (lambda x: x + x) - | Pregel.send_to({value: lambda x: x if len(x) < 10 else None}) + | Pregel.send_to(value=lambda x: x if len(x) < 10 else None) ) -pubsub = Pregel(grow_value, input=value, output=value) +pubsub = Pregel( + grow_value, + channels={"value": channels.LastValue[str]()}, + input="value", + output="value", +) assert pubsub.invoke("a") == "aaaaaaaa" diff --git a/examples/runnable-pregel.py b/examples/runnable-pregel.py index aa59deb88..a0c60f4c4 100644 --- a/examples/runnable-pregel.py +++ b/examples/runnable-pregel.py @@ -74,53 +74,61 @@ editor_chain = ( | JsonOutputFunctionsParser(args_only=False) ) -# channels - -question = channels.LastValue[str]("question") - -draft = channels.LastValue[str]("draft") - -notes = channels.LastValue[str]("notes") - # application -drafter_node = ( - Pregel.subscribe_to(question=question) | drafter_chain | Pregel.send_to(draft) +drafter = ( + # subscribe to question channel as a dict with a single key, "question" + Pregel.subscribe_to(["question"]) + | drafter_chain + | Pregel.send_to("draft") ) -editor_node = ( - Pregel.subscribe_to(draft=draft) +editor = ( + # subscribe to draft channel as a dict with a single key, "draft" + Pregel.subscribe_to(["draft"]) | editor_chain | Pregel.send_to( - {notes: lambda x: x["arguments"]["notes"] if x["name"] == "revise" else None} + # send to "notes" channel if the editor does not accept the draft + notes=lambda x: x["arguments"]["notes"] + if x["name"] == "revise" + else None ) ) -reviser_node = ( - Pregel.subscribe_to(notes=notes).join(question=question, draft=draft) +reviser = ( + # subscribe to new values of "notes" channel, + # and join them with the current values of "question" and "draft" + Pregel.subscribe_to(["notes"]).join("question", "draft") | reviser_chain - | Pregel.send_to(draft) + | Pregel.send_to("draft") ) draft_revise_loop = Pregel( - drafter_node, - reviser_node, - editor_node, - input=question, - output=draft, + [drafter, reviser, editor], + channels={ + "question": channels.LastValue[str](), + "draft": channels.LastValue[str](), + "notes": channels.LastValue[str](), + }, + # output will be a dict with keys "draft" and "notes" + output=["draft", "notes"], + # input can be a dict with any of the channels as keys + input=None, ) # run -# for draft in draft_revise_loop.stream("What food do turtles eat?"): -# print('Draft: "' + draft + '"') -# print("---") +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("What food do turtles eat?"): - print('Draft: "' + draft + '"') + async for draft in draft_revise_loop.astream( + {"question": "What food do turtles eat?"} + ): + print(draft) print("---") -asyncio.run(main()) +# asyncio.run(main()) diff --git a/permchain/channels.py b/permchain/channels.py index dd4680d30..0d6a89fe6 100644 --- a/permchain/channels.py +++ b/permchain/channels.py @@ -2,7 +2,7 @@ import json from abc import ABC, abstractmethod from typing import Callable, FrozenSet, Generic, Optional, Sequence, TypeVar -from typing_extensions import Self +from typing_extensions import Self, get_args Value = TypeVar("Value") Update = TypeVar("Update") @@ -17,14 +17,20 @@ class InvalidUpdateError(Exception): class Channel(Generic[Value, Update], ABC): - def __init__(self, name: str) -> None: - self.name = name + # TODO: add type hints for ValueType and UpdateType + # @property + # def ValueType(self) -> type[Value]: + # """The type of the value stored in the channel.""" + # type_args = get_args(self.__class__.__orig_bases__[-1]) # type: ignore[attr-defined] + # if type_args and len(type_args) == 2: + # return type_args[0] - def __repr__(self) -> str: - return f"{self.__class__.__name__}({self.name})" - - def __str__(self) -> str: - return self.name + # @property + # def UpdateType(self) -> type[Update]: + # """The type of the update received by the channel.""" + # type_args = get_args(self.__class__.__orig_bases__[-1]) # type: ignore[attr-defined] + # if type_args and len(type_args) == 2: + # return type_args[1] @abstractmethod def _empty(self, checkpoint: Optional[str] = None) -> Self: @@ -53,12 +59,12 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]): ``` """ - def __init__(self, name: str, operator: Callable[[Value, Value], Value]): - super().__init__(name) + def __init__(self, operator: Callable[[Value, Value], Value]): + super().__init__() self.operator = operator def _empty(self, checkpoint: Optional[str] = None) -> Self: - empty = self.__class__(self.name, self.operator) + empty = self.__class__(self.operator) if checkpoint is not None: empty.value = json.loads(checkpoint) return empty @@ -85,7 +91,7 @@ class LastValue(Generic[Value], Channel[Value, Value]): """Stores the last value received.""" def _empty(self, checkpoint: Optional[str] = None) -> Self: - empty = self.__class__(self.name) + empty = self.__class__() if checkpoint is not None: empty.value = json.loads(checkpoint) return empty @@ -110,7 +116,7 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value]): """Stores all values received, resets in each step.""" def _empty(self, checkpoint: Optional[str] = None) -> Self: - empty = self.__class__(self.name) + empty = self.__class__() if checkpoint is not None: empty.queue = tuple(json.loads(checkpoint)) return empty @@ -134,7 +140,7 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]): set: set[Value] def _empty(self, checkpoint: Optional[str] = None) -> Self: - empty = self.__class__(self.name) + empty = self.__class__() if checkpoint is not None: empty.set = set(json.loads(checkpoint)) return empty diff --git a/permchain/pregel.py b/permchain/pregel.py index 2c1569751..10f42532f 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -13,6 +13,7 @@ from typing import ( Mapping, Optional, Sequence, + cast, overload, ) @@ -50,11 +51,11 @@ CONFIG_KEY_STEP = "__pregel_step" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" -TYPE_SEND = Callable[[Sequence[tuple[Channel, Any]]], None] +TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] class PregelRead(RunnableLambda): - channel: Channel + channel: str @property def config_specs(self) -> Sequence[ConfigurableFieldSpec]: @@ -68,13 +69,13 @@ class PregelRead(RunnableLambda): ), ] - def __init__(self, channel: Channel) -> None: + 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[[Channel], Any] = config["configurable"][CONFIG_KEY_READ] + read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] except KeyError: raise RuntimeError( f"Runnable {self} is not configured with a read function" @@ -84,7 +85,7 @@ class PregelRead(RunnableLambda): async def _aread(self, _: Any, config: RunnableConfig) -> Any: try: - read: Callable[[Channel], Any] = config["configurable"][CONFIG_KEY_READ] + read: Callable[[str], Any] = config["configurable"][CONFIG_KEY_READ] except KeyError: raise RuntimeError( f"Runnable {self} is not configured with a read function" @@ -94,15 +95,15 @@ class PregelRead(RunnableLambda): class PregelInvoke(RunnableBinding): - channels: Mapping[None, Channel] | Mapping[str, Channel] + channels: Mapping[None, str] | Mapping[str, str] bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) kwargs: Mapping[str, Any] = Field(default_factory=dict) - def join(self, **channels: Channel) -> PregelInvoke: + def join(self, *channels: str) -> PregelInvoke: joiner = RunnablePassthrough.assign( - **{k: PregelRead(chan) for k, chan in channels.items()} + **{chan: PregelRead(chan) for chan in channels} ) if isinstance(self.bound, RunnablePassthrough): return PregelInvoke(channels=self.channels, bound=joiner) @@ -130,7 +131,7 @@ class PregelInvoke(RunnableBinding): class PregelBatch(RunnableEach): - channel: Inbox + channel: str bound: Runnable[Any, Any] = Field(default_factory=RunnablePassthrough) @@ -155,7 +156,7 @@ class PregelBatch(RunnableEach): class PregelSink(RunnableLambda): - channels: Sequence[tuple[Channel, Runnable]] + channels: Sequence[tuple[str, Runnable]] """ Mapping of write channels to Runnables that return the value to be written, or None to skip writing. @@ -166,7 +167,7 @@ class PregelSink(RunnableLambda): def __init__( self, *, - channels: Sequence[tuple[Channel, Runnable]], + channels: Sequence[tuple[str, Runnable]], max_steps: Optional[int] = None, ): super().__init__(func=self._write, afunc=self._awrite) @@ -221,12 +222,14 @@ class PregelSink(RunnableLambda): return input -class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): - input: Channel[Any, Input] +class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]): + channels: Mapping[str, Channel] - output: Channel[Output, Any] + chains: Sequence[PregelInvoke | PregelBatch] - processes: Sequence[PregelInvoke | PregelBatch] + output: str | Sequence[str] + + input: str | None step_timeout: Optional[float] = None @@ -235,42 +238,34 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): def __init__( self, - *processes: PregelInvoke | PregelBatch, - input: Channel[Input, Any], - output: Channel[Output, Any], + *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch, + channels: Mapping[str, Channel], + output: str | Sequence[str], + input: str | None = None, step_timeout: Optional[float] = None, - **kwargs: Any, - ): + ) -> None: + chains_flat: list[PregelInvoke | PregelBatch] = [] + for chain in chains: + if isinstance(chain, (list, tuple)): + chains_flat.extend(chain) + else: + chains_flat.append(chain) super().__init__( - processes=processes, - input=input, + chains=chains_flat, + channels=channels, output=output, + input=input, step_timeout=step_timeout, - **kwargs, ) - @overload @classmethod - def subscribe_to(cls, __channel: Channel) -> PregelInvoke: - ... - - @overload - @classmethod - def subscribe_to( - cls, __channel: Mapping[str, Channel] | None = None, **kwargs: Channel - ) -> PregelInvoke: - ... - - @classmethod - def subscribe_to( - cls, __channel: Channel | Mapping[str, Channel] | None = None, **kwargs: Channel - ) -> PregelInvoke: - """Runs process.invoke() each time channels are updated.""" - __channel = __channel or {} - return ( - PregelInvoke(channels={None: __channel}) - if isinstance(__channel, Channel) - else PregelInvoke(channels={**__channel, **kwargs}) + def subscribe_to(cls, channels: str | Sequence[str]) -> PregelInvoke: + """Runs process.invoke() each time channels are updated, + with a dict of the channel values as input.""" + return PregelInvoke( + channels={None: channels} + if isinstance(channels, str) + else {chan: chan for chan in channels} ) @classmethod @@ -281,54 +276,34 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): @classmethod def send_to( cls, - channels: Channel | Mapping[Channel, RunnableLike], - *, - max_steps: Optional[int] = None, + *channels: str, + _max_steps: Optional[int] = None, + **kwargs: RunnableLike, ) -> PregelSink: """Writes to channels the result of the lambda, or None to skip writing.""" return PregelSink( channels=( - [(channels, RunnablePassthrough())] - if isinstance(channels, Channel) - else [(k, coerce_to_runnable(v)) for k, v in channels.items()] + [(c, RunnablePassthrough()) for c in channels] + + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()] ), - max_steps=max_steps, + max_steps=_max_steps, ) - def _prepare_channels(self) -> Mapping[Channel, Channel]: - channels: dict[Channel, Channel] = {self.output: self.output._empty()} - for proc in self.processes: - if isinstance(proc, PregelInvoke): - for chan in proc.channels.values(): - if chan not in channels: - channels[chan] = chan._empty() - elif isinstance(proc, PregelBatch): - if proc.channel not in channels: - channels[proc.channel] = proc.channel._empty() - else: - raise TypeError( - f"Received process {proc}, expected instance of PregelInvoke or PregelBatch" - ) - - if not channels: - raise ValueError("Found 0 channels for Pregel run") - - if self.input not in channels: - raise ValueError("Input channel not being read from") - - return channels - def _transform( self, - input: Iterator[Input], + input: Iterator[dict[str, Any] | Any], run_manager: CallbackManagerForChainRun, config: RunnableConfig, ) -> Iterator[Output]: - processes = tuple(self.processes) + processes = tuple(self.chains) # TODO this is where we'd restore from checkpoint - channels = self._prepare_channels() + channels = {k: v._empty() for k, v in self.channels.items()} next_tasks = _apply_writes_and_prepare_next_tasks( - processes, channels, deque((self.input, chunk) for chunk in input) + processes, + channels, + deque((self.input, chunk) for chunk in input) + if self.input is not None + else deque((k, v) for chunk in input for k, v in chunk.items()), ) def read(chan: Channel) -> Any: @@ -406,14 +381,18 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): async def _atransform( self, - input: AsyncIterator[Input], + input: AsyncIterator[dict[str, Any] | Any], run_manager: AsyncCallbackManagerForChainRun, config: RunnableConfig, ) -> AsyncIterator[Output]: - processes = tuple(self.processes) - channels = self._prepare_channels() + processes = tuple(self.chains) + channels = {k: v._empty() for k, v in self.channels.items()} next_tasks = _apply_writes_and_prepare_next_tasks( - processes, channels, [(self.input, chunk) async for chunk in input] + processes, + channels, + deque((self.input, chunk) async for chunk in input) + if self.input is not None + else deque((k, v) async for chunk in input for k, v in chunk.items()), ) def read(chan: Channel) -> Any: @@ -480,15 +459,22 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): ) # if any write to output channel in this step, yield current value - if any(chan is self.output for chan, _ in pending_writes): - yield channels[self.output]._get() + if isinstance(self.output, str): + if any(chan is self.output for chan, _ in pending_writes): + yield channels[self.output]._get() + else: + if updated := {c for c, _ in pending_writes if c in self.output}: + yield {chan: channels[chan]._get() for chan in updated} # if no more tasks, we're done if not next_tasks: break def invoke( - self, input: Input, config: RunnableConfig | None = None, **kwargs: Any + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + **kwargs: Any, ) -> Output: latest: Output | None = None for chunk in self.stream(input, config, **kwargs): @@ -496,13 +482,16 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): return latest def stream( - self, input: Input, config: RunnableConfig | None = None, **kwargs: Any + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + **kwargs: Any, ) -> Iterator[Output]: return self.transform(iter([input]), config, **kwargs) def transform( self, - input: Iterator[Input], + input: Iterator[dict[str, Any] | Any], config: RunnableConfig | None = None, **kwargs: Any | None, ) -> Iterator[Output]: @@ -511,7 +500,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): ) async def ainvoke( - self, input: Input, config: RunnableConfig | None = None, **kwargs: Any + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + **kwargs: Any, ) -> Output: latest: Output | None = None async for chunk in self.astream(input, config, **kwargs): @@ -519,7 +511,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): return latest async def astream( - self, input: Input, config: RunnableConfig | None = None, **kwargs: Any + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + **kwargs: Any, ) -> AsyncIterator[Output]: async def input_stream() -> AsyncIterator[Input]: yield input @@ -529,7 +524,7 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): async def atransform( self, - input: AsyncIterator[Input], + input: AsyncIterator[dict[str, Any] | Any], config: RunnableConfig | None = None, **kwargs: Any | None, ) -> AsyncIterator[Output]: @@ -541,10 +536,10 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): def _apply_writes_and_prepare_next_tasks( processes: Sequence[PregelInvoke | PregelBatch], - channels: Mapping[Channel, Channel], - pending_writes: Sequence[tuple[Channel, Any]], + channels: Mapping[str, Channel], + pending_writes: Sequence[tuple[str, Any]], ) -> list[tuple[Runnable, Any]]: - pending_writes_by_channel: dict[Channel, list[Any]] = defaultdict(list) + pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: pending_writes_by_channel[chan].append(val)