Add Topic.reduce()

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

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."""
         ...

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(

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()

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]
 

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")