Improve Channel.publish() to support both streaming and non-streaming

Author: RussellLuo

coagent/agents/chat_agent.py

@@ -125,7 +125,7 @@ def __init__(self, host_agent: ChatAgent, agent_type: str):
 
     async def handle(self, msg: ChatHistory) -> AsyncIterator[ChatMessage]:
         addr = Address(name=self.agent_type, id=self.host_agent.address.id)
-        result = self.host_agent.channel.publish_multi(addr, msg.encode())
+        result = await self.host_agent.channel.publish(addr, msg.encode(), stream=True)
         full_content = ""
         async for chunk in result:
             resp = ChatMessage.decode(chunk)

coagent/cli/main.py

@@ -67,17 +67,18 @@ async def run(
     async with runtime:
         addr = Address(name=agent_type, id=session_id)
         try:
+            response = await runtime.channel.publish(
+                addr,
+                msg,
+                stream=stream,
+                request=True,
+                timeout=10,
+                probe=probe,
+            )
             if not stream:
-                response = await runtime.channel.publish(
-                    addr,
-                    msg,
-                    request=True,
-                    timeout=10,
-                    probe=probe,
-                )
                 print_msg(response, oneline, filter)
             else:
-                async for chunk in runtime.channel.publish_multi(addr, msg):
+                async for chunk in response:
                     print_msg(chunk, oneline, filter)
         except asyncio.CancelledError:
             await runtime.channel.cancel(addr)

coagent/core/runtime.py

@@ -1,3 +1,4 @@
+import abc
 import asyncio
 from typing import AsyncIterator
 
@@ -71,19 +72,63 @@ def channel(self) -> Channel:
 
 
 class BaseChannel(Channel):
-    async def publish_multi(
+    async def publish(
+        self,
+        addr: Address,
+        msg: RawMessage,
+        stream: bool = False,
+        request: bool = False,
+        reply: str = "",
+        timeout: float = 0.5,
+        probe: bool = True,
+    ) -> AsyncIterator[RawMessage] | RawMessage | None:
+        if stream:
+            return self._publish_stream(addr, msg, probe=probe)
+        else:
+            return await self._publish(
+                addr,
+                msg,
+                request=request,
+                stream=stream,
+                reply=reply,
+                timeout=timeout,
+                probe=probe,
+            )
+
+    @abc.abstractmethod
+    async def _publish(
+        self,
+        addr: Address,
+        msg: RawMessage,
+        request: bool = False,
+        stream: bool = False,
+        reply: str = "",
+        timeout: float = 0.5,
+        probe: bool = True,
+    ) -> RawMessage | None:
+        pass
+
+    async def _publish_stream(
         self,
         addr: Address,
         msg: RawMessage,
         probe: bool = True,
     ) -> AsyncIterator[RawMessage]:
-        """A default implementation that leverages the channel's own subscribe and publish methods."""
+        """Publish a message and wait for multiple reply messages.
+
+        Args:
+            addr (Address): The address of the agent.
+            msg (RawMessage): The raw message to send.
+            probe (bool, optional): Whether to probe the agent before sending the message. Defaults to True.
+
+        This is a default implementation that leverages the channel's own subscribe and _publish methods.
+        """
         queue: QueueSubscriptionIterator = QueueSubscriptionIterator()
 
         inbox = await self.new_reply_topic()
         sub = await self.subscribe(addr=Address(name=inbox), handler=queue.receive)
 
-        await self.publish(
+        await self._publish(
             addr,
             msg,
             request=True,

coagent/core/types.py

@@ -236,41 +236,25 @@ async def publish(
         self,
         addr: Address,
         msg: RawMessage,
-        request: bool = False,
         stream: bool = False,
+        request: bool = False,
         reply: str = "",
         timeout: float = 0.5,
         probe: bool = True,
-    ) -> RawMessage | None:
-        """Publish a message.
+    ) -> AsyncIterator[RawMessage] | RawMessage | None:
+        """Publish a message to the given address.
 
         Args:
             addr (Address): The address of the agent.
             msg (RawMessage): The raw message to send.
-            request (bool, optional): Whether this is a request. Defaults to False.
             stream (bool, optional): Whether to request a streaming result. Defaults to False.
+            request (bool, optional): Whether this is a request. Defaults to False. If `stream` is True, then this is always True.
             reply (str, optional): If `request` is True, then this will be the subject to reply to. Defaults to "".
             timeout (float, optional): If `request` is True, then this will be the timeout for the response. Defaults to 0.5.
             probe (bool, optional): Whether to probe the agent before sending the message. Defaults to True.
         """
         pass
 
-    @abc.abstractmethod
-    async def publish_multi(
-        self,
-        addr: Address,
-        msg: RawMessage,
-        probe: bool = True,
-    ) -> AsyncIterator[RawMessage]:
-        """Publish a message and wait for multiple reply messages.
-
-        Args:
-            addr (Address): The address of the agent.
-            msg (RawMessage): The raw message to send.
-            probe (bool, optional): Whether to probe the agent before sending the message. Defaults to True.
-        """
-        pass
-
     @abc.abstractmethod
     async def subscribe(
         self,
@@ -310,20 +294,17 @@ async def run(
         stream: bool = False,
         session_id: str = "",
         timeout: float = 0.5,
-    ) -> RawMessage | AsyncIterator[RawMessage]:
+    ) -> AsyncIterator[RawMessage] | RawMessage | None:
         """Create an agent and run it with the given message."""
         if self.__runtime is None:
             raise ValueError(f"AgentSpec {self.name} is not registered to a runtime.")
 
         session_id = session_id or uuid.uuid4().hex
         addr = Address(name=self.name, id=session_id)
 
-        if stream:
-            return self.__runtime.channel.publish_multi(addr, msg)
-        else:
-            return await self.__runtime.channel.publish(
-                addr, msg, request=True, timeout=timeout
-            )
+        return await self.__runtime.channel.publish(
+            addr, msg, stream=stream, request=True, timeout=timeout
+        )
 
 
 class Runtime(abc.ABC):

coagent/cos/runtime.py

@@ -161,9 +161,10 @@ async def publish_multi(self, request: Request):
         await self._update_message_header_extensions(msg, request)
 
         addr = Address.decode(data["addr"])
-        msgs = self._runtime.channel.publish_multi(
+        msgs = await self._runtime.channel.publish(
             addr=addr,
             msg=msg,
+            stream=True,
             probe=data.get("probe", True),
         )
 

coagent/runtimes/http_runtime.py

@@ -43,8 +43,8 @@ def from_server(cls, server: str, auth: str = "") -> HTTPRuntime:
 class HTTPChannel(BaseChannel):
     """An HTTP-based channel.
 
-    publish: POST /publish
-    publish_multi: POST /publish_multi
+    _publish: POST /publish
+    _publish_stream: POST /publish_multi
     subscribe: POST /subscribe
     new_reply_topic: POST /reply-topics
     """
@@ -64,7 +64,7 @@ async def connect(self) -> None:
     async def close(self) -> None:
         pass
 
-    async def publish(
+    async def _publish(
         self,
         addr: Address,
         msg: RawMessage,
@@ -95,7 +95,7 @@ async def publish(
         if resp.is_error:
             raise_http_error(resp, resp.text)
 
-    async def publish_multi(
+    async def _publish_stream(
         self,
         addr: Address,
         msg: RawMessage,
@@ -290,9 +290,10 @@ async def publish_multi(
         msg: RawMessage,
         probe: bool = True,
     ) -> AsyncIterator[RawMessage]:
-        msgs = self._channel.publish_multi(
+        msgs = await self._channel.publish(
             addr,
             msg,
+            stream=True,
             probe=probe,
         )
         async for msg in msgs:

coagent/runtimes/local_runtime.py

@@ -44,26 +44,6 @@ async def connect(self) -> None:
     async def close(self) -> None:
         pass
 
-    async def publish(
-        self,
-        addr: Address,
-        msg: RawMessage,
-        request: bool = False,
-        stream: bool = False,
-        reply: str = "",
-        timeout: float = 0.5,
-        probe: bool = True,
-    ) -> RawMessage | None:
-        return await self._publish(
-            addr,
-            msg,
-            request=request,
-            stream=stream,
-            reply=reply,
-            timeout=timeout,
-            probe=probe,
-        )
-
     async def subscribe(
         self,
         addr: Address,

coagent/runtimes/nats_runtime.py

@@ -57,26 +57,6 @@ async def close(self) -> None:
         except ConnectionClosedError:
             pass
 
-    async def publish(
-        self,
-        addr: Address,
-        msg: RawMessage,
-        request: bool = False,
-        stream: bool = False,
-        reply: str = "",
-        timeout: float = 0.5,
-        probe: bool = True,
-    ) -> RawMessage | None:
-        return await self._publish(
-            addr,
-            msg,
-            request=request,
-            stream=stream,
-            reply=reply,
-            timeout=timeout,
-            probe=probe,
-        )
-
     async def subscribe(
         self,
         addr: Address,
@@ -102,14 +82,13 @@ async def _publish(
         reply: str = "",
         timeout: float = 0.5,
         probe: bool = True,
-        nonblocking: bool = False,
     ) -> RawMessage | None:
         if addr.is_reply or not probe or await self._probe(addr):
             return await self._nats_publish(
                 addr, msg, request=request, stream=stream, reply=reply, timeout=timeout
             )
 
-        if request or not nonblocking:
+        if request:
             # If in request-reply (or non-blocking) mode, always wait for the reply.
             return await self._create_and_publish(
                 addr,

examples/rich_client.py

@@ -40,9 +40,10 @@ async def ainit(cls):
     async def asend(cls, query: str) -> AsyncIterator[ChatMessage]:
         msg = ChatMessage(role="user", content=query)
         cls.history.messages.append(msg)
-        result = cls.runtime.channel.publish_multi(
+        result = await cls.runtime.channel.publish(
             cls.addr,
             cls.history.encode(),
+            stream=True,
         )
         full_reply = ChatMessage(role="assistant", content="")
         async for chunk in result:

examples/rich_client_textarea.py

@@ -37,9 +37,10 @@ async def ainit(cls):
     async def asend(cls, query: str) -> AsyncIterator[str]:
         msg = ChatMessage(role="user", content=query)
         cls.history.messages.append(msg)
-        result = cls.runtime.channel.publish_multi(
+        result = await cls.runtime.channel.publish(
             cls.addr,
             cls.history.encode(),
+            stream=True,
         )
         content = ""
         async for chunk in result: