pydantic-ai-slim 0.0.54__py3-none-any.whl → 0.1.0__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -30,6 +30,7 @@ from ..messages import (
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
+    VideoUrl,
 )
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -39,7 +40,9 @@ from . import (
     StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
+    get_user_agent,
 )
+from ._json_schema import JsonSchema, WalkJsonSchema
 
 try:
     from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream, NotGiven
@@ -251,15 +254,12 @@ class OpenAIModel(Model):
         # standalone function to make it easier to override
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
-        elif not model_request_parameters.allow_text_result:
+        elif not model_request_parameters.allow_text_output:
             tool_choice = 'required'
         else:
             tool_choice = 'auto'
 
-        openai_messages: list[chat.ChatCompletionMessageParam] = []
-        for m in messages:
-            async for msg in self._map_message(m):
-                openai_messages.append(msg)
+        openai_messages = await self._map_messages(messages)
 
         try:
             return await self.client.chat.completions.create(
@@ -282,6 +282,7 @@ class OpenAIModel(Model):
                 logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
                 reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN),
                 user=model_settings.get('openai_user', NOT_GIVEN),
+                extra_headers={'User-Agent': get_user_agent()},
             )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
@@ -315,35 +316,40 @@ class OpenAIModel(Model):
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
-    async def _map_message(self, message: ModelMessage) -> AsyncIterable[chat.ChatCompletionMessageParam]:
+    async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
-        if isinstance(message, ModelRequest):
-            async for item in self._map_user_message(message):
-                yield item
-        elif isinstance(message, ModelResponse):
-            texts: list[str] = []
-            tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
-            for item in message.parts:
-                if isinstance(item, TextPart):
-                    texts.append(item.content)
-                elif isinstance(item, ToolCallPart):
-                    tool_calls.append(self._map_tool_call(item))
-                else:
-                    assert_never(item)
-            message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
-            if texts:
-                # Note: model responses from this model should only have one text item, so the following
-                # shouldn't merge multiple texts into one unless you switch models between runs:
-                message_param['content'] = '\n\n'.join(texts)
-            if tool_calls:
-                message_param['tool_calls'] = tool_calls
-            yield message_param
-        else:
-            assert_never(message)
+        openai_messages: list[chat.ChatCompletionMessageParam] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                async for item in self._map_user_message(message):
+                    openai_messages.append(item)
+            elif isinstance(message, ModelResponse):
+                texts: list[str] = []
+                tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
+                for item in message.parts:
+                    if isinstance(item, TextPart):
+                        texts.append(item.content)
+                    elif isinstance(item, ToolCallPart):
+                        tool_calls.append(self._map_tool_call(item))
+                    else:
+                        assert_never(item)
+                message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
+                if texts:
+                    # Note: model responses from this model should only have one text item, so the following
+                    # shouldn't merge multiple texts into one unless you switch models between runs:
+                    message_param['content'] = '\n\n'.join(texts)
+                if tool_calls:
+                    message_param['tool_calls'] = tool_calls
+                openai_messages.append(message_param)
+            else:
+                assert_never(message)
+        if instructions := self._get_instructions(messages):
+            openai_messages.insert(0, chat.ChatCompletionSystemMessageParam(content=instructions, role='system'))
+        return openai_messages
 
     @staticmethod
     def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
@@ -446,6 +452,8 @@ class OpenAIModel(Model):
                     # file_data = f'data:{media_type};base64,{base64_encoded}'
                     # file = File(file={'file_data': file_data, 'file_name': item.url, 'file_id': item.url}, type='file')
                     # content.append(file)
+                elif isinstance(item, VideoUrl):  # pragma: no cover
+                    raise NotImplementedError('VideoUrl is not supported for OpenAI')
                 else:
                     assert_never(item)
         return chat.ChatCompletionUserMessageParam(role='user', content=content)
@@ -589,19 +597,19 @@ class OpenAIResponsesModel(Model):
         # standalone function to make it easier to override
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
-        elif not model_request_parameters.allow_text_result:
+        elif not model_request_parameters.allow_text_output:
             tool_choice = 'required'
         else:
             tool_choice = 'auto'
 
-        system_prompt, openai_messages = await self._map_message(messages)
+        instructions, openai_messages = await self._map_messages(messages)
         reasoning = self._get_reasoning(model_settings)
 
         try:
             return await self.client.responses.create(
                 input=openai_messages,
                 model=self._model_name,
-                instructions=system_prompt,
+                instructions=instructions,
                 parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
                 tools=tools or NOT_GIVEN,
                 tool_choice=tool_choice or NOT_GIVEN,
@@ -613,6 +621,7 @@ class OpenAIResponsesModel(Model):
                 timeout=model_settings.get('timeout', NOT_GIVEN),
                 reasoning=reasoning,
                 user=model_settings.get('openai_user', NOT_GIVEN),
+                extra_headers={'User-Agent': get_user_agent()},
             )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
@@ -629,8 +638,8 @@ class OpenAIResponsesModel(Model):
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[responses.FunctionToolParam]:
         tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.result_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        if model_request_parameters.output_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
     @staticmethod
@@ -644,15 +653,16 @@ class OpenAIResponsesModel(Model):
             'strict': f.strict or False,
         }
 
