Support async in dspy.Tool (#8106)

* async Tool

* fix tests

Author: chenmoneygithub

dspy/primitives/tool.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 from typing import Any, Callable, Optional, get_origin, get_type_hints
 
@@ -132,10 +133,26 @@ def _parse_function(self, func: Callable, arg_desc: dict[str, str] = None):
         self.arg_types = self.arg_types or arg_types
         self.has_kwargs = any([param.kind == param.VAR_KEYWORD for param in sig.parameters.values()])
 
-    def _parse_args(self, **kwargs):
+    def _validate_and_parse_args(self, **kwargs):
+        # Validate the args value comply to the json schema.
+        for k, v in kwargs.items():
+            if k not in self.args:
+                if self.has_kwargs:
+                    continue
+                else:
+                    raise ValueError(f"Arg {k} is not in the tool's args.")
+            try:
+                instance = v.model_dump() if hasattr(v, "model_dump") else v
+                type_str = self.args[k].get("type")
+                if type_str is not None and type_str != "Any":
+                    validate(instance=instance, schema=self.args[k])
+            except ValidationError as e:
+                raise ValueError(f"Arg {k} is invalid: {e.message}")
+
+        # Parse the args to the correct type.
         parsed_kwargs = {}
         for k, v in kwargs.items():
-            if k in self.arg_types and self.arg_types[k] != any:
+            if k in self.arg_types and self.arg_types[k] != Any:
                 # Create a pydantic model wrapper with a dummy field `value` to parse the arg to the correct type.
                 # This is specifically useful for handling nested Pydantic models like `list[list[MyPydanticModel]]`
                 pydantic_wrapper = create_model("Wrapper", value=(self.arg_types[k], ...))
@@ -147,19 +164,15 @@ def _parse_args(self, **kwargs):
 
     @with_callbacks
     def __call__(self, **kwargs):
-        for k, v in kwargs.items():
-            if k not in self.args:
-                if self.has_kwargs:
-                    # If the tool has kwargs, skip validation for unknown args
-                    continue
-                raise ValueError(f"Arg {k} is not in the tool's args.")
-            try:
-                instance = v.model_dump() if hasattr(v, "model_dump") else v
-                type_str = self.args[k].get("type")
-                if type_str is not None and type_str != "Any":
-                    validate(instance=instance, schema=self.args[k])
-            except ValidationError as e:
-                raise ValueError(f"Arg {k} is invalid: {e.message}")
-
-        parsed_kwargs = self._parse_args(**kwargs)
-        return self.func(**parsed_kwargs)
+        parsed_kwargs = self._validate_and_parse_args(**kwargs)
+        result = self.func(**parsed_kwargs)
+        if asyncio.iscoroutine(result):
+            raise ValueError("You are calling `__call__` on an async tool, please use `acall` instead.")
+        return result
+
+    async def acall(self, **kwargs):
+        parsed_kwargs = self._validate_and_parse_args(**kwargs)
+        result = self.func(**parsed_kwargs)
+        if not asyncio.iscoroutine(result):
+            raise ValueError("You are calling `acall` on a non-async tool, please use `__call__` instead.")
+        return await result

tests/primitives/test_tool.py

@@ -2,6 +2,7 @@
 from pydantic import BaseModel
 from dspy.primitives.tool import Tool
 from typing import Any, Optional
+import asyncio
 
 
 # Test fixtures
@@ -67,6 +68,52 @@ def complex_dummy_function(profile: UserProfile, priority: int, notes: Optional[
     }
 
 
+async def async_dummy_function(x: int, y: str = "hello") -> str:
+    """An async dummy function for testing.
+
+    Args:
+        x: An integer parameter
+        y: A string parameter
+    """
+    await asyncio.sleep(0.1)  # Simulate some async work
+    return f"{y} {x}"
+
+
+async def async_dummy_with_pydantic(model: DummyModel) -> str:
+    """An async dummy function that accepts a Pydantic model."""
+    await asyncio.sleep(0.1)  # Simulate some async work
+    return f"{model.field1} {model.field2}"
+
+
+async def async_complex_dummy_function(
+    profile: UserProfile, priority: int, notes: Optional[str] = None
+) -> dict[str, Any]:
+    """Process user profile with complex nested structure asynchronously.
+
+    Args:
+        profile: User profile containing nested contact and address information
+        priority: Priority level of the processing
+        notes: Optional processing notes
+    """
+    # Simulate some async processing work
+    await asyncio.sleep(0.1)
+
+    primary_address = next(
+        (addr for addr in profile.contact.addresses if addr.is_primary), profile.contact.addresses[0]
+    )
+
+    # Simulate more async work after finding primary address
+    await asyncio.sleep(0.1)
+
+    return {
+        "user_id": profile.user_id,
+        "name": profile.name,
+        "priority": priority,
+        "primary_address": primary_address.model_dump(),
+        "notes": notes,
+    }
+
+
 def test_basic_initialization():
     tool = Tool(name="test_tool", desc="A test tool", args={"param1": {"type": "string"}}, func=lambda x: x)
     assert tool.name == "test_tool"
@@ -198,7 +245,107 @@ def dummy_function(x: list[list[DummyModel]]):
 def test_tool_call_kwarg():
     def fn(x: int, **kwargs):
         return kwargs
+
     tool = Tool(fn)
 
     assert tool(x=1, y=2, z=3) == {"y": 2, "z": 3}
 
+
+@pytest.mark.asyncio
+async def test_async_tool_from_function():
+    tool = Tool(async_dummy_function)
+
+    assert tool.name == "async_dummy_function"
+    assert "An async dummy function for testing" in tool.desc
+    assert "x" in tool.args
+    assert "y" in tool.args
+    assert tool.args["x"]["type"] == "integer"
+    assert tool.args["y"]["type"] == "string"
+    assert tool.args["y"]["default"] == "hello"
+
+    # Test async call
+    result = await tool.acall(x=42, y="hello")
+    assert result == "hello 42"
+
+
+@pytest.mark.asyncio
+async def test_async_tool_with_pydantic():
+    tool = Tool(async_dummy_with_pydantic)
+
+    assert tool.name == "async_dummy_with_pydantic"
+    assert "model" in tool.args
+    assert tool.args["model"]["type"] == "object"
+    assert "field1" in tool.args["model"]["properties"]
+    assert "field2" in tool.args["model"]["properties"]
+
+    # Test async call with pydantic model
+    model = DummyModel(field1="test", field2=123)
+    result = await tool.acall(model=model)
+    assert result == "test 123"
+
+    # Test async call with dict
+    result = await tool.acall(model={"field1": "test", "field2": 123})
+    assert result == "test 123"
+
+
+@pytest.mark.asyncio
+async def test_async_tool_with_complex_pydantic():
+    tool = Tool(async_complex_dummy_function)
+
+    profile = UserProfile(
+        user_id=1,
+        name="Test User",
+        contact=ContactInfo(
+            email="[email protected]",
+            addresses=[
+                Address(street="123 Main St", city="Test City", zip_code="12345", is_primary=True),
+                Address(street="456 Side St", city="Test City", zip_code="12345"),
+            ],
+        ),
+    )
+
+    result = await tool.acall(profile=profile, priority=1, notes="Test note")
+    assert result["user_id"] == 1
+    assert result["name"] == "Test User"
+    assert result["priority"] == 1
+    assert result["notes"] == "Test note"
+    assert result["primary_address"]["street"] == "123 Main St"
+
+
+@pytest.mark.asyncio
+async def test_async_tool_invalid_call():
+    tool = Tool(async_dummy_function)
+    with pytest.raises(ValueError):
+        await tool.acall(x="not an integer", y="hello")
+
+
+@pytest.mark.asyncio
+async def test_async_tool_with_kwargs():
+    async def fn(x: int, **kwargs):
+        return kwargs
+
+    tool = Tool(fn)
+
+    result = await tool.acall(x=1, y=2, z=3)
+    assert result == {"y": 2, "z": 3}
+
+
+@pytest.mark.asyncio
+async def test_async_concurrent_calls():
+    """Test that multiple async tools can run concurrently."""
+    tool = Tool(async_dummy_function)
+
+    # Create multiple concurrent calls
+    tasks = [tool.acall(x=i, y=f"hello{i}") for i in range(5)]
+
+    # Run them concurrently and measure time
+    start_time = asyncio.get_event_loop().time()
+    results = await asyncio.gather(*tasks)
+    end_time = asyncio.get_event_loop().time()
+
+    # Verify results, `asyncio.gather` returns results in the order of the tasks
+    assert results == [f"hello{i} {i}" for i in range(5)]
+
+    # Check that it ran concurrently (should take ~0.1s, not ~0.5s)
+    # We use 0.3s as threshold to account for some overhead
+    assert end_time - start_time < 0.3