From c37f7e2c2ab49f5edee15a5dc3e829301d21452f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Sep 2023 17:38:11 +0100 Subject: [PATCH] Add Topic.reduce() --- README.md | 5 ++- permchain/connection.py | 4 +- permchain/connection_inmemory.py | 8 +++- permchain/pubsub.py | 64 ++++++++++++++++++++++++----- permchain/topic.py | 34 ++++++++++++++++ tests/test_invoke.py | 69 ++++++++++++++++++++++++++++++++ 6 files changed, 169 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 13c088005..ea0925f18 100644 --- a/README.md +++ b/README.md @@ -29,13 +29,14 @@ Check `tests` and `examples` for more examples. - [x] Add Connection.peek() to monitor past messages from all topics - [x] Enable resuming PubSub from the "middle" of the computation - [x] Add test for .peek() +- [x] Add "wait until topic X is done" pattern, aka. `Topic.reduce()` - [ ] Move tracking of inflight processes/messages to Connection - [ ] Use this to build retry mechanism, where any inflight messages are moved back to the respective topics when restarting + - [ ] But this would require being able to replay a message for a single listener only, which maybe requires a larger redesign of PubSub<>Connection contract than what I wanted to do here - [ ] Detect cycles (aka. infinite loops) and throw an error - [ ] Allow user to catch that error (by subcribing to an error topic?) -- [ ] Add "human in the loop" pattern, one of the two below +- [ ] Add example for "human in the loop" pattern, one of the two below - [ ] Example with one permchain, which runs until it produces either 1. request for input or 2. output. The consumer code then gets the needed info, and restarts the permchain with answer, and same state id - [ ] Allow interrupting execution by breaking out of the iterator returned by .stream() - [ ] Build example showing a simple "human in the loop" pattern using this, ie. if a certain message asking for input is published the consumer of the iterator breaks out, does something and then restarts it -- [ ] Add "wait until topic X is done" pattern, aka. `Topic.reduce()` - [ ] Add Redis-backed Connection implementation diff --git a/permchain/connection.py b/permchain/connection.py index 4f960e635..ad96edd96 100644 --- a/permchain/connection.py +++ b/permchain/connection.py @@ -17,7 +17,7 @@ class PubSubConnection(ABC): return f"{prefix}:{topic_name}" @abstractmethod - def iterate(self, prefix: str, topic_name: str) -> Iterator[Any]: + def iterate(self, prefix: str, topic_name: str, wait: bool) -> Iterator[Any]: """Iterate over all currently queued messages for a topic, consuming them.""" ... @@ -56,4 +56,6 @@ class PubSubConnection(ABC): @abstractmethod def peek(self, prefix: str) -> Iterator[LogMessage]: + """Iterate over all previously published messages for all topics, + without consuming them.""" ... diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py index 0fdf7f86c..d18d2577d 100644 --- a/permchain/connection_inmemory.py +++ b/permchain/connection_inmemory.py @@ -37,14 +37,20 @@ class InMemoryPubSubConnection(PubSubConnection): def peek(self, prefix: str) -> Iterator[LogMessage]: return iter(self.logs[str(prefix)]) - def iterate(self, prefix: str, topic_name: str) -> Iterator[Any]: + def iterate(self, prefix: str, topic_name: str, wait: bool) -> Iterator[Any]: topic = self.full_topic_name(prefix, topic_name) + + # This connection doesn't support iterating over topics with listeners connected with self.lock: if self.listeners[topic]: raise RuntimeError( f"Cannot iterate over topic {topic} while listeners are connected" ) + # If wait is False, add sentinel to queue to ensure the iterator terminates + if not wait: + self.topics[topic].close() + return iter(self.topics[topic]) def listen( diff --git a/permchain/pubsub.py b/permchain/pubsub.py index 933e14599..0d9408283 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -7,27 +7,41 @@ from itertools import groupby from typing import Any, Iterator, List, Optional, Sequence, Set, TypeVar from langchain.callbacks.manager import CallbackManagerForChainRun -from langchain.load.serializable import Serializable from langchain.schema.runnable import Runnable, RunnableConfig, patch_config from langchain.schema.runnable.base import Runnable from langchain.schema.runnable.config import get_executor_for_config from permchain.connection import PubSubConnection from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY -from permchain.topic import INPUT_TOPIC, OUTPUT_TOPIC, RunnableSubscriber +from permchain.topic import ( + INPUT_TOPIC, + OUTPUT_TOPIC, + RunnableReducer, + RunnableSubscriber, +) T = TypeVar("T") T_in = TypeVar("T_in") T_out = TypeVar("T_out") -class PubSub(Serializable, Runnable[Any, Any], ABC): - processes: Sequence[RunnableSubscriber[Any]] +class PubSub(Runnable[Any, Any], ABC): + processes: Sequence[RunnableSubscriber[Any] | RunnableReducer[Any]] connection: PubSubConnection + def __init__( + self, + processes: Sequence[RunnableSubscriber[Any] | RunnableReducer[Any]], + connection: PubSubConnection, + ) -> None: + super().__init__() + self.processes = processes + self.connection = connection + class Config: arbitrary_types_allowed = True + revalidate_instances = False def with_retry(self, **kwargs: Any) -> Runnable[Any, Any]: return self.__class__( @@ -41,6 +55,16 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): run_manager: CallbackManagerForChainRun, config: RunnableConfig, ) -> Iterator[Any]: + subscribers: list[RunnableSubscriber[Any]] = [] + reducers: list[RunnableReducer[Any]] = [] + for process in self.processes: + if isinstance(process, RunnableReducer): + reducers.append(process) + elif isinstance(process, RunnableSubscriber): + subscribers.append(process) + else: + raise ValueError(f"Unknown process type: {process}") + # Consume input iterator into a single value input_value = None for chunk in input: @@ -57,6 +81,25 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # Track exceptions exceptions: List[Exception] = [] + def on_idle(): + if reducers: + reducers_by_topic = groupby( + sorted(reducers, key=lambda p: p.topic.name), + lambda p: p.topic.name, + ) + for topic_name, processes in reducers_by_topic: + messages = list( + self.connection.iterate( + topic_prefix, topic_name, wait=False + ) + ) + if messages: + for process in processes: + run_once(process, messages) + + if not inflight: + self.connection.disconnect(topic_prefix) + def run_once(process: RunnableSubscriber[Any], value: Any) -> None: """Run a process once.""" @@ -77,7 +120,7 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # - all processes are done, or # - an exception occurred if not inflight or exc is not None: - self.connection.disconnect(topic_prefix) + on_idle() def get(topic_name: str) -> Any: if topic_name == INPUT_TOPIC: @@ -110,7 +153,7 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # Listen on all subscribed topics processes_by_topic = groupby( - sorted(self.processes, key=lambda p: p.topic.name), + sorted(subscribers, key=lambda p: p.topic.name), lambda p: p.topic.name, ) for topic_name, processes in processes_by_topic: @@ -128,10 +171,12 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # Yield output until all processes are done # This blocks the current thread, all other work needs to go # through the executor - for chunk in self.connection.iterate(topic_prefix, OUTPUT_TOPIC): + for chunk in self.connection.iterate( + topic_prefix, OUTPUT_TOPIC, wait=True + ): yield chunk else: - self.connection.disconnect(topic_prefix) + on_idle() finally: # Cancel all inflight futures while inflight: @@ -156,6 +201,3 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): for chunk in self.stream(input, config): collected.append(chunk) return collected - - -PubSub.update_forward_refs() diff --git a/permchain/topic.py b/permchain/topic.py index 8b8d7f861..4a98fbc5c 100644 --- a/permchain/topic.py +++ b/permchain/topic.py @@ -21,6 +21,7 @@ from langchain.schema.runnable import ( RunnableSequence, ) from langchain.schema.runnable.base import Other, coerce_to_runnable +from langchain.schema.runnable.config import RunnableConfig from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY @@ -45,6 +46,12 @@ class Topic(Serializable, Generic[T], ABC): return RunnableSubscriber(topic=self) + def reduce(self) -> RunnableReducer[T]: + if self.name == OUTPUT_TOPIC: + raise ValueError("Cannot reduce on output topic") + + return RunnableReducer(topic=self) + def current(self) -> RunnableCurrentValue[T]: if self.name == OUTPUT_TOPIC: raise ValueError("Cannot subscribe to output topic") @@ -106,6 +113,33 @@ class RunnableSubscriber(RunnableBinding[T, Any]): raise NotImplementedError() +class RunnableReducer(RunnableBinding[list[T], Any]): + topic: Topic[T] + + bound: Runnable[list[T], Any] = Field(default_factory=RunnablePassthrough) + + kwargs: Mapping[str, Any] = Field(default_factory=dict) + + def __or__( + self, + other: Runnable[Any, Other] + | Callable[[Any], Other] + | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], + ) -> RunnableSequence[list[T], Other]: + if isinstance(self.bound, RunnablePassthrough): + return RunnableReducer(topic=self.topic, bound=coerce_to_runnable(other)) + else: + return RunnableReducer(topic=self.topic, bound=self.bound | other) + + def __ror__( + self, + other: Runnable[Other, Any] + | Callable[[Any], Other] + | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], + ) -> RunnableSequence[Other, Any]: + raise NotImplementedError() + + class RunnablePublisher(Serializable, Runnable[T, T]): topic: Topic[T] diff --git a/tests/test_invoke.py b/tests/test_invoke.py index 83713bd2b..9cd2af9a7 100644 --- a/tests/test_invoke.py +++ b/tests/test_invoke.py @@ -169,6 +169,75 @@ def test_invoke_two_processes_two_in_two_out(mocker: MockerFixture): assert conn.listeners == {} +def test_invoke_two_processes_two_in_reduce_two_out(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) + topic_one = Topic("one") + topic_two = Topic("two") + chain_one = Topic.IN.subscribe() | add_one | topic_one.publish() + chain_two = topic_one.subscribe() | add_one | topic_two.publish() + chain_three = Topic.IN.subscribe() | add_one | topic_two.publish() + chain_four = topic_two.reduce() | add_10_each | Topic.OUT.publish() + + # Chains can be invoked directly for testing + assert chain_one.invoke(2) == 3 + assert chain_four.invoke([2, 3]) == [12, 13] + + conn = InMemoryPubSubConnection() + pubsub = PubSub((chain_one, chain_two, chain_three, chain_four), connection=conn) + + # Using in-memory conn internals to make assertions about pubsub + # If we start with 0 listeners + assert conn.listeners == {} + + # Then invoke pubsub + # We get a single array result as chain_four waits for all publishers to finish + # before operating on all elements published to topic_two as an array + assert pubsub.invoke(2) == [[13, 14]] + + # After invoke returns the listeners were cleaned up + assert conn.listeners == {} + + +def test_invoke_reduce_then_subscribe(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) + + topic_one = Topic("one") + topic_two = Topic("two") + + chain_one = Topic.IN.subscribe() | add_10_each | topic_one.publish_each() + chain_two = topic_one.reduce() | sum | topic_two.publish() + chain_three = topic_two.subscribe() | add_one | Topic.OUT.publish() + + # Chains can be invoked directly for testing + assert chain_two.invoke([2, 3]) == 5 + assert chain_three.invoke(5) == 6 + + state_id = uuid4() + conn = InMemoryPubSubConnection(clear_on_disconnect=False) + pubsub = PubSub((chain_one, chain_two, chain_three), connection=conn) + + # Using in-memory conn internals to make assertions about pubsub + # If we start with 0 listeners + assert conn.listeners == {} + + # Then invoke pubsub + # We get a single array result as chain_four waits for all publishers to finish + # before operating on all elements published to topic_two as an array + assert pubsub.invoke([2, 3], {"state_id": state_id}) == [26] + assert [{**m, "started_at": None} for m in conn.peek(state_id)] == [ + {"message": [2, 3], "topic_name": "__in__", "started_at": None}, + {"message": 12, "topic_name": "one", "started_at": None}, + {"message": 13, "topic_name": "one", "started_at": None}, + {"message": 25, "topic_name": "two", "started_at": None}, + {"message": 26, "topic_name": "__out__", "started_at": None}, + ] + + # After invoke returns the listeners were cleaned up + assert conn.listeners == {} + + def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) topic_one = Topic("one")