Repository: stanfordnlp/dspy
Pull request #7658 (2025-02-03 19:59:44)

Improve `dspy.Tool`
Some changes to `dspy.Tool`:

1. Move it under `primitives` to separate it from `dspy.ReAct` definition.
2. Resolve relative model json schema with absolute schema, i.e., expanding all those `$ref`s. 
3. Provide a way for users to specify parameter description, which is proved to be useful for improving function calling quality. 


Patch:
diff --git a/dspy/predict/__init__.py b/dspy/predict/__init__.py
index 21f59bc4..bc85ba7b 100644
--- a/dspy/predict/__init__.py
+++ b/dspy/predict/__init__.py
@@ -5,7 +5,7 @@ from dspy.predict.knn import KNN
 from dspy.predict.multi_chain_comparison import MultiChainComparison
 from dspy.predict.predict import Predict
 from dspy.predict.program_of_thought import ProgramOfThought
-from dspy.predict.react import ReAct, Tool
+from dspy.predict.react import ReAct
 from dspy.predict.parallel import Parallel
 
 __all__ = [
@@ -17,6 +17,5 @@ __all__ = [
     "Predict",
     "ProgramOfThought",
     "ReAct",
-    "Tool",
     "Parallel",
 ]
diff --git a/dspy/predict/react.py b/dspy/predict/react.py
index fbbb6fa1..e14ba843 100644
--- a/dspy/predict/react.py
+++ b/dspy/predict/react.py
@@ -1,62 +1,12 @@
-import inspect
-from typing import Any, Callable, Literal, get_origin, get_type_hints
+from typing import Any, Callable, Literal, get_origin
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel
 
 import dspy
 from dspy.primitives.program import Module
+from dspy.primitives.tool import Tool
 from dspy.signatures.signature import ensure_signature
-from dspy.utils.callback import with_callbacks
-
-class Tool:
-    def __init__(
-        self,
-        func: Callable,
-        name: str = None,
-        desc: str = None,
-        args: dict[str, Any] = None,
-    ):
-        annotations_func = func if inspect.isfunction(func) or inspect.ismethod(func) else func.__call__
-        
-        self.func = func
-        self.name = name or getattr(func, "__name__", type(func).__name__)
-        self.desc = (
-            desc
-            or getattr(func, "__doc__", None)
-            or getattr(annotations_func, "__doc__", "")
-        )
-        self.args = {}
-        self.arg_types = {}
-
-        # If an explicit args dict is passed, use that; otherwise, extract from the function.
-        if args is not None:
-            hints = args
-        else:
-            # Use inspect.signature to get all parameter names
-            sig = inspect.signature(annotations_func)
-            # Get available type hints
-            available_hints = get_type_hints(annotations_func)
-            # Build a dictionary of parameter name -> type (defaulting to Any when missing)
-            hints = {
-                param_name: available_hints.get(param_name, Any)
-                for param_name in sig.parameters.keys()
-            }
-
-        # Process each argument's type to generate its JSON schema.
-        for k, v in hints.items():
-            self.arg_types[k] = v
-            if k == "return":
-                continue
-            # Check if the type (or its origin) is a subclass of Pydantic's BaseModel
-            origin = get_origin(v) or v
-            if isinstance(origin, type) and issubclass(origin, BaseModel):
-                self.args[k] = v.model_json_schema()
-            else:
-                self.args[k] = TypeAdapter(v).json_schema() or "Any"
-
-    @with_callbacks
-    def __call__(self, *args, **kwargs):
-        return self.func(*args, **kwargs)
+
 
 class ReAct(Module):
     def __init__(self, signature, tools: list[Callable], max_iters=5):
@@ -67,7 +17,7 @@ class ReAct(Module):
         self.signature = signature = ensure_signature(signature)
         self.max_iters = max_iters
 
-        tools = [t if isinstance(t, Tool) or hasattr(t, "input_variable") else Tool(t) for t in tools]
+        tools = [t if isinstance(t, Tool) else Tool(t) for t in tools]
         tools = {tool.name: tool for tool in tools}
 
         inputs = ", ".join([f"`{k}`" for k in signature.input_fields.keys()])
@@ -94,7 +44,7 @@ class ReAct(Module):
         )
 
         for idx, tool in enumerate(tools.values()):
-            args = tool.args if hasattr(tool, "args") else str({tool.input_variable: str})
+            args = getattr(tool, "args")
             desc = (f", whose description is <desc>{tool.desc}</desc>." if tool.desc else ".").replace("\n", "  ")
             desc += f" It takes arguments {args} in JSON format."
             instr.append(f"({idx+1}) {tool.name}{desc}")
diff --git a/dspy/primitives/__init__.py b/dspy/primitives/__init__.py
index 21ee25be..c7cc2c00 100644
--- a/dspy/primitives/__init__.py
+++ b/dspy/primitives/__init__.py
@@ -4,7 +4,7 @@ from dspy.primitives.module import BaseModule
 from dspy.primitives.prediction import Prediction, Completions
 from dspy.primitives.program import Program, Module
 from dspy.primitives.python_interpreter import PythonInterpreter
