Handle tools that return async iterators more properly

Author: RussellLuo

coagent/agents/aswarm/core.py

@@ -180,19 +180,30 @@ async def handle_tool_calls(
             function_result = func(**args)
 
             if is_async_iterator(function_result):
-                # NOTE(luopeng)
+                # NOTE(luopeng):
                 #
-                # If the function returns an async iterator, we assume that
-                # the function is actually a sub-agent, then we should return
-                # the stream directly to the user.
+                # The function is an async generator function. We assume that
+                # it's better to return the stream directly to the user.
                 #
                 # Note that this only works if there's one tool call in the batch.
                 async for chunk in function_result:
                     yield normalize_function_result(chunk)
                 return
 
+            function_result_after_await = await function_result
+            if is_async_iterator(function_result_after_await):
+                # NOTE(luopeng):
+                #
+                # The function returns an async iterator internally. We assume
+                # that it's better to return the stream directly to the user.
+                #
+                # Note that this only works if there's one tool call in the batch.
+                async for chunk in function_result_after_await:
+                    yield normalize_function_result(chunk)
+                return
+
             # Non-streaming results are handled here.
-            raw_result = normalize_function_result(await function_result)
+            raw_result = normalize_function_result(function_result_after_await)
             if raw_result.to_user:
                 # Return the reply directly to the user.
                 yield raw_result

coagent/agents/chat_agent.py

@@ -7,6 +7,7 @@
 from typing import Any, AsyncIterator, Callable
 
 from coagent.core import Address, BaseAgent, Context, handler, logger
+from coagent.core.agent import is_async_iterator
 from pydantic_core import PydanticUndefined
 from pydantic.fields import FieldInfo
 
@@ -44,7 +45,9 @@ def confirm(template: str):
 
     def wrapper(func):
         @functools.wraps(func)
-        async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
+        async def run(
+            *args: Any, **kwargs: Any
+        ) -> AsyncIterator[ChatMessage | str] | ChatMessage | str:
             # Ask the user to confirm if not yet.
             ctx = kwargs.get("ctx", None)
             if ctx and not RunContext(ctx).user_confirmed:
@@ -58,9 +61,11 @@ async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
                     to_user=True,
                 )
 
-            # Note that we assume that the tool is not an async generator,
-            # so we always use `await` here.
-            return await func(*args, **kwargs)
+            result = func(*args, **kwargs)
+            if is_async_iterator(result):
+                return result
+            else:
+                return await result
 
         return run
 
@@ -88,7 +93,9 @@ def submit(template: str = ""):
 
     def wrapper(func):
         @functools.wraps(func)
-        async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
+        async def run(
+            *args: Any, **kwargs: Any
+        ) -> AsyncIterator[ChatMessage | str] | ChatMessage | str:
             # Ask the user to fill in the input form if not yet.
             ctx = kwargs.get("ctx", None)
             if ctx and not RunContext(ctx).user_submitted:
@@ -107,9 +114,11 @@ async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
                     to_user=True,
                 )
 
-            # Note that we assume that the tool is not an async generator,
-            # so we always use `await` here.
-            return await func(*args, **kwargs)
+            result = func(*args, **kwargs)
+            if is_async_iterator(result):
+                return result
+            else:
+                return await result
 
         return run
 
@@ -148,7 +157,9 @@ def wrap_error(func):
     """Decorator to capture and return the possible error when running the given tool."""
 
     @functools.wraps(func)
-    async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
+    async def run(
+        *args: Any, **kwargs: Any
+    ) -> AsyncIterator[ChatMessage | str] | ChatMessage | str:
         try:
             # Fill in the missing arguments with default values if possible.
             #
@@ -163,9 +174,12 @@ async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
                     else:
                         kwargs[name] = default
 
-            # Note that we assume that the tool is not an async generator,
-            # so we always use `await` here.
-            return await func(*args, **kwargs)
+            result = func(*args, **kwargs)
+            if is_async_iterator(result):
+                return result
+            else:
+                return await result
+
         except Exception as exc:
             logger.exception(exc)
             return f"Error: {exc}"