-    async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[responses.ResponseInputItemParam]]:
+    async def _map_messages(
+        self, messages: list[ModelMessage]
+    ) -> tuple[str | NotGiven, list[responses.ResponseInputItemParam]]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.responses.ResponseInputParam`."""
-        system_prompt: str = ''
         openai_messages: list[responses.ResponseInputItemParam] = []
         for message in messages:
             if isinstance(message, ModelRequest):
                 for part in message.parts:
                     if isinstance(part, SystemPromptPart):
-                        system_prompt += part.content
+                        openai_messages.append(responses.EasyInputMessageParam(role='system', content=part.content))
                     elif isinstance(part, UserPromptPart):
                         openai_messages.append(await self._map_user_prompt(part))
                     elif isinstance(part, ToolReturnPart):
@@ -689,7 +699,8 @@ class OpenAIResponsesModel(Model):
                         assert_never(item)
             else:
                 assert_never(message)
-        return system_prompt, openai_messages
+        instructions = self._get_instructions(messages) or NOT_GIVEN
+        return instructions, openai_messages
 
     @staticmethod
     def _map_tool_call(t: ToolCallPart) -> responses.ResponseFunctionToolCallParam:
@@ -762,6 +773,8 @@ class OpenAIResponsesModel(Model):
                             filename=f'filename.{item.format}',
                         )
                     )
+                elif isinstance(item, VideoUrl):  # pragma: no cover
+                    raise NotImplementedError('VideoUrl is not supported for OpenAI.')
                 else:
                     assert_never(item)
         return responses.EasyInputMessageParam(role='user', content=content)
@@ -919,137 +932,79 @@ def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk | responses.R
         )
 
 
-class _StrictSchemaHelper:
-    def make_schema_strict(self, schema: dict[str, Any]) -> dict[str, Any]:
-        """Recursively handle the schema to make it compatible with OpenAI strict mode.
-
-        See https://platform.openai.com/docs/guides/function-calling?api-mode=responses#strict-mode for more details,
-        but this basically just requires:
-        * `additionalProperties` must be set to false for each object in the parameters
-        * all fields in properties must be marked as required
-        """
-        assert isinstance(schema, dict), 'Schema must be a dictionary, this is probably a bug'
+@dataclass
+class _OpenAIJsonSchema(WalkJsonSchema):
+    """Recursively handle the schema to make it compatible with OpenAI strict mode.
 
-        # Create a copy to avoid modifying the original schema
-        schema = schema.copy()
+    See https://platform.openai.com/docs/guides/function-calling?api-mode=responses#strict-mode for more details,
+    but this basically just requires:
+    * `additionalProperties` must be set to false for each object in the parameters
+    * all fields in properties must be marked as required
+    """
 
-        # Handle $defs
-        if defs := schema.get('$defs'):
-            schema['$defs'] = {k: self.make_schema_strict(v) for k, v in defs.items()}
+    def __init__(self, schema: JsonSchema, strict: bool | None):
+        super().__init__(schema)
+        self.strict = strict
+        self.is_strict_compatible = True
+
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        # Remove unnecessary keys
+        schema.pop('title', None)
+        schema.pop('default', None)
+        schema.pop('$schema', None)
+        schema.pop('discriminator', None)
+
+        # Remove incompatible keys, but note their impact in the description provided to the LLM
+        description = schema.get('description')
+        min_length = schema.pop('minLength', None)
+        max_length = schema.pop('minLength', None)
+        if description is not None:
+            notes = list[str]()
+            if min_length is not None:  # pragma: no cover
+                notes.append(f'min_length={min_length}')
+            if max_length is not None:  # pragma: no cover
+                notes.append(f'max_length={max_length}')
+            if notes:  # pragma: no cover
+                schema['description'] = f'{description} ({", ".join(notes)})'
 
-        # Process schema based on its type
         schema_type = schema.get('type')
         if schema_type == 'object':
-            # Handle object type by setting additionalProperties to false
-            # and adding all properties to required list
-            self._make_object_schema_strict(schema)
-        elif schema_type == 'array':
-            # Handle array types by processing their items
-            if 'items' in schema:
-                items: Any = schema['items']
-                schema['items'] = self.make_schema_strict(items)
-            if 'prefixItems' in schema:
-                prefix_items: list[Any] = schema['prefixItems']
-                schema['prefixItems'] = [self.make_schema_strict(item) for item in prefix_items]
-
-        elif schema_type in {'string', 'number', 'integer', 'boolean', 'null'}:
-            pass  # Primitive types need no special handling
-        elif 'oneOf' in schema:
-            schema['oneOf'] = [self.make_schema_strict(item) for item in schema['oneOf']]
-        elif 'anyOf' in schema:
-            schema['anyOf'] = [self.make_schema_strict(item) for item in schema['anyOf']]
-
+            if self.strict is True:
+                # additional properties are disallowed
+                schema['additionalProperties'] = False
+
+                # all properties are required
+                if 'properties' not in schema:
+                    schema['properties'] = dict[str, Any]()
+                schema['required'] = list(schema['properties'].keys())
+
+            elif self.strict is None:
+                if (
+                    schema.get('additionalProperties') is not False
+                    or 'properties' not in schema
+                    or 'required' not in schema
+                ):
+                    self.is_strict_compatible = False
+                else:
+                    required = schema['required']
+                    for k in schema['properties'].keys():
+                        if k not in required:
+                            self.is_strict_compatible = False
         return schema
 
-    def _make_object_schema_strict(self, schema: dict[str, Any]) -> None:
-        schema['additionalProperties'] = False
-
-        # Handle patternProperties; note this may not be compatible with strict mode but is included for completeness
-        if 'patternProperties' in schema and isinstance(schema['patternProperties'], dict):
-            pattern_props: dict[str, Any] = schema['patternProperties']
-            schema['patternProperties'] = {str(k): self.make_schema_strict(v) for k, v in pattern_props.items()}
-
-        # Handle properties — update their schemas recursively, and make all properties required
-        if 'properties' in schema and isinstance(schema['properties'], dict):
-            properties: dict[str, Any] = schema['properties']
-            schema['properties'] = {k: self.make_schema_strict(v) for k, v in properties.items()}
-            schema['required'] = list(properties.keys())
-
-    def is_schema_strict(self, schema: dict[str, Any]) -> bool:
-        """Check if the schema is strict-mode-compatible.
-
-        A schema is compatible if:
-        * `additionalProperties` is set to false for each object in the parameters
-        * all fields in properties are marked as required
-
-        See https://platform.openai.com/docs/guides/function-calling?api-mode=responses#strict-mode for more details.
-        """
-        assert isinstance(schema, dict), 'Schema must be a dictionary, this is probably a bug'
-
-        # Note that checking the defs first is usually the fastest way to proceed, but
-        # it makes it hard/impossible to hit coverage below, hence all the pragma no covers.
-        # I still included the handling below because I'm not _confident_ those code paths can't be hit.
-        if defs := schema.get('$defs'):
-            if not all(self.is_schema_strict(v) for v in defs.values()):  # pragma: no branch
-                return False
-
-        schema_type = schema.get('type')
-        if schema_type == 'object':
-            if not self._is_object_schema_strict(schema):
-                return False
-        elif schema_type == 'array':
-            if 'items' in schema:
-                items: Any = schema['items']
-                if not self.is_schema_strict(items):  # pragma: no cover
-                    return False
-            if 'prefixItems' in schema:
-                prefix_items: list[Any] = schema['prefixItems']
-                if not all(self.is_schema_strict(item) for item in prefix_items):  # pragma: no cover
-                    return False
-        elif schema_type in {'string', 'number', 'integer', 'boolean', 'null'}:
-            pass
-        elif 'oneOf' in schema:  # pragma: no cover
-            if not all(self.is_schema_strict(item) for item in schema['oneOf']):
-                return False
-
-        elif 'anyOf' in schema:  # pragma: no cover
-            if not all(self.is_schema_strict(item) for item in schema['anyOf']):
-                return False
-
-        return True
-
-    def _is_object_schema_strict(self, schema: dict[str, Any]) -> bool:
-        """Check if the schema is an object and has additionalProperties set to false."""
-        if schema.get('additionalProperties') is not False:
-            return False
-        if 'properties' not in schema:  # pragma: no cover
-            return False
-        if 'required' not in schema:  # pragma: no cover
-            return False
-
-        for k, v in schema['properties'].items():
-            if k not in schema['required']:
-                return False
-            if not self.is_schema_strict(v):  # pragma: no cover
-                return False
-
-        return True
-
 
 def _customize_request_parameters(model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
     """Customize the request parameters for OpenAI models."""
 
     def _customize_tool_def(t: ToolDefinition):
-        if t.strict is True:
-            parameters_json_schema = _StrictSchemaHelper().make_schema_strict(t.parameters_json_schema)
-            return replace(t, parameters_json_schema=parameters_json_schema)
-        elif t.strict is None:
-            strict = _StrictSchemaHelper().is_schema_strict(t.parameters_json_schema)
-            return replace(t, strict=strict)
-        return t
+        schema_transformer = _OpenAIJsonSchema(t.parameters_json_schema, strict=t.strict)
+        parameters_json_schema = schema_transformer.walk()
+        if t.strict is None:
+            t = replace(t, strict=schema_transformer.is_strict_compatible)
+        return replace(t, parameters_json_schema=parameters_json_schema)
 
     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],
+        allow_text_output=model_request_parameters.allow_text_output,
+        output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
     )

pydantic_ai/models/test.py

@@ -22,9 +22,9 @@ from ..messages import (
     ToolCallPart,
     ToolReturnPart,
 )
-from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
+from ..usage import Usage
 from . import (
     Model,
     ModelRequestParameters,
@@ -34,15 +34,15 @@ from .function import _estimate_string_tokens, _estimate_usage  # pyright: ignor
 
 
 @dataclass
-class _TextResult:
-    """A private wrapper class to tag a result that came from the custom_result_text field."""
+class _WrappedTextOutput:
+    """A private wrapper class to tag an output that came from the custom_output_text field."""
 
     value: str | None
 
 
 @dataclass
-class _FunctionToolResult:
-    """A wrapper class to tag a result that came from the custom_result_args field."""
+class _WrappedToolOutput:
+    """A wrapper class to tag an output that came from the custom_output_args field."""
 
     value: Any | None
 
@@ -65,16 +65,16 @@ class TestModel(Model):
 
     call_tools: list[str] | Literal['all'] = 'all'
     """List of tools to call. If `'all'`, all tools will be called."""
-    custom_result_text: str | None = None
-    """If set, this text is returned as the final result."""
-    custom_result_args: Any | None = None
-    """If set, these args will be passed to the result tool."""
+    custom_output_text: str | None = None
+    """If set, this text is returned as the final output."""
+    custom_output_args: Any | None = None
+    """If set, these args will be passed to the output tool."""
     seed: int = 0
     """Seed for generating random data."""
     last_model_request_parameters: ModelRequestParameters | None = field(default=None, init=False)
     """The last ModelRequestParameters passed to the model in a request.
 
-    The ModelRequestParameters contains information about the function and result tools available during request handling.
+    The ModelRequestParameters contains information about the function and output tools available during request handling.
 
     This is set when a request is made, so will reflect the function tools from the last step of the last run.
     """
@@ -88,7 +88,6 @@ class TestModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, Usage]:
         self.last_model_request_parameters = model_request_parameters
-
         model_response = self._request(messages, model_settings, model_request_parameters)
         usage = _estimate_usage([*messages, model_response])
         return model_response, usage
@@ -128,29 +127,29 @@ class TestModel(Model):
             tools_to_call = (function_tools_lookup[name] for name in self.call_tools)
             return [(r.name, r) for r in tools_to_call]
 
-    def _get_result(self, model_request_parameters: ModelRequestParameters) -> _TextResult | _FunctionToolResult:
-        if self.custom_result_text is not None:
-            assert model_request_parameters.allow_text_result, (
-                'Plain response not allowed, but `custom_result_text` is set.'
+    def _get_output(self, model_request_parameters: ModelRequestParameters) -> _WrappedTextOutput | _WrappedToolOutput:
+        if self.custom_output_text is not None:
+            assert model_request_parameters.allow_text_output, (
+                'Plain response not allowed, but `custom_output_text` is set.'
             )
-            assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
-            return _TextResult(self.custom_result_text)
-        elif self.custom_result_args is not None:
-            assert model_request_parameters.result_tools is not None, (
-                'No result tools provided, but `custom_result_args` is set.'
+            assert self.custom_output_args is None, 'Cannot set both `custom_output_text` and `custom_output_args`.'
+            return _WrappedTextOutput(self.custom_output_text)
+        elif self.custom_output_args is not None:
+            assert model_request_parameters.output_tools is not None, (
+                'No output tools provided, but `custom_output_args` is set.'
             )
-            result_tool = model_request_parameters.result_tools[0]
+            output_tool = model_request_parameters.output_tools[0]
 
-            if k := result_tool.outer_typed_dict_key:
-                return _FunctionToolResult({k: self.custom_result_args})
+            if k := output_tool.outer_typed_dict_key:
+                return _WrappedToolOutput({k: self.custom_output_args})
             else:
-                return _FunctionToolResult(self.custom_result_args)
-        elif model_request_parameters.allow_text_result:
-            return _TextResult(None)
-        elif model_request_parameters.result_tools:
-            return _FunctionToolResult(None)
+                return _WrappedToolOutput(self.custom_output_args)
+        elif model_request_parameters.allow_text_output:
+            return _WrappedTextOutput(None)
+        elif model_request_parameters.output_tools:
+            return _WrappedToolOutput(None)
         else:
-            return _TextResult(None)
+            return _WrappedTextOutput(None)  # pragma: no cover
 
     def _request(
         self,
@@ -159,8 +158,8 @@ class TestModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         tool_calls = self._get_tool_calls(model_request_parameters)
-        result = self._get_result(model_request_parameters)
-        result_tools = model_request_parameters.result_tools
+        output_wrapper = self._get_output(model_request_parameters)
+        output_tools = model_request_parameters.output_tools
 
         # if there are tools, the first thing we want to do is call all of them
         if tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
@@ -176,29 +175,29 @@ class TestModel(Model):
             # check if there are any retry prompts, if so retry them
             new_retry_names = {p.tool_name for p in last_message.parts if isinstance(p, RetryPromptPart)}
             if new_retry_names:
-                # Handle retries for both function tools and result tools
+                # Handle retries for both function tools and output tools
                 # Check function tools first
                 retry_parts: list[ModelResponsePart] = [
                     ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls if name in new_retry_names
                 ]
-                # Check result tools
-                if result_tools:
+                # Check output tools
+                if output_tools:
                     retry_parts.extend(
                         [
                             ToolCallPart(
                                 tool.name,
-                                result.value
-                                if isinstance(result, _FunctionToolResult) and result.value is not None
+                                output_wrapper.value
+                                if isinstance(output_wrapper, _WrappedToolOutput) and output_wrapper.value is not None
                                 else self.gen_tool_args(tool),
                             )
-                            for tool in result_tools
+                            for tool in output_tools
                             if tool.name in new_retry_names
                         ]
                     )
                 return ModelResponse(parts=retry_parts, model_name=self._model_name)
 
-        if isinstance(result, _TextResult):
-            if (response_text := result.value) is None:
+        if isinstance(output_wrapper, _WrappedTextOutput):
+            if (response_text := output_wrapper.value) is None:
                 # build up details of tool responses
                 output: dict[str, Any] = {}
                 for message in messages:
@@ -215,16 +214,16 @@ class TestModel(Model):
             else:
                 return ModelResponse(parts=[TextPart(response_text)], model_name=self._model_name)
         else:
-            assert result_tools, 'No result tools provided'
-            custom_result_args = result.value
-            result_tool = result_tools[self.seed % len(result_tools)]
-            if custom_result_args is not None:
+            assert output_tools, 'No output tools provided'
+            custom_output_args = output_wrapper.value
+            output_tool = output_tools[self.seed % len(output_tools)]
+            if custom_output_args is not None:
                 return ModelResponse(
-                    parts=[ToolCallPart(result_tool.name, custom_result_args)], model_name=self._model_name
+                    parts=[ToolCallPart(output_tool.name, custom_output_args)], model_name=self._model_name
                 )
             else:
-                response_args = self.gen_tool_args(result_tool)
-                return ModelResponse(parts=[ToolCallPart(result_tool.name, response_args)], model_name=self._model_name)
+                response_args = self.gen_tool_args(output_tool)
+                return ModelResponse(parts=[ToolCallPart(output_tool.name, response_args)], model_name=self._model_name)
 
 
 @dataclass

pydantic_ai/models/wrapper.py

@@ -37,6 +37,9 @@ class WrapperModel(Model):
         async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream:
             yield response_stream
 
+    def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
+        return self.wrapped.customize_request_parameters(model_request_parameters)
+
     @property
     def model_name(self) -> str:
         return self.wrapped.model_name

pydantic_ai/providers/__init__.py

@@ -52,6 +52,10 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .deepseek import DeepSeekProvider
 
         return DeepSeekProvider()
+    elif provider == 'azure':
+        from .azure import AzureProvider
+
+        return AzureProvider()
     elif provider == 'google-vertex':
         from .google_vertex import GoogleVertexProvider
 

pydantic_ai/providers/azure.py

@@ -87,9 +87,9 @@ class AzureProvider(Provider[AsyncOpenAI]):
                     'Must provide one of the `azure_endpoint` argument or the `AZURE_OPENAI_ENDPOINT` environment variable'
                 )
 
-            if not api_key and 'OPENAI_API_KEY' not in os.environ:  # pragma: no cover
+            if not api_key and 'AZURE_OPENAI_API_KEY' not in os.environ:  # pragma: no cover
                 raise UserError(
-                    'Must provide one of the `api_key` argument or the `OPENAI_API_KEY` environment variable'
+                    'Must provide one of the `api_key` argument or the `AZURE_OPENAI_API_KEY` environment variable'
                 )
 
             if not api_version and 'OPENAI_API_VERSION' not in os.environ:  # pragma: no cover