Support composite arg type parsing in dspy.Tool (#8095)

* tool calling args parsing

* fix wrong comment

* Add test

Author: chenmoneygithub

dspy/predict/react.py

@@ -1,8 +1,7 @@
 import logging
-from typing import Any, Callable, Literal, get_origin
+from typing import Any, Callable, Literal
 
 from litellm import ContextWindowExceededError
-from pydantic import BaseModel
 
 import dspy
 from dspy.primitives.program import Module
@@ -82,18 +81,7 @@ def forward(self, **input_args):
             trajectory[f"tool_args_{idx}"] = pred.next_tool_args
 
             try:
-                parsed_tool_args = {}
-                tool = self.tools[pred.next_tool_name]
-                for k, v in pred.next_tool_args.items():
-                    if hasattr(tool, "arg_types") and k in tool.arg_types:
-                        arg_type = tool.arg_types[k]
-                        if isinstance((origin := get_origin(arg_type) or arg_type), type) and issubclass(
-                            origin, BaseModel
-                        ):
-                            parsed_tool_args[k] = arg_type.model_validate(v)
-                            continue
-                    parsed_tool_args[k] = v
-                trajectory[f"observation_{idx}"] = self.tools[pred.next_tool_name](**parsed_tool_args)
+                trajectory[f"observation_{idx}"] = self.tools[pred.next_tool_name](**pred.next_tool_args)
             except Exception as e:
                 trajectory[f"observation_{idx}"] = f"Failed to execute: {e}"
 

dspy/primitives/tool.py

@@ -2,7 +2,7 @@
 from typing import Any, Callable, Optional, get_origin, get_type_hints
 
 from jsonschema import ValidationError, validate
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, TypeAdapter, create_model
 
 from dspy.utils.callback import with_callbacks
 
@@ -130,6 +130,19 @@ def _parse_function(self, func: Callable, arg_desc: dict[str, str] = None):
         self.args = self.args or args
         self.arg_types = self.arg_types or arg_types
 
+    def _parse_args(self, **kwargs):
+        parsed_kwargs = {}
+        for k, v in kwargs.items():
+            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], ...))
+                parsed = pydantic_wrapper.model_validate({"value": v})
+                parsed_kwargs[k] = parsed.value
+            else:
+                parsed_kwargs[k] = v
+        return parsed_kwargs
+
     @with_callbacks
     def __call__(self, **kwargs):
         for k, v in kwargs.items():
@@ -142,4 +155,6 @@ def __call__(self, **kwargs):
                     validate(instance=instance, schema=self.args[k])
             except ValidationError as e:
                 raise ValueError(f"Arg {k} is invalid: {e.message}")
-        return self.func(**kwargs)
+
+        parsed_kwargs = self._parse_args(**kwargs)
+        return self.func(**parsed_kwargs)

tests/primitives/test_tool.py

@@ -159,3 +159,37 @@ def foo(x=100):
     tool = Tool(foo)
     assert tool.args["x"]["default"] == 100
     assert not hasattr(tool.args["x"], "type")
+
+
+def test_tool_call_parses_args():
+    tool = Tool(dummy_with_pydantic)
+
+    args = {
+        "model": {
+            "field1": "hello",
+            "field2": 123,
+        }
+    }
+
+    result = tool(**args)
+    assert result == "hello 123"
+
+
+def test_tool_call_parses_nested_list_of_pydantic_model():
+    def dummy_function(x: list[list[DummyModel]]):
+        return x
+
+    tool = Tool(dummy_function)
+    args = {
+        "x": [
+            [
+                {
+                    "field1": "hello",
+                    "field2": 123,
+                }
+            ]
+        ]
+    }
+
+    result = tool(**args)
+    assert result == [[DummyModel(field1="hello", field2=123)]]