Chat listen() and iterate() to use PubSubMessage wrapper instead of raw value

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

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:

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
         )

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: