Add docstring_format, require_parameter_descriptions, schema_generator to FunctionToolset (#2601)

Author: g-eoj

pydantic_ai_slim/pydantic_ai/toolsets/function.py

@@ -33,26 +33,43 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
     See [toolset docs](../toolsets.md#function-toolset) for more information.
     """
 
-    max_retries: int
     tools: dict[str, Tool[Any]]
+    max_retries: int
     _id: str | None
+    docstring_format: DocstringFormat
+    require_parameter_descriptions: bool
+    schema_generator: type[GenerateJsonSchema]
 
     def __init__(
         self,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
         *,
         max_retries: int = 1,
         id: str | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ):
         """Build a new function toolset.
 
         Args:
             tools: The tools to add to the toolset.
             max_retries: The maximum number of retries for each tool during a run.
-            id: An optional unique ID for the toolset. A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the toolset's activities within the workflow.
+            id: An optional unique ID for the toolset. A toolset needs to have an ID in order to be used in a durable execution environment like Temporal,
+                in which case the ID will be used to identify the toolset's activities within the workflow.
+            docstring_format: Format of tool docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+                Applies to all tools, unless overridden when adding a tool.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+                Applies to all tools, unless overridden when adding a tool.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                Applies to all tools, unless overridden when adding a tool.
         """
         self.max_retries = max_retries
         self._id = id
+        self.docstring_format = docstring_format
+        self.require_parameter_descriptions = require_parameter_descriptions
+        self.schema_generator = schema_generator
 
         self.tools = {}
         for tool in tools:
@@ -76,9 +93,9 @@ def tool(
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
         requires_approval: bool = False,
     ) -> Callable[[ToolFuncEither[AgentDepsT, ToolParams]], ToolFuncEither[AgentDepsT, ToolParams]]: ...
@@ -91,9 +108,9 @@ def tool(
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
         requires_approval: bool = False,
     ) -> Any:
@@ -137,9 +154,11 @@ async def spam(ctx: RunContext[str], y: float) -> float:
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
-                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
-            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
-            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                If `None`, the default value is determined by the toolset.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing.
+                If `None`, the default value is determined by the toolset.
+            schema_generator: The JSON schema generator class to use for this tool.
+                If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
@@ -173,9 +192,9 @@ def add_function(
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
         requires_approval: bool = False,
     ) -> None:
@@ -196,14 +215,23 @@ def add_function(
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
-                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
-            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
-            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                If `None`, the default value is determined by the toolset.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing.
+                If `None`, the default value is determined by the toolset.
+            schema_generator: The JSON schema generator class to use for this tool.
+                If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
         """
+        if docstring_format is None:
+            docstring_format = self.docstring_format
+        if require_parameter_descriptions is None:
+            require_parameter_descriptions = self.require_parameter_descriptions
+        if schema_generator is None:
+            schema_generator = self.schema_generator
+
         tool = Tool[AgentDepsT](
             func,
             takes_ctx=takes_ctx,

tests/test_toolsets.py

@@ -130,6 +130,29 @@ def subtract(a: int, b: int) -> int:
     assert await bar_toolset.handle_call(ToolCallPart(tool_name='bar_add', args={'a': 1, 'b': 2})) == 3
 
 
+async def test_function_toolset_with_defaults():
+    defaults_toolset = FunctionToolset[None](require_parameter_descriptions=True)
+
+    with pytest.raises(
+        UserError,
+        match=re.escape('Missing parameter descriptions for'),
+    ):
+
+        @defaults_toolset.tool
+        def add(a: int, b: int) -> int:
+            """Add two numbers"""
+            return a + b  # pragma: no cover
+
+
+async def test_function_toolset_with_defaults_overridden():
+    defaults_toolset = FunctionToolset[None](require_parameter_descriptions=True)
+
+    @defaults_toolset.tool(require_parameter_descriptions=False)
+    def subtract(a: int, b: int) -> int:
+        """Subtract two numbers"""
+        return a - b  # pragma: no cover
+
+
 async def test_prepared_toolset_user_error_add_new_tools():
     """Test that PreparedToolset raises UserError when prepare function tries to add new tools."""
     context = build_run_context(None)