Support tool and toolcall in input and output field types

Author: chenmoneygithub

dspy/adapters/base.py

@@ -352,4 +352,4 @@ def parse(self, signature: Type[Signature], completion: str) -> dict[str, Any]:
         Returns:
             A dictionary of the output fields.
         """
-        raise NotImplementedError
+        raise NotImplementedError

dspy/adapters/types/base_type.py

@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Any
+from typing import Any, Union, get_args, get_origin
 
 import json_repair
 import pydantic
@@ -26,12 +26,42 @@ def format(self) -> list[dict[str, Any]]:
         ```
     """
 
-    def format(self) -> list[dict[str, Any]]:
+    def format(self) -> Union[list[dict[str, Any]], str]:
         raise NotImplementedError
 
+    @classmethod
+    def description(cls) -> str:
+        """Description of the custom type"""
+        return ""
+
+    @classmethod
+    def extract_custom_type_from_annotation(cls, annotation):
+        """Extract all custom types from the annotation.
+
+        This is used to extract all custom types from the annotation of a field, while the annotation can
+        have arbitrary level of nesting. For example, we detect `Tool` is in `list[dict[str, Tool]]`.
+        """
+        # Direct match
+        if isinstance(annotation, type) and issubclass(annotation, cls):
+            return [annotation]
+
+        origin = get_origin(annotation)
+        if origin is None:
+            return []
+
+        result = []
+        # Recurse into all type args
+        for arg in get_args(annotation):
+            result.extend(cls.extract_custom_type_from_annotation(arg))
+
+        return result
+
     @pydantic.model_serializer()
     def serialize_model(self):
-        return f"{CUSTOM_TYPE_START_IDENTIFIER}{self.format()}{CUSTOM_TYPE_END_IDENTIFIER}"
+        formatted = self.format()
+        if isinstance(formatted, list):
+            return f"{CUSTOM_TYPE_START_IDENTIFIER}{self.format()}{CUSTOM_TYPE_END_IDENTIFIER}"
+        return formatted
 
 
 def split_message_content_for_custom_types(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:

dspy/adapters/types/tool.py

@@ -19,6 +19,14 @@ class Tool(BaseType):
     functions for now.
     """
 
+    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
+    has_kwargs: bool = False
+
     def __init__(
         self,
         func: Callable,
@@ -56,15 +64,7 @@ def foo(x: int, y: str = "hello"):
         # Expected output: {'x': {'type': 'integer'}, 'y': {'type': 'string', 'default': 'hello'}}
         ```
         """
-        super().__init__()  # Initialize the Pydantic BaseModel
-        self.func = func
-        self.name = name
-        self.desc = desc
-        self.args = args
-        self.arg_types = arg_types
-        self.arg_desc = arg_desc
-        self.has_kwargs = False
-
+        super().__init__(func=func, name=name, desc=desc, args=args, arg_types=arg_types, arg_desc=arg_desc)
         self._parse_function(func, arg_desc)
 
     def _parse_function(self, func: Callable, arg_desc: Optional[dict[str, str]] = None):
@@ -141,17 +141,7 @@ def _validate_and_parse_args(self, **kwargs):
         return parsed_kwargs
 
     def format(self):
-        return [
-            {
-                "type": "function",
-                "function": {
-                    "name": self.name,
-                    "description": self.desc,
-                    "parameters": self.args,
-                    "requirements": "Arguments must be provided in JSON format.",
-                },
-            }
-        ]
+        return str(self)
 
     @with_callbacks
     def __call__(self, **kwargs):
@@ -200,6 +190,13 @@ class ToolCall(BaseType):
     name: str
     args: dict[str, Any]
 
+    @classmethod
+    def description(cls) -> str:
+        return (
+            "Tool call information, including the name of the tool and the arguments to be passed to it. "
+            "Arguments must be provided in JSON format."
+        )
+
 
 def resolve_json_schema_reference(schema: dict) -> dict:
     """Recursively resolve json model schema, expanding all references."""

dspy/adapters/utils.py

@@ -10,6 +10,7 @@
 from pydantic import TypeAdapter
 from pydantic.fields import FieldInfo
 
+from dspy.adapters.types.base_type import BaseType
 from dspy.signatures.utils import get_dspy_field_type
 
 
@@ -200,7 +201,14 @@ def get_field_description_string(fields: dict) -> str:
     for idx, (k, v) in enumerate(fields.items()):
         field_message = f"{idx + 1}. `{k}`"
         field_message += f" ({get_annotation_name(v.annotation)})"
-        field_message += f": {v.json_schema_extra['desc']}" if v.json_schema_extra["desc"] != f"${{{k}}}" else ""
+        desc = v.json_schema_extra["desc"] if v.json_schema_extra["desc"] != f"${{{k}}}" else ""
+
+        custom_types = BaseType.extract_custom_type_from_annotation(v.annotation)
+        for custom_type in custom_types:
+            if len(custom_type.description()) > 0:
+                desc += f"\n    Type description of {get_annotation_name(custom_type)}: {custom_type.description()}"
+
+        field_message += f": {desc}"
         field_message += (
             f"\nConstraints: {v.json_schema_extra['constraints']}" if v.json_schema_extra.get("constraints") else ""
         )

tests/adapters/test_chat_adapter.py

@@ -5,6 +5,7 @@
 
 import dspy
 import pydantic
+import json
 
 
 @pytest.mark.parametrize(
@@ -213,3 +214,40 @@ class MySignature(dspy.Signature):
 
     # The query image is formatted in the last user message
     assert {"type": "image_url", "image_url": {"url": "https://example.com/image4.jpg"}} in messages[-1]["content"]
+
+
+def test_chat_adapter_with_tool():
+    class MySignature(dspy.Signature):
+        """Answer question with the help of the tools"""
+
+        question: str = dspy.InputField()
+        tools: list[dspy.Tool] = dspy.InputField()
+        answer: str = dspy.OutputField()
+        tool_calls: list[dspy.ToolCall] = dspy.OutputField()
+
+    def get_weather(city: str) -> str:
+        """Get the weather for a city"""
+        return f"The weather in {city} is sunny"
+
+    def get_population(country: str, year: int) -> str:
+        """Get the population for a country"""
+        return f"The population of {country} in {year} is 1000000"
+
+    tools = [dspy.Tool(get_weather), dspy.Tool(get_population)]
+
+    adapter = dspy.ChatAdapter()
+    messages = adapter.format(MySignature, [], {"question": "What is the weather in Tokyo?", "tools": tools})
+
+    assert len(messages) == 2
+
+    # The output field type description should be included in the system message even if the output field is nested
+    assert dspy.ToolCall.description() in messages[0]["content"]
+
+    # The user message should include the question and the tools
+    assert "What is the weather in Tokyo?" in messages[1]["content"]
+    assert "get_weather" in messages[1]["content"]
+    assert "get_population" in messages[1]["content"]
+
+    # Tool arguments format should be included in the user message
+    assert "{'city': {'type': 'string'}}" in messages[1]["content"]
+    assert "{'country': {'type': 'string'}, 'year': {'type': 'integer'}}" in messages[1]["content"]

tests/adapters/test_json_adapter.py

@@ -5,6 +5,7 @@
 from litellm.utils import Choices, Message, ModelResponse
 
 import dspy
+import json
 
 
 def test_json_adapter_passes_structured_output_when_supported_by_model():
@@ -337,3 +338,40 @@ class MySignature(dspy.Signature):
 
     # The query image is formatted in the last user message
     assert {"type": "image_url", "image_url": {"url": "https://example.com/image4.jpg"}} in messages[-1]["content"]
+
+
+def test_json_adapter_with_tool():
+    class MySignature(dspy.Signature):
+        """Answer question with the help of the tools"""
+
+        question: str = dspy.InputField()
+        tools: list[dspy.Tool] = dspy.InputField()
+        answer: str = dspy.OutputField()
+        tool_calls: list[dspy.ToolCall] = dspy.OutputField()
+
+    def get_weather(city: str) -> str:
+        """Get the weather for a city"""
+        return f"The weather in {city} is sunny"
+
+    def get_population(country: str, year: int) -> str:
+        """Get the population for a country"""
+        return f"The population of {country} in {year} is 1000000"
+
+    tools = [dspy.Tool(get_weather), dspy.Tool(get_population)]
+
+    adapter = dspy.JSONAdapter()
+    messages = adapter.format(MySignature, [], {"question": "What is the weather in Tokyo?", "tools": tools})
+
+    assert len(messages) == 2
+
+    # The output field type description should be included in the system message even if the output field is nested
+    assert dspy.ToolCall.description() in messages[0]["content"]
+
+    # The user message should include the question and the tools
+    assert "What is the weather in Tokyo?" in messages[1]["content"]
+    assert "get_weather" in messages[1]["content"]
+    assert "get_population" in messages[1]["content"]
+
+    # Tool arguments format should be included in the user message
+    assert "{'city': {'type': 'string'}}" in messages[1]["content"]
+    assert "{'country': {'type': 'string'}, 'year': {'type': 'integer'}}" in messages[1]["content"]

tests/primitives/test_tool.py renamed to tests/adapters/test_tool.py

@@ -1,10 +1,12 @@
 import asyncio
 from typing import Any, Optional
 
+import dspy
 import pytest
 from pydantic import BaseModel
 
-from dspy.primitives.tool import Tool
+from dspy.adapters.types.tool import Tool
+from unittest import mock
 
 
 # Test fixtures