Refactor the streaming mechanism

Author: RussellLuo

coagent/agents/messages.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any
 
 from pydantic import Field
@@ -14,6 +16,15 @@ class ChatMessage(Message):
         default=False, description="Whether the message is sent directly to user."
     )
 
+    def __add__(self, other: ChatMessage) -> ChatMessage:
+        return ChatMessage(
+            role=self.role,
+            content=self.content + other.content,
+            type=self.type,
+            sender=self.sender,
+            to_user=self.to_user,
+        )
+
     def model_dump(self, **kwargs) -> dict[str, Any]:
         return super().model_dump(include={"role", "content"}, **kwargs)
 

coagent/agents/parallel.py

@@ -7,13 +7,14 @@
     GenericMessage,
     Message,
     RawMessage,
+    Reply,
     SetReplyAgent,
 )
 
 
 class StartAggregation(Message):
     candidates: list[str]
-    reply_addr: Address
+    reply_info: Reply | None
 
 
 class AggregationStatus(Message):
@@ -62,17 +63,17 @@ async def handle(self, msg: GenericMessage, ctx: Context) -> None:
         self._results.append(msg.encode())
 
         if len(self._results) == len(self._data.candidates):
-            if self._data.reply_addr:
+            if self._data.reply_info:
                 result = await self._aggregate(self._results)
-                await self.channel.publish(self._data.reply_addr, result)
+                await self.send_reply(self._data.reply_info, result)
             self._busy = False
 
-    async def aggregate(self, results: list[RawMessage]) -> RawMessage:
+    async def aggregate(self, results: list[RawMessage]) -> Message:
         """Aggregate the results to a single one.
 
         Override this method to provide custom aggregation logic.
         """
-        return AggregationResult(results=results).encode()
+        return AggregationResult(results=results)
 
 
 class Parallel(BaseAgent):
@@ -87,12 +88,13 @@ def __init__(self, *agent_types: str, aggregator: str = ""):
 
     async def started(self) -> None:
         aggregator_addr = Address(name=self._aggregator_type, id=self.address.id)
+        aggregator_reply = Reply(address=aggregator_addr)
         # Make each agent reply to the aggregator agent.
         for agent_type in self._agent_types:
             addr = Address(name=agent_type, id=self.address.id)
             await self.channel.publish(
                 addr,
-                SetReplyAgent(address=aggregator_addr).encode(),
+                SetReplyAgent(reply_info=aggregator_reply).encode(),
             )
 
     @handler
@@ -101,16 +103,14 @@ async def handle(self, msg: GenericMessage, ctx: Context) -> None:
             return
 
         # Let the aggregator agent reply to the sending agent, if asked.
-        reply_address = self.reply_address or msg.reply
-        if reply_address:
+        reply = self.reply or msg.reply
+        if reply:
             # Reset the reply address of the message, since it will be replied by the aggregator agent.
             msg.reply = None
 
         result = await self.channel.publish(
             Address(name=self._aggregator_type, id=self.address.id),
-            StartAggregation(
-                candidates=self._agent_types, reply_addr=reply_address
-            ).encode(),
+            StartAggregation(candidates=self._agent_types, reply_info=reply).encode(),
             request=True,
         )
         status = AggregationStatus.decode(result)

coagent/agents/sequential.py

@@ -3,6 +3,7 @@
     BaseAgent,
     Context,
     GenericMessage,
+    Reply,
     SetReplyAgent,
     handler,
 )
@@ -20,9 +21,10 @@ async def started(self) -> None:
             # Set the reply address of the current agent to be the next agent.
             addr = Address(name=self._agent_types[i], id=self.address.id)
             next_addr = Address(name=self._agent_types[i + 1], id=self.address.id)
+            reply = Reply(address=next_addr)
             await self.channel.publish(
                 addr,
-                SetReplyAgent(address=next_addr).encode(),
+                SetReplyAgent(reply_info=reply).encode(),
             )
 
     @handler
@@ -31,12 +33,12 @@ async def handle(self, msg: GenericMessage, ctx: Context) -> None:
             return
 
         # Let the last agent reply to the sending agent, if asked.
