Support tools to add parameter descriptions using typing.Annotated

Author: RussellLuo

coagent/agents/aswarm/util.py

@@ -1,4 +1,5 @@
 import inspect
+import typing
 from datetime import datetime
 from typing import Any
 
@@ -172,15 +173,34 @@ def greet(
         # Determine the default value
         default = param.default
 
-        # Check if the parameter has a Field with description
-        if isinstance(param.default, FieldInfo):
-            field = param.default
-            fields[param_name] = (annotation, field)
-        elif param.default != inspect.Parameter.empty:
-            fields[param_name] = (annotation, param.default)
+        # Determine the field information.
+        if isinstance(default, FieldInfo):
+            # The default value is already a Field.
+            typ = annotation
+            field = default
+        elif default != inspect.Parameter.empty:
+            # Normal default value
+            if typing.get_origin(annotation) == typing.Annotated:
+                # The parameter is annotated with metadata.
+                typ, description = typing.get_args(annotation)[:2]
+                field = Field(default=default, description=str(description))
+            else:
+                # No metadata, use Field without description.
+                typ = annotation
+                field = Field(default=default)
         else:
-            # If no default value, use Field without default
-            fields[param_name] = (annotation, Field(...))
+            # No default value
+            if typing.get_origin(annotation) == typing.Annotated:
+                # The parameter is annotated with metadata.
+                # Always treat the first metadata element as the description.
+                typ, description = typing.get_args(annotation)[:2]
+                field = Field(..., description=str(description))
+            else:
+                # No default value and no metadata, use Field without default.
+                typ = annotation
+                field = Field(...)
+
+        fields[param_name] = (typ, field)
 
     # 3. Create the Pydantic model
     model_name = f"{func.__name__}"

tests/agents/test_util.py

@@ -1,6 +1,10 @@
+from typing import Annotated
+
 import pytest
 
 from coagent.agents.util import chat
+from coagent.agents.aswarm.util import function_to_jsonschema
+from pydantic import Field
 
 
 @pytest.mark.asyncio
@@ -25,3 +29,88 @@ async def test_chat_stream(mock_model_client):
         # Only one chunk.
         chunk = _chunk
     assert chunk and chunk.content == "hello"
+
+
+def test_function_to_jsonschema_normal():
+    def func(a: int, b: str = "ok") -> None:
+        """This is a test function."""
+        pass
+
+    schema = function_to_jsonschema(func)
+    assert schema == {
+        "function": {
+            "description": "This is a test function.",
+            "name": "func",
+            "parameters": {
+                "properties": {
+                    "a": {"title": "A", "type": "integer"},
+                    "b": {"default": "ok", "title": "B", "type": "string"},
+                },
+                "required": ["a"],
+                "title": "func",
+                "type": "object",
+            },
+        },
+        "type": "function",
+    }
+
+
+def test_function_to_jsonschema_annotated():
+    def func(a: Annotated[int, "Param a"], b: Annotated[str, "Param b"] = "ok") -> None:
+        """This is a test function."""
+        pass
+
+    schema = function_to_jsonschema(func)
+    assert schema == {
+        "function": {
+            "description": "This is a test function.",
+            "name": "func",
+            "parameters": {
+                "properties": {
+                    "a": {"description": "Param a", "title": "A", "type": "integer"},
+                    "b": {
+                        "default": "ok",
+                        "description": "Param b",
+                        "title": "B",
+                        "type": "string",
+                    },
+                },
+                "required": ["a"],
+                "title": "func",
+                "type": "object",
+            },
+        },
+        "type": "function",
+    }
+
+
+def test_function_to_jsonschema_pydantic_field():
+    def func(
+        a: int = Field(description="Param a"),
+        b: str = Field(default="ok", description="Param b"),
+    ) -> None:
+        """This is a test function."""
+        pass
+
+    schema = function_to_jsonschema(func)
+    assert schema == {
+        "function": {
+            "description": "This is a test function.",
+            "name": "func",
+            "parameters": {
+                "properties": {
+                    "a": {"description": "Param a", "title": "A", "type": "integer"},
+                    "b": {
+                        "default": "ok",
+                        "description": "Param b",
+                        "title": "B",
+                        "type": "string",
+                    },
+                },
+                "required": ["a"],
+                "title": "func",
+                "type": "object",
+            },
+        },
+        "type": "function",
+    }