diff --git a/README.md b/README.md index ab533757e..0e00491ca 100644 --- a/README.md +++ b/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 diff --git a/examples/example.ipynb b/examples/example.ipynb index 9c514a1f3..223fd90ee 100644 --- a/examples/example.ipynb +++ b/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", + "]" ] } ], diff --git a/permchain/__init__.py b/permchain/__init__.py index 845641df9..770967f30 100644 --- a/permchain/__init__.py +++ b/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", diff --git a/permchain/connection.py b/permchain/connection.py index fd99d67c9..d988e5862 100644 --- a/permchain/connection.py +++ b/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]: + ... diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py index 4aefd8d0f..0762a3d89 100644 --- a/permchain/connection_inmemory.py +++ b/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: diff --git a/permchain/pubsub.py b/permchain/pubsub.py index 33d952561..123cbb171 100644 --- a/permchain/pubsub.py +++ b/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) diff --git a/tests/test_invoke.py b/tests/test_invoke.py index f397f22ab..d0f48c29d 100644 --- a/tests/test_invoke.py +++ b/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