Move prefix management to Connection, Add Connection.peek(), Add test for resuming ongoing execution

README.md

@@ -26,11 +26,12 @@ Check `tests` and `examples` for more examples.
 - [x] Add initial retry support (pending changes in `langchain`)
 - [x] Implement OUT as regular topic
 - [x] Implement IN as regular topic
+- [x] Add Connection.peek() to monitor past messages from all topics
+- [ ] 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
 - [ ] Detect cycles (aka. infinite loops) and throw an error
   - [ ] Allow user to catch that error (by subcribing to an error topic?)
-- [ ] Replace Queue data structure with a Log data structure (this will enable checking the status of the readers, etc.)
-  - [ ] eg. https://anyio.readthedocs.io/en/3.x/streams.html
-- [ ] Enable resuming PubSub from the "middle" of the computation
+- [x] Enable resuming PubSub from the "middle" of the computation
 - [ ] Add "human in the loop" pattern
-- [ ] Add "wait until topic X is done" pattern
+- [ ] Add "wait until topic X is done" pattern, aka. `Topic.reduce()`
 - [ ] Add Redis-backed Connection implementation

examples/example.ipynb

@@ -349,7 +349,14 @@
     }
    ],
    "source": [
-    "[*web_researcher.batch([{\"question\": \"What food do turtles eat?\"}, {\"question\": \"Where do bears live?\"}])]"
+    "[\n",
+    "    *web_researcher.batch(\n",
+    "        [\n",
+    "            {\"question\": \"What food do turtles eat?\"},\n",
+    "            {\"question\": \"Where do bears live?\"},\n",
+    "        ]\n",
+    "    )\n",
+    "]"
    ]
   }
  ],

permchain/__init__.py

@@ -1,6 +1,6 @@
+from permchain.connection_inmemory import InMemoryPubSubConnection
 from permchain.pubsub import PubSub
 from permchain.topic import Topic
-from permchain.connection_inmemory import InMemoryPubSubConnection
 
 __all__ = [
     "PubSub",

permchain/connection.py

@@ -1,38 +1,57 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Iterator
+from typing import Any, Callable, Iterator, TypedDict
 
 PubSubListener = Callable[[Any], None]
 
 
+class LogMessage(TypedDict):
+    message: Any
+    topic_name: str
+    started_at: str
+
+
 class PubSubConnection(ABC):
-    @abstractmethod
-    def iterate(self, topic_name: str) -> Iterator[Any]:
-        ...
+    def full_topic_name(self, prefix: str, topic_name: str) -> str:
+        return f"{prefix}:{topic_name}"
 
     @abstractmethod
-    def listen(self, topic_name: str, listener: PubSubListener) -> None:
+    def iterate(self, prefix: str, topic_name: str) -> Iterator[Any]:
         ...
 
-    async def alisten(self, topic_name: str, listener: PubSubListener) -> None:
+    # TODO add aiterate() method
+
+    @abstractmethod
+    def listen(
+        self, prefix: str, topic_name: str, listeners: list[PubSubListener]
+    ) -> None:
+        ...
+
+    async def alisten(
+        self, prefix: str, topic_name: str, listeners: list[PubSubListener]
+    ) -> None:
         return await asyncio.get_event_loop().run_in_executor(
-            None, self.listen, topic_name, listener
+            None, self.listen, prefix, topic_name, listeners
         )
 
     @abstractmethod
-    def send(self, topic_name: str, message: Any) -> None:
+    def send(self, prefix: str, topic_name: str, message: Any) -> None:
         ...
 
-    async def asend(self, topic_name: str, message: Any) -> None:
+    async def asend(self, prefix: str, topic_name: str, message: Any) -> None:
         return await asyncio.get_event_loop().run_in_executor(
-            None, self.send, topic_name, message
+            None, self.send, prefix, topic_name, message
         )
 
     @abstractmethod
-    def disconnect(self, topic_name: str) -> None:
+    def disconnect(self, prefix: str) -> None:
         ...
 
-    async def adisconnect(self, topic_name: str) -> None:
+    async def adisconnect(self, prefix: str) -> None:
         return await asyncio.get_event_loop().run_in_executor(
-            None, self.disconnect, topic_name
+            None, self.disconnect, prefix
         )
+
+    @abstractmethod
+    def peek(self, prefix: str) -> Iterator[LogMessage]:
+        ...

permchain/connection_inmemory.py

@@ -1,11 +1,10 @@
-import threading
 import queue
+import threading
 from collections import defaultdict
+from datetime import datetime
 from typing import Any, Iterator
 
-q = queue.Queue()
-
-from permchain.connection import PubSubConnection, PubSubListener
+from permchain.connection import LogMessage, PubSubConnection, PubSubListener
 
 
 class IterableQueue(queue.SimpleQueue):
@@ -22,58 +21,90 @@ class IterableQueue(queue.SimpleQueue):
 
 
 class InMemoryPubSubConnection(PubSubConnection):
+    clear_on_disconnect: bool
+    logs: defaultdict[str, list[Any]]
     topics: defaultdict[str, IterableQueue]
     listeners: defaultdict[str, list[PubSubListener]]
     lock: threading.RLock
 
-    def __init__(self) -> None:
+    def __init__(self, clear_on_disconnect: bool = True) -> None:
+        self.clear_on_disconnect = clear_on_disconnect
+        self.logs = defaultdict(list)
         self.topics = defaultdict(IterableQueue)
         self.listeners = defaultdict(list)
         self.lock = threading.RLock()
 
-    def iterate(self, topic_name: str) -> Iterator[Any]:
+    def peek(self, prefix: str) -> Iterator[LogMessage]:
+        return iter(self.logs[prefix])
+
+    def iterate(self, prefix: str, topic_name: str) -> Iterator[Any]:
+        topic = self.full_topic_name(prefix, topic_name)
         with self.lock:
-            if self.listeners[topic_name]:
+            if self.listeners[topic]:
                 raise RuntimeError(
-                    f"Cannot iterate over topic {topic_name} while listeners are connected"
+                    f"Cannot iterate over topic {topic} while listeners are connected"
                 )
 
-        return iter(self.topics[topic_name])
+        return iter(self.topics[topic])
 
-    def listen(self, topic_name: str, listeners: list[PubSubListener]) -> None:
-        self.disconnect(topic_name)
+    def listen(
+        self, prefix: str, topic_name: str, listeners: list[PubSubListener]
+    ) -> None:
+        topic = self.full_topic_name(prefix, topic_name)
+        self.disconnect(topic)
 
         with self.lock:
-            self.listeners[topic_name].extend(listeners)
-            topic_queue = self.topics[topic_name]
+            # Add the listeners for future messages
+            self.listeners[topic].extend(listeners)
+
+            # Send any pending messages to the listeners
+            topic_queue = self.topics[topic]
             while not topic_queue.empty():
                 message = topic_queue.get()
-                for listener in self.listeners[topic_name]:
-                    listener(message)
+                if message is not topic_queue.done_sentinel:
+                    for listener in self.listeners[topic]:
+                        listener(message)
+
+    def send(self, prefix: str, topic_name: str, message: Any) -> None:
+        topic = self.full_topic_name(prefix, topic_name)
 
-    def send(self, topic_name: str, message: Any) -> None:
         with self.lock:
-            listeners = self.listeners[topic_name]
+            # Add the message to the log
+            self.logs[prefix].append(
+                LogMessage(
+                    message=message,
+                    topic_name=topic_name,
+                    started_at=datetime.now().isoformat(),
+                )
+            )
+            listeners = self.listeners[topic]
             if listeners:
+                # Send the message to listeners if any are connected
                 for listener in listeners:
                     listener(message)
             else:
-                self.topics[topic_name].put(message)
+                # Otherwise add the message to the topic queue for later
+                self.topics[topic].put(message)
 
-    def disconnect(self, prefix: str) -> None:
+    def disconnect(self, prefix_or_topic: str) -> None:
         with self.lock:
+            if self.clear_on_disconnect:
+                if prefix_or_topic in self.logs:
+                    del self.logs[prefix_or_topic]
+
             to_delete = []
             for topic, queue in self.topics.items():
-                if topic.startswith(prefix):
+                if topic.startswith(prefix_or_topic):
                     queue.close()
-                    to_delete.append(topic)
+                    if self.clear_on_disconnect:
+                        to_delete.append(topic)
             # can't delete while iterating
             for topic in to_delete:
                 del self.topics[topic]
 
             to_delete = []
             for topic in self.listeners:
-                if topic.startswith(prefix):
+                if topic.startswith(prefix_or_topic):
                     to_delete.append(topic)
             # can't delete while iterating
             for topic in to_delete:

permchain/pubsub.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from itertools import groupby
 from abc import ABC
 from concurrent.futures import CancelledError, Future
 from functools import partial
+from itertools import groupby
 from typing import Any, Iterator, List, Optional, Sequence, Set, TypeVar
 
 from langchain.callbacks.manager import CallbackManagerForChainRun
@@ -50,20 +50,13 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
                 input_value += chunk
 
         with get_executor_for_config(config) as executor:
-            topic_prefix = str(run_manager.run_id) + ":"
+            # Namespace topics for each run
+            topic_prefix = str(run_manager.parent_run_id or run_manager.run_id)
             # Track inflight futures
             inflight: Set[Future] = set()
             # Track exceptions
             exceptions: List[Exception] = []
 
-            # Namespace topics for each run
-            def prefix_topic_name(topic_name: str) -> str:
-                return f"{topic_prefix}{topic_name}"
-
-            def send(topic_name: str, message: Any) -> None:
-                """Send a message to a topic. Injected into config."""
-                self.connection.send(prefix_topic_name(topic_name), message)
-
             def run_once(process: RunnableSubscriber[Any], value: Any) -> None:
                 """Run a process once."""
 
@@ -106,7 +99,7 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
                             callbacks=run_manager.get_child(),
                             run_name=f"Topic: {process.topic.name}",
                         ),
-                        CONFIG_SEND_KEY: send,
+                        CONFIG_SEND_KEY: partial(self.connection.send, topic_prefix),
                         CONFIG_GET_KEY: get,
                     },
                 )
@@ -122,19 +115,18 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
             )
             for topic_name, processes in listeners_by_topic:
                 self.connection.listen(
-                    prefix_topic_name(topic_name),
+                    topic_prefix,
+                    topic_name,
                     [partial(run_once, process) for process in processes],
                 )
 
             # Send input to input processes
-            send(INPUT_TOPIC, input_value)
+            self.connection.send(topic_prefix, INPUT_TOPIC, input_value)
 
             try:
                 if inflight:
                     # Yield output until all processes are done
-                    for chunk in self.connection.iterate(
-                        prefix_topic_name(OUTPUT_TOPIC)
-                    ):
+                    for chunk in self.connection.iterate(topic_prefix, OUTPUT_TOPIC):
                         yield chunk
                 else:
                     self.connection.disconnect(topic_prefix)

tests/test_invoke.py

@@ -1,8 +1,9 @@
 import pytest
 from pytest_mock import MockerFixture
+
+from langchain.callbacks.manager import trace_as_chain_group
 from permchain.connection_inmemory import InMemoryPubSubConnection
 from permchain.pubsub import PubSub
-
 from permchain.topic import RunnableSubscriber, Topic
 
 
@@ -47,6 +48,32 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture):
     assert conn.listeners == {}
 
 
+def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
+    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection(clear_on_disconnect=False)
+    pubsub_one = PubSub(processes=(chain_one,), connection=conn)
+    pubsub_two = PubSub(processes=(chain_two,), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+    # Then invoke both pubsubs, as a group
+    # The second picks up where the first left off
+    with trace_as_chain_group("PubSubGroup") as cm:
+        assert pubsub_one.invoke(2, {"callbacks": cm}) == []
+        assert pubsub_two.invoke(None, {"callbacks": cm}) == [4]
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
 def test_invoke_many_processes_in_out(mocker: MockerFixture):
     test_size = 100