-        reply_address = self.reply_address or msg.reply
-        if reply_address:
+        reply = self.reply or msg.reply
+        if reply:
             last_addr = Address(name=self._agent_types[-1], id=self.address.id)
             await self.channel.publish(
                 last_addr,
-                SetReplyAgent(address=reply_address).encode(),
+                SetReplyAgent(reply_info=reply).encode(),
             )
             # Reset the reply address of the message, since it will be replied by the last agent.
             msg.reply = None

coagent/core/__init__.py

@@ -17,6 +17,7 @@
     MessageHeader,
     new,
     RawMessage,
+    Reply,
     Subscription,
 )
 from .util import idle_loop

coagent/core/agent.py

@@ -4,7 +4,7 @@
 
 from pydantic import BaseModel, ValidationError
 
-from .exceptions import MessageDecodeError, InternalError
+from .exceptions import MessageDecodeError, InternalError, StreamError
 from .logger import logger
 from .messages import (
     Cancel,
@@ -17,7 +17,7 @@
     SetReplyAgent,
     StopIteration,
 )
-from .types import Address, Agent, Channel, RawMessage, State, Subscription
+from .types import Address, Agent, Channel, RawMessage, Reply, State, Subscription
 
 
 class Context:
@@ -122,8 +122,8 @@ def __init__(self, timeout: float = 60):
         # this would result in a lot of messages.
         self._lock: asyncio.Lock = asyncio.Lock()
 
-        # Normally reply_address is set by an orchestration agent by sending a `SetReplyAgent` message.
-        self.reply_address: Address | None = None
+        # Normally `reply` is set by an orchestration agent by sending a `SetReplyAgent` message.
+        self.reply: Reply | None = None
 
         handlers, message_types = self.__collect_handlers()
         # A list of handlers that are registered to handle messages.
@@ -269,7 +269,7 @@ async def _handle_data(self) -> None:
                     await self.stopped()
 
                 case SetReplyAgent():
-                    self.reply_address = msg.address
+                    self.reply = msg.reply_info
 
                 case ProbeAgent() | Empty():
                     # Do not handle probes and empty messages.
@@ -282,18 +282,41 @@ async def _handle_data_custom(self, msg: Message, ctx: Context) -> None:
         """Handle user-defined DATA messages."""
         h: Handler = self.__get_handler(msg)
         result = h(self, msg, ctx)
+        if not is_async_iterator(result):
+            result = await result or Empty()
+        await self.__send_reply(msg.reply, result)
+
+    async def __send_reply(
+        self, in_msg_reply: Reply, result: Message | AsyncIterator[Message]
+    ) -> bool:
+        reply = self.reply or in_msg_reply
+        if not reply:
+            return False
+
+        # Reply to the sender if asked.
+        await self.send_reply(reply, result)
+        return True
 
-        async def pub(x: Message):
-            await self.__send_reply(msg.reply, x)
+    async def send_reply(
+        self,
+        to: Reply,
+        result: Message | AsyncIterator[Message],
+    ) -> None:
+        async def pub(msg: Message):
+            await self.channel.publish(to.address, msg.encode())
 
         async def pub_exc(exc: BaseException):
             err = InternalError.from_exception(exc)
             await pub(err.encode_message())
 
-        if is_async_iterator(result):
+        if to.stream:
             try:
-                async for x in result:
-                    await pub(x)
+                if is_async_iterator(result):
+                    async for msg in result:
+                        await pub(msg)
+                else:
+                    msg = result
+                    await pub(msg)
             except asyncio.CancelledError as exc:
                 await pub_exc(exc)
                 raise
@@ -304,23 +327,26 @@ async def pub_exc(exc: BaseException):
                 await pub(StopIteration())
         else:
             try:
-                x = await result or Empty()
-                await pub(x)
+                if is_async_iterator(result):
+                    accumulated: RawMessage | None = None
+                    async for msg in result:
+                        if not accumulated:
+                            accumulated = msg
+                        else:
+                            try:
+                                accumulated += msg
+                            except TypeError:
+                                await pub_exc(StreamError("Streaming mode is required"))
+                    await pub(accumulated)
+                else:
+                    msg = result
+                    await pub(msg)
             except asyncio.CancelledError as exc:
                 await pub_exc(exc)
                 raise
             except Exception as exc:
                 await pub_exc(exc)
 
-    async def __send_reply(self, in_msg_reply: Address, out_msg: Message) -> bool:
-        reply_address = self.reply_address or in_msg_reply
-        if not reply_address:
-            return False
-
-        # Reply to the sending agent if asked.
-        await self.channel.publish(reply_address, out_msg.encode())
-        return True
-
     def __get_handler(self, msg: Message) -> Handler | None:
         msg_type: Type[Any] = type(msg)
 

coagent/core/exceptions.py

@@ -99,3 +99,7 @@ def from_exception(
 
 class DeadlineExceededError(BaseError):
     """Raised when a context deadline is exceeded."""
+
+
+class StreamError(BaseError):
+    """Raised when the sender requests a non-streaming result but the receiver sends a stream."""

coagent/core/messages.py

@@ -3,17 +3,26 @@
 import json
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
-from .types import Address, MessageHeader, RawMessage
+from .types import MessageHeader, RawMessage, Reply
 
 
 class Message(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    reply: Address | None = Field(default=None, description="Reply address.")
+    reply: Reply | None = Field(default=None, description="Reply information.")
     extensions: dict = Field(
         default_factory=dict, description="Extension fields from RawMessage header."
     )
 
+    def __add__(self, other: Message) -> Message:
+        """Concatenate two messages.
+
+        This binary operator is mainly used to aggregate multiple streaming
+        messages into one message when the sender requests to receive a
+        non-streaming result.
+        """
+        return NotImplemented
+
     def encode(
         self, content_type: str = "application/json", exclude_defaults: bool = True
     ) -> RawMessage:
@@ -114,9 +123,12 @@ class ProbeAgent(Message):
 
 
 class SetReplyAgent(Message):
-    """A message to set the agent to reply to."""
+    """A message to set the reply information of an agent.
+
+    This is mainly useful when orchestrating multiple agents to work together.
+    """
 
-    address: Address
+    reply_info: Reply
 
 
 class Empty(Message):

coagent/core/runtime.py

@@ -87,6 +87,7 @@ async def publish_multi(
             addr,
             msg,
             request=True,
+            stream=True,
             reply=inbox,
             probe=probe,
         )

coagent/core/types.py

@@ -88,6 +88,13 @@ def decode(cls, data: dict) -> Address:
         return cls.model_validate(data)
 
 
+class Reply(BaseModel):
+    address: Address = Field(..., description="Reply address.")
+    stream: bool = Field(
+        False, description="Whether the sender requests a streaming result."
+    )
+
+
 class MessageHeader(BaseModel):
     type: str = Field(..., description="Message type name.")
     content_type: str = Field(
@@ -98,7 +105,7 @@ class MessageHeader(BaseModel):
 
 class RawMessage(BaseModel):
     header: MessageHeader = Field(..., description="Message header.")
-    reply: Address | None = Field(default=None, description="Reply address.")
+    reply: Reply | None = Field(default=None, description="Reply information.")
     content: bytes = Field(default=b"", description="Message content.")
 
     def encode(self, mode: str = "python", exclude_defaults: bool = True) -> dict:
@@ -223,6 +230,7 @@ async def publish(
         addr: Address,
         msg: RawMessage,
         request: bool = False,
+        stream: bool = False,
         reply: str = "",
         timeout: float = 0.5,
         probe: bool = True,
@@ -233,6 +241,7 @@ async def publish(
             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.
             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.

coagent/runtimes/http_runtime.py

@@ -69,6 +69,7 @@ async def publish(
         addr: Address,
         msg: RawMessage,
         request: bool = False,
+        stream: bool = False,
         reply: str = "",
         timeout: float = 5.0,
         probe: bool = True,
@@ -77,6 +78,7 @@ async def publish(
             addr=addr.encode(mode="json"),
             msg=msg.encode(mode="json"),
             request=request,
+            stream=stream,
             reply=reply,
             timeout=timeout,
             probe=probe,
@@ -267,6 +269,7 @@ async def publish(
         addr: Address,
         msg: RawMessage,
         request: bool = False,
+        stream: bool = False,
         reply: str = "",
         timeout: float = 5.0,
         probe: bool = True,
@@ -275,6 +278,7 @@ async def publish(
             addr,
             msg,
             request=request,
+            stream=stream,
             reply=reply,
             timeout=timeout,
             probe=probe,