Make it possible to override JSON schema generation for tools (#1108)

Author: dmontagu

docs/tools.md

@@ -265,19 +265,18 @@ def print_schema(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse
     print(tool.parameters_json_schema)
     """
     {
+        'additionalProperties': False,
         'properties': {
-            'a': {'description': 'apple pie', 'title': 'A', 'type': 'integer'},
-            'b': {'description': 'banana cake', 'title': 'B', 'type': 'string'},
+            'a': {'description': 'apple pie', 'type': 'integer'},
+            'b': {'description': 'banana cake', 'type': 'string'},
             'c': {
                 'additionalProperties': {'items': {'type': 'number'}, 'type': 'array'},
                 'description': 'carrot smoothie',
-                'title': 'C',
                 'type': 'object',
             },
         },
         'required': ['a', 'b', 'c'],
         'type': 'object',
-        'additionalProperties': False,
     }
     """
     return ModelResponse(parts=[TextPart('foobar')])
@@ -328,9 +327,9 @@ print(test_model.last_model_request_parameters.function_tools)
         description='This is a Foobar',
         parameters_json_schema={
             'properties': {
-                'x': {'title': 'X', 'type': 'integer'},
-                'y': {'title': 'Y', 'type': 'string'},
-                'z': {'default': 3.14, 'title': 'Z', 'type': 'number'},
+                'x': {'type': 'integer'},
+                'y': {'type': 'string'},
+                'z': {'default': 3.14, 'type': 'number'},
             },
             'required': ['x', 'y'],
             'title': 'Foobar',
@@ -432,16 +431,12 @@ print(test_model.last_model_request_parameters.function_tools)
         name='greet',
         description='',
         parameters_json_schema={
+            'additionalProperties': False,
             'properties': {
-                'name': {
-                    'title': 'Name',
-                    'type': 'string',
-                    'description': 'Name of the human to greet.',
-                }
+                'name': {'type': 'string', 'description': 'Name of the human to greet.'}
             },
             'required': ['name'],
             'type': 'object',
-            'additionalProperties': False,
         },
         outer_typed_dict_key=None,
     )

pydantic_ai_slim/pydantic_ai/_pydantic.py

@@ -44,6 +44,7 @@ def function_schema(  # noqa: C901
     takes_ctx: bool,
     docstring_format: DocstringFormat,
     require_parameter_descriptions: bool,
+    schema_generator: type[GenerateJsonSchema],
 ) -> FunctionSchema:
     """Build a Pydantic validator and JSON schema from a tool function.
 
@@ -52,6 +53,7 @@ def function_schema(  # noqa: C901
         takes_ctx: Whether the function takes a `RunContext` first argument.
         docstring_format: The docstring format to use.
         require_parameter_descriptions: Whether to require descriptions for all tool function parameters.
+        schema_generator: The JSON schema generator class to use.
 
     Returns:
         A `FunctionSchema` instance.
@@ -150,14 +152,12 @@ def function_schema(  # noqa: C901
     )
     # PluggableSchemaValidator is api compatible with SchemaValidator
     schema_validator = cast(SchemaValidator, schema_validator)
-    json_schema = GenerateJsonSchema().generate(schema)
+    json_schema = schema_generator().generate(schema)
 
     # workaround for https://github.com/pydantic/pydantic/issues/10785
-    # if we build a custom TypeDict schema (matches when `single_arg_name is None`), we manually set
+    # if we build a custom TypedDict schema (matches when `single_arg_name is None`), we manually set
     # `additionalProperties` in the JSON Schema
-    if single_arg_name is None:
-        json_schema['additionalProperties'] = bool(var_kwargs_schema)
-    elif not description:
+    if single_arg_name is not None and not description:
         # if the tool description is not set, and we have a single parameter, take the description from that
         # and set it on the tool
         description = json_schema.pop('description', None)
@@ -218,6 +218,7 @@ def _build_schema(
     td_schema = core_schema.typed_dict_schema(
         fields,
         config=core_config,
+        total=var_kwargs_schema is None,
         extras_schema=gen_schema.generate_schema(var_kwargs_schema) if var_kwargs_schema else None,
     )
     return td_schema, None

pydantic_ai_slim/pydantic_ai/agent.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, cast, final, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
+from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import TypeGuard, TypeVar, deprecated
 
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
@@ -31,6 +32,7 @@
 from .tools import (
     AgentDepsT,
     DocstringFormat,
+    GenerateToolJsonSchema,
     RunContext,
     Tool,
     ToolFuncContext,
@@ -936,6 +938,7 @@ def tool(
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ...
 
     def tool(
@@ -948,6 +951,7 @@ def tool(
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -989,6 +993,7 @@ async def spam(ctx: RunContext[str], y: float) -> float:
             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 func is None:
 
@@ -997,15 +1002,22 @@ def tool_decorator(
             ) -> ToolFuncContext[AgentDepsT, ToolParams]:
                 # noinspection PyTypeChecker
                 self._register_function(
-                    func_, True, name, retries, prepare, docstring_format, require_parameter_descriptions
+                    func_,
+                    True,
+                    name,
+                    retries,
+                    prepare,
+                    docstring_format,
+                    require_parameter_descriptions,
+                    schema_generator,
                 )
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
             self._register_function(
-                func, True, name, retries, prepare, docstring_format, require_parameter_descriptions
+                func, True, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
             )
             return func
 
@@ -1022,6 +1034,7 @@ def tool_plain(
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
     def tool_plain(
@@ -1034,6 +1047,7 @@ def tool_plain(
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -1075,20 +1089,28 @@ async def spam(ctx: RunContext[str]) -> float:
             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 func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
                 self._register_function(
-                    func_, False, name, retries, prepare, docstring_format, require_parameter_descriptions
+                    func_,
+                    False,
+                    name,
+                    retries,
+                    prepare,
+                    docstring_format,
+                    require_parameter_descriptions,
+                    schema_generator,
                 )
                 return func_
 
             return tool_decorator
         else:
             self._register_function(
-                func, False, name, retries, prepare, docstring_format, require_parameter_descriptions
+                func, False, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
             )
             return func
 
@@ -1101,6 +1123,7 @@ def _register_function(
         prepare: ToolPrepareFunc[AgentDepsT] | None,
         docstring_format: DocstringFormat,
         require_parameter_descriptions: bool,
+        schema_generator: type[GenerateJsonSchema],
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
@@ -1112,6 +1135,7 @@ def _register_function(
             prepare=prepare,
             docstring_format=docstring_format,
             require_parameter_descriptions=require_parameter_descriptions,
+            schema_generator=schema_generator,
         )
         self._register_tool(tool)
 

pydantic_ai_slim/pydantic_ai/tools.py

@@ -7,7 +7,8 @@
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
 
 from pydantic import ValidationError
-from pydantic_core import SchemaValidator
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypeVar
 
 from . import _pydantic, _utils, messages as _messages, models
@@ -142,6 +143,22 @@ def hitchhiker(ctx: RunContext[int], answer: str) -> str:
 A = TypeVar('A')
 
 
+class GenerateToolJsonSchema(GenerateJsonSchema):
+    def typed_dict_schema(self, schema: core_schema.TypedDictSchema) -> JsonSchemaValue:
+        s = super().typed_dict_schema(schema)
+        total = schema.get('total')
+        if total is not None:
+            s['additionalProperties'] = not total
+        return s
+
+    def _named_required_fields_schema(self, named_required_fields: Sequence[tuple[str, bool, Any]]) -> JsonSchemaValue:
+        # Remove largely-useless property titles
+        s = super()._named_required_fields_schema(named_required_fields)
+        for p in s.get('properties', {}):
+            s['properties'][p].pop('title', None)
+        return s
+
+
 @dataclass(init=False)
 class Tool(Generic[AgentDepsT]):
     """A tool function for an agent."""
@@ -176,6 +193,7 @@ def __init__(
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ):
         """Create a new tool instance.
 
@@ -225,11 +243,14 @@ async def prep_my_tool(
             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. Defaults to `GenerateToolJsonSchema`.
         """
         if takes_ctx is None:
             takes_ctx = _pydantic.takes_ctx(function)
 
-        f = _pydantic.function_schema(function, takes_ctx, docstring_format, require_parameter_descriptions)
+        f = _pydantic.function_schema(
+            function, takes_ctx, docstring_format, require_parameter_descriptions, schema_generator
+        )
         self.function = function
         self.takes_ctx = takes_ctx
         self.max_retries = max_retries

tests/graph/test_graph.py

@@ -393,7 +393,7 @@ async def run(self, ctx: GraphRunContext) -> End[None]:
 
         assert isinstance(n, BaseNode)
         n = await run.next()
-        assert n == snapshot(End(None))
+        assert n == snapshot(End(data=None))
 
         with pytest.raises(TypeError, match=r'`next` must be called with a `BaseNode` instance, got End\(data=None\).'):
             await run.next()

tests/graph/test_persistence.py

@@ -287,7 +287,7 @@ async def run(self, ctx: GraphRunContext) -> End[int]:
     node = Foo()
     async with graph.iter(node, persistence=sp) as run:
         end = await run.next()
-        assert end == snapshot(End(123))
+        assert end == snapshot(End(data=123))
 
     msg = "Incorrect snapshot status 'success', must be 'created' or 'pending'."
     with pytest.raises(GraphNodeStatusError, match=msg):