Feature/add openai strict mode (#1304)

Co-authored-by: Moritz Scherrmann <[email protected]>
Co-authored-by: David Montague <[email protected]>
Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: 4 people

docs/tools.md

@@ -353,6 +353,7 @@ print(test_model.last_model_request_parameters.function_tools)
             'type': 'object',
         },
         outer_typed_dict_key=None,
+        strict=None,
     )
 ]
 """
@@ -456,6 +457,7 @@ print(test_model.last_model_request_parameters.function_tools)
             'type': 'object',
         },
         outer_typed_dict_key=None,
+        strict=None,
     )
 ]
 """

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -311,6 +311,7 @@ async def _make_request(
             return self._result
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
+        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
         model_response, request_usage = await ctx.deps.model.request(
             ctx.state.message_history, model_settings, model_request_parameters
         )

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -50,7 +50,11 @@ def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
     if schema.get('type') == 'object':
         return schema
     elif schema.get('$ref') is not None:
-        return schema.get('$defs', {}).get(schema['$ref'][8:])  # This removes the initial "#/$defs/".
+        maybe_result = schema.get('$defs', {}).get(schema['$ref'][8:])  # This removes the initial "#/$defs/".
+
+        if "'$ref': '#/$defs/" in str(maybe_result):
+            return schema  # We can't remove the $defs because the schema contains other references
+        return maybe_result
     else:
         raise UserError('Schema must be an object')
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -940,6 +940,7 @@ def tool(
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
     ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ...
 
     def tool(
@@ -953,6 +954,7 @@ def tool(
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -995,6 +997,8 @@ async def spam(ctx: RunContext[str], y: float) -> float:
                 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`.
+            strict: Whether to enforce JSON schema compliance (only affects OpenAI).
+                See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
         """
         if func is None:
 
@@ -1011,14 +1015,23 @@ def tool_decorator(
                     docstring_format,
                     require_parameter_descriptions,
                     schema_generator,
+                    strict,
                 )
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
             self._register_function(
-                func, True, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
+                func,
+                True,
+                name,
+                retries,
+                prepare,
+                docstring_format,
+                require_parameter_descriptions,
+                schema_generator,
+                strict,
             )
             return func
 
@@ -1036,6 +1049,7 @@ def tool_plain(
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
     def tool_plain(
@@ -1049,6 +1063,7 @@ def tool_plain(
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -1091,6 +1106,8 @@ async def spam(ctx: RunContext[str]) -> float:
                 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`.
+            strict: Whether to enforce JSON schema compliance (only affects OpenAI).
+                See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
         """
         if func is None:
 
@@ -1105,13 +1122,22 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams
                     docstring_format,
                     require_parameter_descriptions,
                     schema_generator,
+                    strict,
                 )
                 return func_
 
             return tool_decorator
         else:
             self._register_function(
-                func, False, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
+                func,
+                False,
+                name,
+                retries,
+                prepare,
+                docstring_format,
+                require_parameter_descriptions,
+                schema_generator,
+                strict,
             )
             return func
 
@@ -1125,6 +1151,7 @@ def _register_function(
         docstring_format: DocstringFormat,
         require_parameter_descriptions: bool,
         schema_generator: type[GenerateJsonSchema],
+        strict: bool | None,
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
@@ -1137,6 +1164,7 @@ def _register_function(
             docstring_format=docstring_format,
             require_parameter_descriptions=require_parameter_descriptions,
             schema_generator=schema_generator,
+            strict=strict,
         )
         self._register_tool(tool)
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -274,6 +274,15 @@ async def request_stream(
         # noinspection PyUnreachableCode
         yield  # pragma: no cover
 
+    def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
+        """Customize the request parameters for the model.
+
+        This method can be overridden by subclasses to modify the request parameters before sending them to the model.
+        In particular, this method can be used to make modifications to the generated tool JSON schemas if necessary
+        for vendor/model-specific reasons.
+        """
+        return model_request_parameters
+
     @property
     @abstractmethod
     def model_name(self) -> str:

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -5,7 +5,7 @@
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Annotated, Any, Literal, Protocol, Union, cast
 from uuid import uuid4
@@ -152,6 +152,16 @@ async def request_stream(
         ) as http_response:
             yield await self._process_streamed_response(http_response)
 
+    def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
+        def _customize_tool_def(t: ToolDefinition):
+            return replace(t, parameters_json_schema=_GeminiJsonSchema(t.parameters_json_schema).simplify())
+
+        return ModelRequestParameters(
+            function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
+            allow_text_result=model_request_parameters.allow_text_result,
+            result_tools=[_customize_tool_def(tool) for tool in model_request_parameters.result_tools],
+        )
+
     @property
     def model_name(self) -> GeminiModelName:
         """The model name."""
@@ -640,7 +650,7 @@ class _GeminiFunction(TypedDict):
 
 
 def _function_from_abstract_tool(tool: ToolDefinition) -> _GeminiFunction:
-    json_schema = _GeminiJsonSchema(tool.parameters_json_schema).simplify()
+    json_schema = tool.parameters_json_schema
     f = _GeminiFunction(name=tool.name, description=tool.description)
     if json_schema.get('properties'):
         f['parameters'] = json_schema