diff --git a/permchain/connection.py b/permchain/connection.py index 87228880c..9c8528816 100644 --- a/permchain/connection.py +++ b/permchain/connection.py @@ -2,30 +2,33 @@ import asyncio from abc import ABC, abstractmethod from typing import Any, Callable, Iterator, TypedDict -PubSubListener = Callable[[Any], None] - -class PubSubLog(TypedDict): +class PubSubMessage(TypedDict): topic: str value: Any published_at: str correlation_id: str +PubSubListener = Callable[[PubSubMessage], None] + + class PubSubConnection(ABC): def full_name(self, prefix: str, *parts: str) -> str: """Return the full topic name for a given prefix and topic name.""" return ":".join(map(str, [prefix, *parts])) @abstractmethod - def observe(self, prefix: str) -> Iterator[PubSubLog]: + def observe(self, prefix: str) -> Iterator[PubSubMessage]: """Iterate over messages for all topics under this prefix, without affecting listeners/iterators on each topic. This method waits for new messages to arrive.""" ... @abstractmethod - def iterate(self, prefix: str, topic: str, *, wait: bool) -> Iterator[Any]: + def iterate( + self, prefix: str, topic: str, *, wait: bool + ) -> Iterator[PubSubMessage]: """Iterate over all currently queued messages for a topic, consuming them. Optionally wait for new messages to arrive.""" ... @@ -44,12 +47,12 @@ class PubSubConnection(ABC): ) @abstractmethod - def send(self, prefix: str, topic: str, message: Any) -> None: + def send(self, prefix: str, topic: str, value: Any) -> None: ... - async def asend(self, prefix: str, topic: str, message: Any) -> None: + async def asend(self, prefix: str, topic: str, value: Any) -> None: return await asyncio.get_event_loop().run_in_executor( - None, self.send, prefix, topic, message + None, self.send, prefix, topic, value ) @abstractmethod diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py index d3b39525a..53b288ddf 100644 --- a/permchain/connection_inmemory.py +++ b/permchain/connection_inmemory.py @@ -2,18 +2,25 @@ import queue import threading from collections import defaultdict from datetime import datetime -from typing import Any, Iterator +from typing import Any, Iterator, cast -from permchain.connection import PubSubConnection, PubSubListener, PubSubLog +from permchain.connection import PubSubConnection, PubSubListener, PubSubMessage class IterableQueue(queue.SimpleQueue): done_sentinel = object() - def get(self, block: bool = True, timeout: float | None = None) -> Any: + def put( + self, item: PubSubMessage, block: bool = True, timeout: float | None = None + ) -> None: + return super().put(item, block, timeout) + + def get( + self, block: bool = True, timeout: float | None = None + ) -> PubSubMessage | object: return super().get(block=block, timeout=timeout) - def __iter__(self) -> Iterator[Any]: + def __iter__(self) -> Iterator[PubSubMessage]: return iter(self.get, self.done_sentinel) def close(self) -> None: @@ -34,10 +41,12 @@ class InMemoryPubSubConnection(PubSubConnection): self.listeners = defaultdict(list) self.lock = threading.RLock() - def observe(self, prefix: str) -> Iterator[PubSubLog]: + def observe(self, prefix: str) -> Iterator[PubSubMessage]: return iter(self.logs[str(prefix)]) - def iterate(self, prefix: str, topic: str, *, wait: bool) -> Iterator[Any]: + def iterate( + self, prefix: str, topic: str, *, wait: bool + ) -> Iterator[PubSubMessage]: topic = self.full_name(prefix, topic) # This connection doesn't support iterating over topics with listeners connected @@ -67,20 +76,19 @@ class InMemoryPubSubConnection(PubSubConnection): message = topic_queue.get() if message is not topic_queue.done_sentinel: for listener in self.listeners[full_name]: - listener(message) + listener(cast(PubSubMessage, message)) - def send(self, prefix: str, topic: str, message: Any) -> None: + def send(self, prefix: str, topic: str, value: Any) -> None: full_name = self.full_name(prefix, topic) + message = PubSubMessage( + value=value, + topic=topic, + correlation_id=str(prefix), + published_at=datetime.now().isoformat(), + ) # Add the message to the log - self.logs[str(prefix)].put( - PubSubLog( - value=message, - topic=topic, - correlation_id=str(prefix), - published_at=datetime.now().isoformat(), - ) - ) + self.logs[str(prefix)].put(message) with self.lock: listeners = self.listeners[full_name] if listeners: diff --git a/permchain/pubsub.py b/permchain/pubsub.py index 4bac44846..3769bafd5 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -10,7 +10,7 @@ from langchain.callbacks.manager import CallbackManagerForChainRun from langchain.schema.runnable import Runnable, RunnableConfig, patch_config from langchain.schema.runnable.config import get_executor_for_config -from permchain.connection import PubSubConnection, PubSubLog +from permchain.connection import PubSubConnection, PubSubMessage from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY from permchain.topic import ( INPUT_TOPIC, @@ -126,9 +126,15 @@ class PubSub(Runnable[Any, Any], ABC): on_idle() def run_once( - process: RunnableSubscriber[Any] | RunnableReducer[Any], value: Any + process: RunnableSubscriber[Any] | RunnableReducer[Any], + messages: PubSubMessage | list[PubSubMessage], ) -> None: """Run a process once.""" + value = ( + [m["value"] for m in messages] + if isinstance(messages, list) + else messages["value"] + ) def get(topic_name: str) -> Any: if topic_name == INPUT_TOPIC: @@ -201,7 +207,7 @@ class PubSub(Runnable[Any, Any], ABC): input: Any, config: Optional[RunnableConfig] = None, **kwargs: Optional[Any], - ) -> Iterator[PubSubLog]: + ) -> Iterator[PubSubMessage]: yield from self._transform_stream_with_config( iter([input]), self._transform, config, **kwargs ) diff --git a/tests/test_invoke.py b/tests/test_invoke.py index 14332a36c..914fdeb1d 100644 --- a/tests/test_invoke.py +++ b/tests/test_invoke.py @@ -4,15 +4,15 @@ from uuid import uuid4 import pytest from pytest_mock import MockerFixture -from permchain.connection import PubSubLog +from permchain.connection import PubSubMessage from permchain.connection_inmemory import InMemoryPubSubConnection from permchain.pubsub import PubSub from permchain.topic import RunnableSubscriber, Topic def clean_log( - logs: Iterator[PubSubLog], correlation_id: bool | None = None -) -> list[PubSubLog]: + logs: Iterator[PubSubMessage], correlation_id: bool | None = None +) -> list[PubSubMessage]: if correlation_id is False: return [{**m, "published_at": None, "correlation_id": None} for m in logs] else: