pydantic-ai-slim 0.6.2__py3-none-any.whl → 0.7.1__py3-none-any.whl

pydantic_ai/models/huggingface.py

@@ -5,16 +5,15 @@ from collections.abc import AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Literal, Union, cast, overload
+from typing import Any, Literal, Union, cast, overload
 
 from typing_extensions import assert_never
 
-from pydantic_ai._thinking_part import split_content_into_text_and_thinking
-from pydantic_ai.exceptions import UserError
-from pydantic_ai.providers import Provider, infer_provider
-
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
+from .._run_context import RunContext
+from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc
+from ..exceptions import UserError
 from ..messages import (
     AudioUrl,
     BinaryContent,
@@ -37,9 +36,15 @@ from ..messages import (
     VideoUrl,
 )
 from ..profiles import ModelProfile
+from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests
+from . import (
+    Model,
+    ModelRequestParameters,
+    StreamedResponse,
+    check_allow_model_requests,
+)
 
 try:
     import aiohttp
@@ -150,12 +155,13 @@ class HuggingFaceModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._completions_create(
             messages, True, cast(HuggingFaceModelSettings, model_settings or {}), model_request_parameters
         )
-        yield await self._process_streamed_response(response)
+        yield await self._process_streamed_response(response, model_request_parameters)
 
     @property
     def model_name(self) -> HuggingFaceModelName:
@@ -263,7 +269,9 @@ class HuggingFaceModel(Model):
             vendor_id=response.id,
         )
 
-    async def _process_streamed_response(self, response: AsyncIterable[ChatCompletionStreamOutput]) -> StreamedResponse:
+    async def _process_streamed_response(
+        self, response: AsyncIterable[ChatCompletionStreamOutput], model_request_parameters: ModelRequestParameters
+    ) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
@@ -273,6 +281,7 @@ class HuggingFaceModel(Model):
             )
 
         return HuggingFaceStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _model_profile=self.profile,
             _response=peekable_response,
