Support dspy.Tool as input field type and dspy.ToolCall as output field type (#8242)

* init

* Support tool and toolcall in input and output field types

* move file

* increment

* better ways

* return a list from lm for backward compatibility

* fix tests

* better name

* fix tests

* fix invalid import path

Author: chenmoneygithub

dspy/__init__.py

@@ -8,7 +8,7 @@
 
 from dspy.evaluate import Evaluate  # isort: skip
 from dspy.clients import *  # isort: skip
-from dspy.adapters import Adapter, ChatAdapter, JSONAdapter, TwoStepAdapter, Image, Audio, History, BaseType  # isort: skip
+from dspy.adapters import Adapter, ChatAdapter, JSONAdapter, TwoStepAdapter, Image, Audio, History, BaseType, Tool, ToolCalls  # isort: skip
 from dspy.utils.logging_utils import configure_dspy_loggers, disable_logging, enable_logging
 from dspy.utils.asyncify import asyncify
 from dspy.utils.saving import load

dspy/adapters/__init__.py

@@ -2,7 +2,7 @@
 from dspy.adapters.chat_adapter import ChatAdapter
 from dspy.adapters.json_adapter import JSONAdapter
 from dspy.adapters.two_step_adapter import TwoStepAdapter
-from dspy.adapters.types import History, Image, Audio, BaseType
+from dspy.adapters.types import History, Image, Audio, BaseType, Tool, ToolCalls
 
 __all__ = [
     "Adapter",
@@ -13,4 +13,6 @@
     "Audio",
     "JSONAdapter",
     "TwoStepAdapter",
+    "Tool",
+    "ToolCalls",
 ]

dspy/adapters/base.py

@@ -1,10 +1,17 @@
-from typing import TYPE_CHECKING, Any, Optional, Type
+import logging
+from typing import TYPE_CHECKING, Any, Optional, Type, get_origin
+
+import json_repair
+import litellm
 
 from dspy.adapters.types import History
 from dspy.adapters.types.base_type import split_message_content_for_custom_types
+from dspy.adapters.types.tool import Tool, ToolCalls
 from dspy.signatures.signature import Signature
 from dspy.utils.callback import BaseCallback, with_callbacks
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from dspy.clients.lm import LM
 
@@ -20,18 +27,78 @@ def __init_subclass__(cls, **kwargs) -> None:
         cls.format = with_callbacks(cls.format)
         cls.parse = with_callbacks(cls.parse)
 
-    def _call_post_process(self, outputs: list[dict[str, Any]], signature: Type[Signature]) -> list[dict[str, Any]]:
+    def _call_preprocess(
+        self,
+        lm: "LM",
+        lm_kwargs: dict[str, Any],
+        signature: Type[Signature],
+        inputs: dict[str, Any],
+        use_native_function_calling: bool = False,
+    ) -> dict[str, Any]:
+        if use_native_function_calling:
+            tool_call_input_field_name = self._get_tool_call_input_field_name(signature)
+            tool_call_output_field_name = self._get_tool_call_output_field_name(signature)
+
+            if tool_call_output_field_name and tool_call_input_field_name is None:
+                raise ValueError(
+                    f"You provided an output field {tool_call_output_field_name} to receive the tool calls information, "
+                    "but did not provide any tools as the input. Please provide a list of tools as the input by adding an "
+                    "input field with type `list[dspy.Tool]`."
+                )
+
+            if tool_call_output_field_name and litellm.supports_function_calling(model=lm.model):
+                tools = inputs[tool_call_input_field_name]
+                tools = tools if isinstance(tools, list) else [tools]
+
+                litellm_tools = []
+                for tool in tools:
+                    litellm_tools.append(tool.format_as_litellm_function_call())
+
+                lm_kwargs["tools"] = litellm_tools
+
+                signature_for_native_function_calling = signature.delete(tool_call_output_field_name)
+
+                return signature_for_native_function_calling
+
+        return signature
+
+    def _call_postprocess(
+        self,
+        signature: Type[Signature],
+        outputs: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
         values = []
 
+        tool_call_output_field_name = self._get_tool_call_output_field_name(signature)
+
         for output in outputs:
             output_logprobs = None
+            tool_calls = None
+            text = output
 
             if isinstance(output, dict):
-                output, output_logprobs = output["text"], output["logprobs"]
-
-            value = self.parse(signature, output)
-
-            if output_logprobs is not None:
+                text = output["text"]
+                output_logprobs = output.get("logprobs")
+                tool_calls = output.get("tool_calls")
+
+            if text:
+                value = self.parse(signature, text)
+            else:
+                value = {}
+                for field_name in signature.output_fields.keys():
+                    value[field_name] = None
+
+            if tool_calls and tool_call_output_field_name:
+                tool_calls = [
+                    {
+                        "name": v["function"]["name"],
+                        "args": json_repair.loads(v["function"]["arguments"]),
+                    }
+                    for v in tool_calls
+                ]
+                value[tool_call_output_field_name] = ToolCalls.from_dict_list(tool_calls)
+
+            if output_logprobs:
                 value["logprobs"] = output_logprobs
 
             values.append(value)
@@ -46,10 +113,11 @@ def __call__(
         demos: list[dict[str, Any]],
         inputs: dict[str, Any],
     ) -> list[dict[str, Any]]:
-        inputs = self.format(signature, demos, inputs)
+        processed_signature = self._call_preprocess(lm, lm_kwargs, signature, inputs)
+        inputs = self.format(processed_signature, demos, inputs)
 
         outputs = lm(messages=inputs, **lm_kwargs)
-        return self._call_post_process(outputs, signature)
+        return self._call_postprocess(signature, outputs)
 
     async def acall(
         self,
@@ -59,10 +127,11 @@ async def acall(
         demos: list[dict[str, Any]],
         inputs: dict[str, Any],
     ) -> list[dict[str, Any]]:
-        inputs = self.format(signature, demos, inputs)
+        processed_signature = self._call_preprocess(lm, lm_kwargs, signature, inputs)
+        inputs = self.format(processed_signature, demos, inputs)
 
         outputs = await lm.acall(messages=inputs, **lm_kwargs)
-        return self._call_post_process(outputs, signature)
+        return self._call_postprocess(signature, outputs)
 
     def format(
         self,
@@ -297,6 +366,22 @@ def _get_history_field_name(self, signature: Type[Signature]) -> bool:
                 return name
         return None
 
+    def _get_tool_call_input_field_name(self, signature: Type[Signature]) -> bool:
+        for name, field in signature.input_fields.items():
+            # Look for annotation `list[dspy.Tool]` or `dspy.Tool`
+            origin = get_origin(field.annotation)
+            if origin is list and field.annotation.__args__[0] == Tool:
+                return name
+            if field.annotation == Tool:
+                return name
+        return None
+
+    def _get_tool_call_output_field_name(self, signature: Type[Signature]) -> bool:
+        for name, field in signature.output_fields.items():
+            if field.annotation == ToolCalls:
+                return name
+        return None
+
     def format_conversation_history(
         self,
         signature: Type[Signature],
@@ -352,4 +437,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/json_adapter.py

@@ -77,6 +77,16 @@ def __call__(
                     f"`response_format` argument. Original error: {e}"
                 ) from e
 
+    def _call_preprocess(
+        self,
+        lm: "LM",
+        lm_kwargs: dict[str, Any],
+        signature: Type[Signature],
+        inputs: dict[str, Any],
+        use_native_function_calling: bool = True,
+    ) -> dict[str, Any]:
+        return super()._call_preprocess(lm, lm_kwargs, signature, inputs, use_native_function_calling)
+
     def format_field_structure(self, signature: Type[Signature]) -> str:
         parts = []
         parts.append("All interactions will be structured in the following way, with the appropriate values filled in.")

dspy/adapters/two_step_adapter.py

@@ -1,7 +1,10 @@
 from typing import Any, Optional, Type
 
+import json_repair
+
 from dspy.adapters.base import Adapter
 from dspy.adapters.chat_adapter import ChatAdapter
+from dspy.adapters.types import ToolCalls
 from dspy.adapters.utils import get_field_description_string
 from dspy.clients import LM
 from dspy.signatures.field import InputField
@@ -115,11 +118,16 @@ async def acall(
 
         values = []
 
+        tool_call_output_field_name = self._get_tool_call_output_field_name(signature)
         for output in outputs:
             output_logprobs = None
+            tool_calls = None
+            text = output
 
             if isinstance(output, dict):
-                output, output_logprobs = output["text"], output["logprobs"]
+                text = output["text"]
+                output_logprobs = output.get("logprobs")
+                tool_calls = output.get("tool_calls")
 
             try:
                 # Call the smaller LM to extract structured data from the raw completion text with ChatAdapter
@@ -128,13 +136,23 @@ async def acall(
                     lm_kwargs={},
                     signature=extractor_signature,
                     demos=[],
-                    inputs={"text": output},
+                    inputs={"text": text},
                 )
                 value = value[0]
 
             except Exception as e:
                 raise ValueError(f"Failed to parse response from the original completion: {output}") from e
 
+            if tool_calls and tool_call_output_field_name:
+                tool_calls = [
+                    {
+                        "name": v["function"]["name"],
+                        "args": json_repair.loads(v["function"]["arguments"]),
+                    }
+                    for v in tool_calls
+                ]
+                value[tool_call_output_field_name] = ToolCalls.from_dict_list(tool_calls)
+
             if output_logprobs is not None:
                 value["logprobs"] = output_logprobs
 

dspy/adapters/types/__init__.py

@@ -2,5 +2,6 @@
 from dspy.adapters.types.image import Image
 from dspy.adapters.types.audio import Audio
 from dspy.adapters.types.base_type import BaseType
+from dspy.adapters.types.tool import Tool, ToolCalls
 
-__all__ = ["History", "Image", "Audio", "BaseType"]
+__all__ = ["History", "Image", "Audio", "BaseType", "Tool", "ToolCalls"]

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]]: