Merge chat_stream() into chat()

Author: RussellLuo

coagent/agents/util.py

@@ -9,8 +9,13 @@
 
 
 async def chat(
-    messages: list[ChatMessage], client: ModelClient = default_model_client
-) -> ChatMessage:
+    messages: list[ChatMessage],
+    stream: bool = False,
+    client: ModelClient = default_model_client,
+) -> AsyncIterator[ChatMessage] | ChatMessage:
+    if stream:
+        return _chat_stream(messages, client)
+
     response = await client.acompletion(
         messages=[m.model_dump() for m in messages],
     )
@@ -21,7 +26,7 @@ async def chat(
     )
 
 
-async def chat_stream(
+async def _chat_stream(
     messages: list[ChatMessage], client: ModelClient = default_model_client
 ) -> AsyncIterator[ChatMessage]:
     response = await client.acompletion(

examples/app-builder/dev.py

@@ -2,7 +2,7 @@
 from typing import AsyncIterator
 
 from coagent.agents.chat_agent import ChatMessage
-from coagent.agents.util import chat_stream
+from coagent.agents.util import chat
 from coagent.core import (
     AgentSpec,
     BaseAgent,
@@ -50,7 +50,8 @@ async def handle(
         ]
 
         reply = ""
-        async for chunk in chat_stream(msgs):
+        response = await chat(msgs, stream=True)
+        async for chunk in response:
             yield ChatMessage(role="assistant", content=chunk.content)
             reply += chunk.content
 

examples/app-builder/qa.py

@@ -2,7 +2,7 @@
 from typing import AsyncIterator
 
 from coagent.agents.chat_agent import ChatMessage
-from coagent.agents.util import chat_stream
+from coagent.agents.util import chat
 from coagent.core import (
     AgentSpec,
     BaseAgent,
@@ -69,7 +69,8 @@ async def handle(
         ]
 
         reply = ""
-        async for chunk in chat_stream(msgs):
+        response = await chat(msgs, stream=True)
+        async for chunk in response:
             yield ChatMessage(role="assistant", content=chunk.content)
             reply += chunk.content