@@ -280,10 +289,7 @@ class HuggingFaceModel(Model):
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ChatCompletionInputTool]:
-        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
-        if model_request_parameters.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     async def _map_messages(
         self, messages: list[ModelMessage]

pydantic_ai/models/instrumented.py

@@ -18,6 +18,7 @@ from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provide
 from opentelemetry.util.types import AttributeValue
 from pydantic import TypeAdapter
 
+from .._run_context import RunContext
 from ..messages import ModelMessage, ModelRequest, ModelResponse
 from ..settings import ModelSettings
 from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse
@@ -218,12 +219,13 @@ class InstrumentedModel(WrapperModel):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         with self._instrument(messages, model_settings, model_request_parameters) as finish:
             response_stream: StreamedResponse | None = None
             try:
                 async with super().request_stream(
-                    messages, model_settings, model_request_parameters
+                    messages, model_settings, model_request_parameters, run_context
                 ) as response_stream:
                     yield response_stream
             finally:

pydantic_ai/models/mcp_sampling.py

@@ -3,9 +3,10 @@ from __future__ import annotations as _annotations
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from .. import _mcp, exceptions, usage
+from .._run_context import RunContext
 from ..messages import ModelMessage, ModelResponse
 from ..settings import ModelSettings
 from . import Model, ModelRequestParameters, StreamedResponse
@@ -76,6 +77,7 @@ class MCPSamplingModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         raise NotImplementedError('MCP Sampling does not support streaming')
         yield

pydantic_ai/models/mistral.py

@@ -11,11 +11,11 @@ import pydantic_core
 from httpx import Timeout
 from typing_extensions import assert_never
 
-from pydantic_ai._thinking_part import split_content_into_text_and_thinking
-from pydantic_ai.exceptions import UserError
-
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
+from .._run_context import RunContext
+from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc, number_to_datetime
+from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
     BuiltinToolCallPart,
@@ -176,6 +176,7 @@ class MistralModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         """Make a streaming request to the model from Pydantic AI call."""
         check_allow_model_requests()
@@ -183,7 +184,7 @@ class MistralModel(Model):
             messages, cast(MistralModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(model_request_parameters.output_tools, response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     @property
     def model_name(self) -> MistralModelName:
@@ -246,11 +247,7 @@ class MistralModel(Model):
         if model_request_parameters.builtin_tools:
             raise UserError('Mistral does not support built-in tools')
 
-        if (
-            model_request_parameters.output_tools
-            and model_request_parameters.function_tools
-            or model_request_parameters.function_tools
-        ):
+        if model_request_parameters.function_tools:
             # Function Calling
             response = await self.client.chat.stream_async(
                 model=str(self._model_name),
@@ -318,16 +315,13 @@ class MistralModel(Model):
 
         Returns None if both function_tools and output_tools are empty.
         """
-        all_tools: list[ToolDefinition] = (
-            model_request_parameters.function_tools + model_request_parameters.output_tools
-        )
         tools = [
             MistralTool(
                 function=MistralFunction(
                     name=r.name, parameters=r.parameters_json_schema, description=r.description or ''
                 )
             )
-            for r in all_tools
+            for r in model_request_parameters.tool_defs.values()
         ]
         return tools if tools else None
 
@@ -359,8 +353,8 @@ class MistralModel(Model):
 
     async def _process_streamed_response(
         self,
-        output_tools: list[ToolDefinition],
         response: MistralEventStreamAsync[MistralCompletionEvent],
+        model_request_parameters: ModelRequestParameters,
     ) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
@@ -376,10 +370,10 @@ class MistralModel(Model):
             timestamp = _now_utc()
 
         return MistralStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _response=peekable_response,
             _model_name=self._model_name,
             _timestamp=timestamp,
-            _output_tools={c.name: c for c in output_tools},
         )
 
     @staticmethod
@@ -586,7 +580,6 @@ class MistralStreamedResponse(StreamedResponse):
     _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
-    _output_tools: dict[str, ToolDefinition]
 
     _delta_content: str = field(default='', init=False)
 
@@ -605,10 +598,11 @@ class MistralStreamedResponse(StreamedResponse):
             text = _map_content(content)
             if text:
                 # Attempt to produce an output tool call from the received text
-                if self._output_tools:
+                output_tools = {c.name: c for c in self.model_request_parameters.output_tools}
+                if output_tools:
                     self._delta_content += text
                     # TODO: Port to native "manual JSON" mode
-                    maybe_tool_call_part = self._try_get_output_tool_from_text(self._delta_content, self._output_tools)
+                    maybe_tool_call_part = self._try_get_output_tool_from_text(self._delta_content, output_tools)
                     if maybe_tool_call_part:
                         yield self._parts_manager.handle_tool_call_part(
                             vendor_part_id='output',

pydantic_ai/models/openai.py

@@ -11,13 +11,13 @@ from typing import Any, Literal, Union, cast, overload
 from pydantic import ValidationError
 from typing_extensions import assert_never
 
-from pydantic_ai.exceptions import UserError
-
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
+from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
 from ..builtin_tools import CodeExecutionTool, WebSearchTool
+from ..exceptions import UserError
 from ..messages import (
     AudioUrl,
     BinaryContent,
@@ -59,6 +59,11 @@ try:
     from openai.types.chat.chat_completion_content_part_image_param import ImageURL
     from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio
     from openai.types.chat.chat_completion_content_part_param import File, FileFile
+    from openai.types.chat.chat_completion_message_custom_tool_call import ChatCompletionMessageCustomToolCall
+    from openai.types.chat.chat_completion_message_function_tool_call import ChatCompletionMessageFunctionToolCall
+    from openai.types.chat.chat_completion_message_function_tool_call_param import (
+        ChatCompletionMessageFunctionToolCallParam,
+    )
     from openai.types.chat.chat_completion_prediction_content_param import ChatCompletionPredictionContentParam
     from openai.types.chat.completion_create_params import (
         WebSearchOptions,
@@ -172,6 +177,14 @@ class OpenAIResponsesModelSettings(OpenAIModelSettings, total=False):
         middle of the conversation.
     """
 
+    openai_text_verbosity: Literal['low', 'medium', 'high']
+    """Constrains the verbosity of the model's text response.
+
+    Lower values will result in more concise responses, while higher values will
+    result in more verbose responses. Currently supported values are `low`,
+    `medium`, and `high`.
+    """
+
 
 @dataclass(init=False)
 class OpenAIModel(Model):
@@ -256,13 +269,14 @@ class OpenAIModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._completions_create(
             messages, True, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     @property
     def model_name(self) -> OpenAIModelName:
@@ -415,7 +429,14 @@ class OpenAIModel(Model):
             items.extend(split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                part = ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id)
+                if isinstance(c, ChatCompletionMessageFunctionToolCall):
+                    part = ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id)
+                elif isinstance(c, ChatCompletionMessageCustomToolCall):  # pragma: no cover
+                    # NOTE: Custom tool calls are not supported.
+                    # See <https://github.com/pydantic/pydantic-ai/issues/2513> for more details.
+                    raise RuntimeError('Custom tool calls are not supported')
+                else:
+                    assert_never(c)
                 part.tool_call_id = _guard_tool_call_id(part)
                 items.append(part)
         return ModelResponse(
@@ -427,7 +448,9 @@ class OpenAIModel(Model):
             vendor_id=response.id,
         )
 
-    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
+    async def _process_streamed_response(
+        self, response: AsyncStream[ChatCompletionChunk], model_request_parameters: ModelRequestParameters
+    ) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
@@ -437,6 +460,7 @@ class OpenAIModel(Model):
             )
 
         return OpenAIStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _model_profile=self.profile,
             _response=peekable_response,
@@ -444,10 +468,7 @@ 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.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     def _get_web_search_options(self, model_request_parameters: ModelRequestParameters) -> WebSearchOptions | None:
         for tool in model_request_parameters.builtin_tools:
@@ -461,8 +482,10 @@ class OpenAIModel(Model):
                         ),
                     )
                 return WebSearchOptions(search_context_size=tool.search_context_size)
-            elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
-                raise UserError('`CodeExecutionTool` is not supported by OpenAI')
+            else:
+                raise UserError(
+                    f'`{tool.__class__.__name__}` is not supported by `OpenAIModel`. If it should be, please file an issue.'
+                )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
@@ -473,7 +496,7 @@ class OpenAIModel(Model):
                     openai_messages.append(item)
             elif isinstance(message, ModelResponse):
                 texts: list[str] = []
-                tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
+                tool_calls: list[ChatCompletionMessageFunctionToolCallParam] = []
                 for item in message.parts:
                     if isinstance(item, TextPart):
                         texts.append(item.content)
@@ -504,8 +527,8 @@ class OpenAIModel(Model):
         return openai_messages
 
     @staticmethod
-    def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-        return chat.ChatCompletionMessageToolCallParam(
+    def _map_tool_call(t: ToolCallPart) -> ChatCompletionMessageFunctionToolCallParam:
+        return ChatCompletionMessageFunctionToolCallParam(
             id=_guard_tool_call_id(t=t),
             type='function',
             function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
@@ -631,14 +654,6 @@ class OpenAIResponsesModel(Model):
     The [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) is the
     new API for OpenAI models.
 
-    The Responses API has built-in tools, that you can use instead of building your own:
-
-    - [Web search](https://platform.openai.com/docs/guides/tools-web-search)
-    - [File search](https://platform.openai.com/docs/guides/tools-file-search)
-    - [Computer use](https://platform.openai.com/docs/guides/tools-computer-use)
-
-    Use the `openai_builtin_tools` setting to add these tools to your model.
-
     If you are interested in the differences between the Responses API and the Chat Completions API,
     see the [OpenAI API docs](https://platform.openai.com/docs/guides/responses-vs-chat-completions).
     """
@@ -702,13 +717,14 @@ class OpenAIResponsesModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
         response = await self._responses_create(
             messages, True, cast(OpenAIResponsesModelSettings, model_settings or {}), model_request_parameters
         )
         async with response:
-            yield await self._process_streamed_response(response)
+            yield await self._process_streamed_response(response, model_request_parameters)
 
     def _process_response(self, response: responses.Response) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -735,7 +751,9 @@ class OpenAIResponsesModel(Model):
         )
 
     async def _process_streamed_response(
-        self, response: AsyncStream[responses.ResponseStreamEvent]
+        self,
+        response: AsyncStream[responses.ResponseStreamEvent],
+        model_request_parameters: ModelRequestParameters,
     ) -> OpenAIResponsesStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
@@ -745,6 +763,7 @@ class OpenAIResponsesModel(Model):
 
         assert isinstance(first_chunk, responses.ResponseCreatedEvent)
         return OpenAIResponsesStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.response.created_at),
@@ -775,8 +794,11 @@ class OpenAIResponsesModel(Model):
         model_settings: OpenAIResponsesModelSettings,
         model_request_parameters: ModelRequestParameters,
     ) -> responses.Response | AsyncStream[responses.ResponseStreamEvent]:
-        tools = self._get_tools(model_request_parameters)
-        tools = self._get_builtin_tools(model_request_parameters) + tools
+        tools = (
+            self._get_builtin_tools(model_request_parameters)
+            + list(model_settings.get('openai_builtin_tools', []))
+            + self._get_tools(model_request_parameters)
+        )
 
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
@@ -805,6 +827,10 @@ class OpenAIResponsesModel(Model):
             openai_messages.insert(0, responses.EasyInputMessageParam(role='system', content=instructions))
             instructions = NOT_GIVEN
 
+        if verbosity := model_settings.get('openai_text_verbosity'):
+            text = text or {}
+            text['verbosity'] = verbosity
+
         sampling_settings = (
             model_settings
             if OpenAIModelProfile.from_profile(self.profile).openai_supports_sampling_settings
@@ -859,10 +885,7 @@ class OpenAIResponsesModel(Model):
         return Reasoning(effort=reasoning_effort, summary=reasoning_summary)
 
     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.output_tools:
-            tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
-        return tools
+        return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
     def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[responses.ToolParam]:
         tools: list[responses.ToolParam] = []
@@ -878,6 +901,10 @@ class OpenAIResponsesModel(Model):
                 tools.append(web_search_tool)
             elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
                 tools.append({'type': 'code_interpreter', 'container': {'type': 'auto'}})
+            else:
+                raise UserError(  # pragma: no cover
+                    f'`{tool.__class__.__name__}` is not supported by `OpenAIResponsesModel`. If it should be, please file an issue.'
+                )
         return tools
 
     def _map_tool_definition(self, f: ToolDefinition) -> responses.FunctionToolParam:

pydantic_ai/models/test.py

@@ -12,6 +12,7 @@ import pydantic_core
 from typing_extensions import assert_never
 
 from .. import _utils
+from .._run_context import RunContext
 from ..exceptions import UserError
 from ..messages import (
     BuiltinToolCallPart,
@@ -121,11 +122,13 @@ class TestModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         self.last_model_request_parameters = model_request_parameters
 
         model_response = self._request(messages, model_settings, model_request_parameters)
         yield TestStreamedResponse(
+            model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _structured_response=model_response,
             _messages=messages,

pydantic_ai/models/wrapper.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from functools import cached_property
 from typing import Any
 
+from .._run_context import RunContext
 from ..messages import ModelMessage, ModelResponse
 from ..profiles import ModelProfile
 from ..settings import ModelSettings
@@ -35,8 +36,11 @@ class WrapperModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
-        async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream:
+        async with self.wrapped.request_stream(
+            messages, model_settings, model_request_parameters, run_context
+        ) as response_stream:
             yield response_stream
 
     def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
@@ -60,4 +64,4 @@ class WrapperModel(Model):
         return self.wrapped.settings
 
     def __getattr__(self, item: str):
-        return getattr(self.wrapped, item)  # pragma: no cover
+        return getattr(self.wrapped, item)

pydantic_ai/profiles/openai.py

@@ -32,7 +32,7 @@ class OpenAIModelProfile(ModelProfile):
 
 def openai_model_profile(model_name: str) -> ModelProfile:
     """Get the model profile for an OpenAI model."""
-    is_reasoning_model = model_name.startswith('o')
+    is_reasoning_model = model_name.startswith('o') or model_name.startswith('gpt-5')
     # Structured Outputs (output mode 'native') is only supported with the gpt-4o-mini, gpt-4o-mini-2024-07-18, and gpt-4o-2024-08-06 model snapshots and later.
     # We leave it in here for all models because the `default_structured_output_mode` is `'tool'`, so `native` is only used
     # when the user specifically uses the `NativeOutput` marker, so an error from the API is acceptable.

pydantic_ai/providers/google.py

@@ -12,8 +12,8 @@ from pydantic_ai.profiles.google import google_model_profile
 from pydantic_ai.providers import Provider
 
 try:
-    from google import genai
     from google.auth.credentials import Credentials
+    from google.genai import Client
     from google.genai.types import HttpOptionsDict
 except ImportError as _import_error:
     raise ImportError(
@@ -22,7 +22,7 @@ except ImportError as _import_error:
     ) from _import_error
 
 
-class GoogleProvider(Provider[genai.Client]):
+class GoogleProvider(Provider[Client]):
     """Provider for Google."""
 
     @property
@@ -34,7 +34,7 @@ class GoogleProvider(Provider[genai.Client]):
         return str(self._client._api_client._http_options.base_url)  # type: ignore[reportPrivateUsage]
 
     @property
-    def client(self) -> genai.Client:
+    def client(self) -> Client:
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
@@ -53,7 +53,7 @@ class GoogleProvider(Provider[genai.Client]):
     ) -> None: ...
 
     @overload
-    def __init__(self, *, client: genai.Client) -> None: ...
+    def __init__(self, *, client: Client) -> None: ...
 
     @overload
     def __init__(self, *, vertexai: bool = False) -> None: ...
@@ -65,7 +65,7 @@ class GoogleProvider(Provider[genai.Client]):
         credentials: Credentials | None = None,
         project: str | None = None,
         location: VertexAILocation | Literal['global'] | None = None,
-        client: genai.Client | None = None,
+        client: Client | None = None,
         vertexai: bool | None = None,
     ) -> None:
         """Create a new Google provider.
@@ -102,9 +102,9 @@ class GoogleProvider(Provider[genai.Client]):
                         'Set the `GOOGLE_API_KEY` environment variable or pass it via `GoogleProvider(api_key=...)`'
                         'to use the Google Generative Language API.'
                     )
-                self._client = genai.Client(vertexai=vertexai, api_key=api_key, http_options=http_options)
+                self._client = Client(vertexai=vertexai, api_key=api_key, http_options=http_options)
             else:
-                self._client = genai.Client(
+                self._client = Client(
                     vertexai=vertexai,
                     project=project or os.environ.get('GOOGLE_CLOUD_PROJECT'),
                     # From https://github.com/pydantic/pydantic-ai/pull/2031/files#r2169682149: