Differentiate ChatAgent from StreamChatAgent

Author: RussellLuo

coagent/agents/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: F401
-from .chat_agent import ChatAgent, confirm, RunContext, tool
+from .chat_agent import ChatAgent, confirm, RunContext, StreamChatAgent, tool
 from .dynamic_triage import DynamicTriage
 from .messages import ChatHistory, ChatMessage
 from .model_client import ModelClient

coagent/agents/chat_agent.py

@@ -114,7 +114,7 @@ async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
     return run
 
 
-class ChatAgent(BaseAgent):
+class StreamChatAgent(BaseAgent):
     def __init__(
         self,
         name: str = "",
@@ -168,7 +168,23 @@ async def agent(self, agent_type: str) -> AsyncIterator[ChatMessage]:
             yield chunk
 
     @handler
-    async def handle(
+    async def handle_history(
+        self, msg: ChatHistory, ctx: Context
+    ) -> AsyncIterator[ChatMessage]:
+        response = self._handle_history(msg, ctx)
+        async for resp in response:
+            yield resp
+
+    @handler
+    async def handle_message(
+        self, msg: ChatMessage, ctx: Context
+    ) -> AsyncIterator[ChatMessage]:
+        history = ChatHistory(messages=[msg])
+        response = self._handle_history(history, ctx)
+        async for resp in response:
+            yield resp
+
+    async def _handle_history(
         self, msg: ChatHistory, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
         # For now, we assume that the agent is processing messages sequentially.
@@ -205,3 +221,23 @@ async def update_user_confirmed(self, history: ChatHistory) -> None:
     async def _has_confirm_message(self, history: ChatHistory) -> bool:
         """Check if the penultimate message is a confirmation message."""
         return len(history.messages) > 1 and history.messages[-2].type == "confirm"
+
+
+class ChatAgent(StreamChatAgent):
+    """Non-streaming ChatAgent."""
+
+    @handler
+    async def handle_history(self, msg: ChatHistory, ctx: Context) -> ChatMessage:
+        accumulated_response = ChatMessage(role="assistant", content="")
+        response = super().handle_history(msg, ctx)
+        async for chunk in response:
+            accumulated_response.content += chunk.content
+        return accumulated_response
+
+    @handler
+    async def handle_message(self, msg: ChatMessage, ctx: Context) -> ChatMessage:
+        accumulated_response = ChatMessage(role="assistant", content="")
+        response = super().handle_message(msg, ctx)
+        async for chunk in response:
+            accumulated_response.content += chunk.content
+        return accumulated_response