Fix wrap_error()

Author: RussellLuo

coagent/agents/chat_agent.py

@@ -156,6 +156,16 @@ def tool(func):
 def wrap_error(func):
     """Decorator to capture and return the possible error when running the given tool."""
 
+    async def __wrap_aiter(
+        aiter_: AsyncIterator[ChatMessage | str],
+    ) -> AsyncIterator[ChatMessage | str]:
+        try:
+            async for chunk in aiter_:
+                yield chunk
+        except Exception as exc:
+            logger.exception(exc)
+            yield f"Error: {exc}"
+
     @functools.wraps(func)
     async def run(
         *args: Any, **kwargs: Any
@@ -180,7 +190,7 @@ async def run(
 
             result = func(*args, **kwargs)
             if is_async_iterator(result):
-                return result
+                return __wrap_aiter(result)
             else:
                 return await result
 

tests/agents/test_chat_agent.py

@@ -1,11 +1,13 @@
+from typing import AsyncIterator
+
 from pydantic import Field
 import pytest
 
 from coagent.agents.chat_agent import wrap_error
 
 
 @pytest.mark.asyncio
-async def test_wrap_error():
+async def test_wrap_error_normal():
     @wrap_error
     async def func(
         a: int = Field(..., description="Argument a"),
@@ -16,3 +18,22 @@ async def func(
     assert await func() == "Error: Missing required argument 'a'"
     assert await func(a=1) == 1
     assert await func(a=1, b=0) == "Error: division by zero"
+
+
+@pytest.mark.asyncio
+async def test_wrap_error_aiter():
+    @wrap_error
+    async def func(
+        a: int = Field(..., description="Argument a"),
+        b: int = Field(1, description="Argument b"),
+    ) -> AsyncIterator[float]:
+        yield a / b
+
+    result = await func()
+    assert result == "Error: Missing required argument 'a'"
+
+    result = await func(a=1)
+    assert await anext(result) == 1
+
+    result = await func(a=1, b=0)
+    assert await anext(result) == "Error: division by zero"