Improve wrap_error to fill in missing arguments with default values

Author: RussellLuo

coagent/agents/chat_agent.py

@@ -6,6 +6,7 @@
 from typing import Any, AsyncIterator, Callable
 
 from coagent.core import Address, BaseAgent, Context, handler, logger
+from pydantic_core import PydanticUndefined
 from pydantic.fields import FieldInfo
 
 from .aswarm import Agent as SwarmAgent, Swarm
@@ -39,14 +40,10 @@ async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
             if ctx and not RunContext(ctx).user_confirmed:
                 # We assume that all meaningful arguments (includes `ctx` but
                 # excepts possible `self`) are keyword arguments. Therefore,
-                # here we use kwargs as the source of template variables.
-                tmpl_vars = {
-                    k: v.default if isinstance(v, FieldInfo) else v
-                    for k, v in kwargs.items()
-                }
+                # here we use kwargs directly as the template variables.
                 return ChatMessage(
                     role="assistant",
-                    content=template.format(**tmpl_vars),
+                    content=template.format(**kwargs),
                     type="confirm",
                     to_user=True,
                 )
@@ -94,6 +91,19 @@ def wrap_error(func):
     @functools.wraps(func)
     async def run(*args: Any, **kwargs: Any) -> ChatMessage | str:
         try:
+            # Fill in the missing arguments with default values if possible.
+            #
+            # Note that we assume that all meaningful arguments (includes `ctx`
+            # but excepts possible `self`) are keyword arguments.
+            sig = inspect.signature(func)
+            for name, param in sig.parameters.items():
+                if name not in kwargs and isinstance(param.default, FieldInfo):
+                    default = param.default.default
+                    if default is PydanticUndefined:
+                        raise ValueError(f"Missing required argument {name!r}")
+                    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)

pyproject.toml

@@ -78,6 +78,7 @@ linting = [
 ]
 testing = [
     "pytest",
+    "pytest-asyncio",
     "pytest-cov",
 ]
 

tests/agents/test_chat_agent.py

@@ -0,0 +1,18 @@
+from pydantic import Field
+import pytest
+
+from coagent.agents.chat_agent import wrap_error
+
+
+@pytest.mark.asyncio
+async def test_wrap_error():
+    @wrap_error
+    async def func(
+        a: int = Field(..., description="Argument a"),
+        b: int = Field(1, description="Argument b"),
+    ) -> float:
+        return a / b
+
+    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"