-
+from dspy.primitives.tool import Tool
 
 __all__ = [
     "assertions",
@@ -15,4 +15,5 @@ __all__ = [
     "Program",
     "Module",
     "PythonInterpreter",
+    "Tool",
 ]
diff --git a/dspy/primitives/tool.py b/dspy/primitives/tool.py
new file mode 100644
index 00000000..56663b43
--- /dev/null
+++ b/dspy/primitives/tool.py
@@ -0,0 +1,126 @@
+import inspect
+from typing import Any, Callable, Optional, get_origin, get_type_hints
+
+from jsonschema import ValidationError, validate
+from pydantic import BaseModel, TypeAdapter
+
+from dspy.utils.callback import with_callbacks
+
+
+class Tool:
+    """Tool class.
+
+    This class is used to simplify the creation of tools for tool calling (function calling) in LLMs. Only supports
+    functions for now.
+    """
+
+    def __init__(
+        self,
+        func: Callable,
+        name: Optional[str] = None,
+        desc: Optional[str] = None,
+        args: Optional[dict[str, Any]] = None,
+        arg_types: Optional[dict[str, Any]] = None,
+        arg_desc: Optional[dict[str, str]] = None,
+    ):
+        """Initialize the Tool class.
+
+        Args:
+            func (Callable): The actual function that is being wrapped by the tool.
+            name (Optional[str], optional): The name of the tool. Defaults to None.
+            desc (Optional[str], optional): The description of the tool. Defaults to None.
+            args (Optional[dict[str, Any]], optional): The args of the tool, represented as a dictionary
+                from arg name to arg's json schema. Defaults to None.
+            arg_types (Optional[dict[str, Any]], optional): The argument types of the tool, represented as a dictionary
+                from arg name to the type of the argument. Defaults to None.
+            arg_desc (Optional[dict[str, str]], optional): Descriptions for each arg, represented as a
+                dictionary from arg name to description string. Defaults to None.
+        """
+        self.func = func
+        self.name = name
+        self.desc = desc
+        self.args = args
+        self.arg_types = arg_types
+        self.arg_desc = arg_desc
+
+        self._parse_function(func, arg_desc)
+
+    def _resolve_pydantic_schema(self, model: type[BaseModel]) -> dict:
+        """Recursively resolve Pydantic model schema, expanding all references."""
+        schema = model.model_json_schema()
+
+        # If there are no definitions to resolve, return the main schema
+        if "$defs" not in schema and "definitions" not in schema:
+            return schema
+
+        def resolve_refs(obj: Any) -> Any:
+            if not isinstance(obj, (dict, list)):
+                return obj
+
+            if isinstance(obj, dict):
+                if "$ref" in obj:
+                    ref_path = obj["$ref"].split("/")[-1]
+                    return resolve_refs(schema["$defs"][ref_path])
+                return {k: resolve_refs(v) for k, v in obj.items()}
+
+            # Must be a list
+            return [resolve_refs(item) for item in obj]
+
+        # Resolve all references in the main schema
+        resolved_schema = resolve_refs(schema)
+        # Remove the $defs key as it's no longer needed
+        resolved_schema.pop("$defs", None)
+        return resolved_schema
+
+    def _parse_function(self, func: Callable, arg_desc: dict[str, str] = None):
+        """Helper method that parses a function to extract the name, description, and args.
+
+        This is a helper function that automatically infers the name, description, and args of the tool from the
+        provided function. In order to make the inference work, the function must have valid type hints.
+        """
+        annotations_func = func if inspect.isfunction(func) or inspect.ismethod(func) else func.__call__
+        name = getattr(func, "__name__", type(func).__name__)
+        desc = getattr(func, "__doc__", None) or getattr(annotations_func, "__doc__", "")
+        args = {}
+        arg_types = {}
+
+        # Use inspect.signature to get all arg names
+        sig = inspect.signature(annotations_func)
+        # Get available type hints
+        available_hints = get_type_hints(annotations_func)
+        # Build a dictionary of arg name -> type (defaulting to Any when missing)
+        hints = {param_name: available_hints.get(param_name, Any) for param_name in sig.parameters.keys()}
+
+        # Process each argument's type to generate its JSON schema.
+        for k, v in hints.items():
+            arg_types[k] = v
+            if k == "return":
+                continue
+            # Check if the type (or its origin) is a subclass of Pydantic's BaseModel
+            origin = get_origin(v) or v
+            if isinstance(origin, type) and issubclass(origin, BaseModel):
+                # Get json schema, and replace $ref with the actual schema
+                v_json_schema = self._resolve_pydantic_schema(v)
+                args[k] = v_json_schema
+            else:
+                args[k] = TypeAdapter(v).json_schema() or "Any"
+            if arg_desc and k in arg_desc:
+                args[k]["description"] = arg_desc[k]
+
+        self.name = self.name or name
+        self.desc = self.desc or desc
+        self.args = self.args or args
+        self.arg_types = self.arg_types or arg_types
+
+    @with_callbacks
+    def __call__(self, **kwargs):
+        for k, v in kwargs.items():
+            if k not in self.args:
+                raise ValueError(f"Arg {k} is not in the tool's args.")
+            try:
+                instance = v.model_dump() if hasattr(v, "model_dump") else v
+                if not self.args[k] == "Any":
+                    validate(instance=instance, schema=self.args[k])
+            except ValidationError as e:
+                raise ValueError(f"Arg {k} is invalid: {e.message}")
+        return self.func(**kwargs)


Test patch:
diff --git a/tests/callback/test_callback.py b/tests/callback/test_callback.py
index ae875388..567b9f87 100644
--- a/tests/callback/test_callback.py
+++ b/tests/callback/test_callback.py
@@ -204,8 +204,8 @@ def test_tool_calls():
             self.tools = [dspy.Tool(tool_1), dspy.Tool(tool_2)]
 
         def forward(self, query: str) -> str:
-            query = self.tools[0](query)
-            return self.tools[1](query)
+            query = self.tools[0](query=query)
+            return self.tools[1](query=query)
 
     module = MyModule()
     result = module("query")
diff --git a/tests/primitives/test_tool.py b/tests/primitives/test_tool.py
new file mode 100644
index 00000000..dc30f2d7
--- /dev/null
+++ b/tests/primitives/test_tool.py
@@ -0,0 +1,135 @@
+import pytest
+from pydantic import BaseModel
+from dspy.primitives.tool import Tool
+from typing import Any, Optional
+
+
+# Test fixtures
+def dummy_function(x: int, y: str) -> str:
+    """A dummy function for testing.
+
+    Args:
+        x: An integer parameter
+        y: A string parameter
+    """
+    return f"{y} {x}"
+
+
+class DummyModel(BaseModel):
+    field1: str
+    field2: int
+
+
+def dummy_with_pydantic(model: DummyModel) -> str:
+    """A dummy function that accepts a Pydantic model."""
+    return f"{model.field1} {model.field2}"
+
+
+class Address(BaseModel):
+    street: str
+    city: str
+    zip_code: str
+    is_primary: bool = False
+
+
+class ContactInfo(BaseModel):
+    email: str
+    phone: Optional[str] = None
+    addresses: list[Address]
+
+
+class UserProfile(BaseModel):
+    user_id: int
+    name: str
+    age: Optional[int] = None
+    contact: ContactInfo
+    tags: list[str] = []
+
+
+def complex_dummy_function(profile: UserProfile, priority: int, notes: Optional[str] = None) -> dict[str, Any]:
+    """Process user profile with complex nested structure.
+
+    Args:
+        profile: User profile containing nested contact and address information
+        priority: Priority level of the processing
+        notes: Optional processing notes
+    """
+    primary_address = next(
+        (addr for addr in profile.contact.addresses if addr.is_primary), profile.contact.addresses[0]
+    )
+
+    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"
+    assert tool.desc == "A test tool"
+    assert tool.args == {"param1": {"type": "string"}}
+    assert callable(tool.func)
+
+
+def test_tool_from_function():
+    tool = Tool(dummy_function)
+
+    assert tool.name == "dummy_function"
+    assert "A 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"
+
+
+def test_tool_from_class():
+    class Foo:
+        def __init__(self, user_id: str):
+            self.user_id = user_id
+
+        def __call__(self, a: int, b: int) -> int:
+            """Add two numbers."""
+            return a + b
+
+    tool = Tool(Foo("123"))
+    assert tool.name == "Foo"
+    assert tool.desc == "Add two numbers."
+    assert tool.args == {"a": {"type": "integer"}, "b": {"type": "integer"}}
+
+
+def test_tool_from_function_with_pydantic():
+    tool = Tool(dummy_with_pydantic)
+
+    assert tool.name == "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"]
+
+
+def test_tool_callable():
+    tool = Tool(dummy_function)
+    result = tool(x=42, y="hello")
+    assert result == "hello 42"
+
+
+def test_tool_with_pydantic_callable():
+    tool = Tool(dummy_with_pydantic)
+    model = DummyModel(field1="test", field2=123)
+    result = tool(model=model)
+    assert result == "test 123"
+
+
+def test_invalid_function_call():
+    tool = Tool(dummy_function)
+    with pytest.raises(ValueError):
+        tool(x="not an integer", y="hello")
+
+
+def test_parameter_desc():
+    tool = Tool(dummy_function, arg_desc={"x": "The x parameter"})
+    assert tool.args["x"]["description"] == "The x parameter"