pydantic-ai-slim 1.2.1__py3-none-any.whl → 1.10.0__py3-none-any.whl

pydantic_ai/models/__init__.py

@@ -27,6 +27,7 @@ from .._run_context import RunContext
 from ..builtin_tools import AbstractBuiltinTool
 from ..exceptions import UserError
 from ..messages import (
+    BaseToolCallPart,
     BinaryImage,
     FilePart,
     FileUrl,
@@ -35,14 +36,18 @@ from ..messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    ModelResponsePart,
     ModelResponseStreamEvent,
+    PartEndEvent,
     PartStartEvent,
     TextPart,
+    ThinkingPart,
     ToolCallPart,
     VideoUrl,
 )
 from ..output import OutputMode
 from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
+from ..providers import infer_provider
 from ..settings import ModelSettings, merge_model_settings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage
@@ -129,15 +134,8 @@ KnownModelName = TypeAliasType(
         'cerebras:qwen-3-235b-a22b-thinking-2507',
         'cohere:c4ai-aya-expanse-32b',
         'cohere:c4ai-aya-expanse-8b',
-        'cohere:command',
-        'cohere:command-light',
-        'cohere:command-light-nightly',
         'cohere:command-nightly',
-        'cohere:command-r',
-        'cohere:command-r-03-2024',
         'cohere:command-r-08-2024',
-        'cohere:command-r-plus',
-        'cohere:command-r-plus-04-2024',
         'cohere:command-r-plus-08-2024',
         'cohere:command-r7b-12-2024',
         'deepseek:deepseek-chat',
@@ -416,9 +414,17 @@ class Model(ABC):
         they need to customize the preparation flow further, but most implementations should simply call
         ``self.prepare_request(...)`` at the start of their ``request`` (and related) methods.
         """
-        merged_settings = merge_model_settings(self.settings, model_settings)
-        customized_parameters = self.customize_request_parameters(model_request_parameters)
-        return merged_settings, customized_parameters
+        model_settings = merge_model_settings(self.settings, model_settings)
+
+        if builtin_tools := model_request_parameters.builtin_tools:
+            # Deduplicate builtin tools
+            model_request_parameters = replace(
+                model_request_parameters,
+                builtin_tools=list({tool.unique_id: tool for tool in builtin_tools}.values()),
+            )
+
+        model_request_parameters = self.customize_request_parameters(model_request_parameters)
+        return model_settings, model_request_parameters
 
     @property
     @abstractmethod
@@ -541,7 +547,44 @@ class StreamedResponse(ABC):
                 async for event in iterator:
                     yield event
 
-            self._event_iterator = iterator_with_final_event(self._get_event_iterator())
+            async def iterator_with_part_end(
+                iterator: AsyncIterator[ModelResponseStreamEvent],
+            ) -> AsyncIterator[ModelResponseStreamEvent]:
+                last_start_event: PartStartEvent | None = None
+
+                def part_end_event(next_part: ModelResponsePart | None = None) -> PartEndEvent | None:
+                    if not last_start_event:
+                        return None
+
+                    index = last_start_event.index
+                    part = self._parts_manager.get_parts()[index]
+                    if not isinstance(part, TextPart | ThinkingPart | BaseToolCallPart):
+                        # Parts other than these 3 don't have deltas, so don't need an end part.
+                        return None
+
+                    return PartEndEvent(
+                        index=index,
+                        part=part,
+                        next_part_kind=next_part.part_kind if next_part else None,
+                    )
+
+                async for event in iterator:
+                    if isinstance(event, PartStartEvent):
+                        if last_start_event:
+                            end_event = part_end_event(event.part)
+                            if end_event:
+                                yield end_event
+
+                            event.previous_part_kind = last_start_event.part.part_kind
+                        last_start_event = event
+
+                    yield event
+
+                end_event = part_end_event()
+                if end_event:
+                    yield end_event
+
+            self._event_iterator = iterator_with_part_end(iterator_with_final_event(self._get_event_iterator()))
         return self._event_iterator
 
     @abstractmethod
@@ -644,41 +687,39 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         return TestModel()
 
     try:
-        provider, model_name = model.split(':', maxsplit=1)
+        provider_name, model_name = model.split(':', maxsplit=1)
     except ValueError:
-        provider = None
+        provider_name = None
         model_name = model
         if model_name.startswith(('gpt', 'o1', 'o3')):
-            provider = 'openai'
+            provider_name = 'openai'
         elif model_name.startswith('claude'):
-            provider = 'anthropic'
+            provider_name = 'anthropic'
         elif model_name.startswith('gemini'):
-            provider = 'google-gla'
+            provider_name = 'google-gla'
 
-        if provider is not None:
+        if provider_name is not None:
             warnings.warn(
-                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider}:{model_name}'.",
+                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider_name}:{model_name}'.",
                 DeprecationWarning,
             )
         else:
             raise UserError(f'Unknown model: {model}')
 
-    if provider == 'vertexai':  # pragma: no cover
+    if provider_name == 'vertexai':  # pragma: no cover
         warnings.warn(
             "The 'vertexai' provider name is deprecated. Use 'google-vertex' instead.",
             DeprecationWarning,
         )
-        provider = 'google-vertex'
+        provider_name = 'google-vertex'
 
-    if provider == 'gateway':
-        from ..providers.gateway import infer_model as infer_model_from_gateway
+    provider = infer_provider(provider_name)
 
-        return infer_model_from_gateway(model_name)
-    elif provider == 'cohere':
-        from .cohere import CohereModel
-
-        return CohereModel(model_name, provider=provider)
-    elif provider in (
+    model_kind = provider_name
+    if model_kind.startswith('gateway/'):
+        model_kind = provider_name.removeprefix('gateway/')
+    if model_kind in (
+        'openai',
         'azure',
         'deepseek',
         'cerebras',
@@ -688,42 +729,50 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'heroku',
         'moonshotai',
         'ollama',
-        'openai',
-        'openai-chat',
         'openrouter',
         'together',
         'vercel',
         'litellm',
         'nebius',
+        'ovhcloud',
     ):
+        model_kind = 'openai-chat'
+    elif model_kind in ('google-gla', 'google-vertex'):
+        model_kind = 'google'
+
+    if model_kind == 'openai-chat':
         from .openai import OpenAIChatModel
 
         return OpenAIChatModel(model_name, provider=provider)
-    elif provider == 'openai-responses':
+    elif model_kind == 'openai-responses':
         from .openai import OpenAIResponsesModel
 
-        return OpenAIResponsesModel(model_name, provider='openai')
-    elif provider in ('google-gla', 'google-vertex'):
+        return OpenAIResponsesModel(model_name, provider=provider)
+    elif model_kind == 'google':
         from .google import GoogleModel
 
         return GoogleModel(model_name, provider=provider)
-    elif provider == 'groq':
+    elif model_kind == 'groq':
         from .groq import GroqModel
 
         return GroqModel(model_name, provider=provider)
-    elif provider == 'mistral':
+    elif model_kind == 'cohere':
+        from .cohere import CohereModel
+
+        return CohereModel(model_name, provider=provider)
+    elif model_kind == 'mistral':
         from .mistral import MistralModel
 
         return MistralModel(model_name, provider=provider)
-    elif provider == 'anthropic':
+    elif model_kind == 'anthropic':
         from .anthropic import AnthropicModel
 
         return AnthropicModel(model_name, provider=provider)
-    elif provider == 'bedrock':
+    elif model_kind == 'bedrock':
         from .bedrock import BedrockConverseModel
 
         return BedrockConverseModel(model_name, provider=provider)
-    elif provider == 'huggingface':
+    elif model_kind == 'huggingface':
         from .huggingface import HuggingFaceModel
 
         return HuggingFaceModel(model_name, provider=provider)

pydantic_ai/models/anthropic.py

@@ -3,7 +3,7 @@ from __future__ import annotations as _annotations
 import io
 from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Any, Literal, cast, overload
 
@@ -13,7 +13,7 @@ from typing_extensions import assert_never
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._run_context import RunContext
 from .._utils import guard_tool_call_id as _guard_tool_call_id
-from ..builtin_tools import CodeExecutionTool, MemoryTool, WebSearchTool
+from ..builtin_tools import CodeExecutionTool, MCPServerTool, MemoryTool, WebSearchTool
 from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
@@ -68,6 +68,9 @@ try:
         BetaContentBlockParam,
         BetaImageBlockParam,
         BetaInputJSONDelta,
+        BetaMCPToolResultBlock,
+        BetaMCPToolUseBlock,
+        BetaMCPToolUseBlockParam,
         BetaMemoryTool20250818Param,
         BetaMessage,
         BetaMessageParam,
@@ -82,6 +85,8 @@ try:
         BetaRawMessageStreamEvent,
         BetaRedactedThinkingBlock,
         BetaRedactedThinkingBlockParam,
+        BetaRequestMCPServerToolConfigurationParam,
+        BetaRequestMCPServerURLDefinitionParam,
         BetaServerToolUseBlock,
         BetaServerToolUseBlockParam,
         BetaSignatureDelta,
@@ -162,7 +167,7 @@ class AnthropicModel(Model):
         self,
         model_name: AnthropicModelName,
         *,
-        provider: Literal['anthropic'] | Provider[AsyncAnthropicClient] = 'anthropic',
+        provider: Literal['anthropic', 'gateway'] | Provider[AsyncAnthropicClient] = 'anthropic',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
@@ -179,7 +184,7 @@ class AnthropicModel(Model):
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            provider = infer_provider('gateway/anthropic' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -264,7 +269,7 @@ class AnthropicModel(Model):
     ) -> BetaMessage | AsyncStream[BetaRawMessageStreamEvent]:
         # standalone function to make it easier to override
         tools = self._get_tools(model_request_parameters)
-        tools, beta_features = self._add_builtin_tools(tools, model_request_parameters)
+        tools, mcp_servers, beta_features = self._add_builtin_tools(tools, model_request_parameters)
 
         tool_choice: BetaToolChoiceParam | None
 
@@ -300,6 +305,7 @@ class AnthropicModel(Model):
                 model=self._model_name,
                 tools=tools or OMIT,
                 tool_choice=tool_choice or OMIT,
+                mcp_servers=mcp_servers or OMIT,
                 stream=stream,
                 thinking=model_settings.get('anthropic_thinking', OMIT),
                 stop_sequences=model_settings.get('stop_sequences', OMIT),
@@ -318,11 +324,14 @@ class AnthropicModel(Model):
     def _process_response(self, response: BetaMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
+        builtin_tool_calls: dict[str, BuiltinToolCallPart] = {}
         for item in response.content:
             if isinstance(item, BetaTextBlock):
                 items.append(TextPart(content=item.text))
             elif isinstance(item, BetaServerToolUseBlock):
-                items.append(_map_server_tool_use_block(item, self.system))
+                call_part = _map_server_tool_use_block(item, self.system)
+                builtin_tool_calls[call_part.tool_call_id] = call_part
+                items.append(call_part)
             elif isinstance(item, BetaWebSearchToolResultBlock):
                 items.append(_map_web_search_tool_result_block(item, self.system))
             elif isinstance(item, BetaCodeExecutionToolResultBlock):
@@ -333,6 +342,13 @@ class AnthropicModel(Model):
                 )
             elif isinstance(item, BetaThinkingBlock):
                 items.append(ThinkingPart(content=item.thinking, signature=item.signature, provider_name=self.system))
+            elif isinstance(item, BetaMCPToolUseBlock):
+                call_part = _map_mcp_server_use_block(item, self.system)
+                builtin_tool_calls[call_part.tool_call_id] = call_part
+                items.append(call_part)
+            elif isinstance(item, BetaMCPToolResultBlock):
+                call_part = builtin_tool_calls.get(item.tool_use_id)
+                items.append(_map_mcp_server_result_block(item, call_part, self.system))
             else:
                 assert isinstance(item, BetaToolUseBlock), f'unexpected item type {type(item)}'
                 items.append(
@@ -383,8 +399,9 @@ class AnthropicModel(Model):
 
     def _add_builtin_tools(
         self, tools: list[BetaToolUnionParam], model_request_parameters: ModelRequestParameters
-    ) -> tuple[list[BetaToolUnionParam], list[str]]:
+    ) -> tuple[list[BetaToolUnionParam], list[BetaRequestMCPServerURLDefinitionParam], list[str]]:
         beta_features: list[str] = []
+        mcp_servers: list[BetaRequestMCPServerURLDefinitionParam] = []
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
                 user_location = UserLocation(type='approximate', **tool.user_location) if tool.user_location else None
@@ -408,11 +425,26 @@ class AnthropicModel(Model):
                 tools = [tool for tool in tools if tool['name'] != 'memory']
                 tools.append(BetaMemoryTool20250818Param(name='memory', type='memory_20250818'))
                 beta_features.append('context-management-2025-06-27')
+            elif isinstance(tool, MCPServerTool) and tool.url:
+                mcp_server_url_definition_param = BetaRequestMCPServerURLDefinitionParam(
+                    type='url',
+                    name=tool.id,
+                    url=tool.url,
+                )
+                if tool.allowed_tools is not None:  # pragma: no branch
+                    mcp_server_url_definition_param['tool_configuration'] = BetaRequestMCPServerToolConfigurationParam(
+                        enabled=bool(tool.allowed_tools),
+                        allowed_tools=tool.allowed_tools,
+                    )
+                if tool.authorization_token:  # pragma: no cover
+                    mcp_server_url_definition_param['authorization_token'] = tool.authorization_token
+                mcp_servers.append(mcp_server_url_definition_param)
+                beta_features.append('mcp-client-2025-04-04')
             else:  # pragma: no cover
                 raise UserError(
                     f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.'
                 )
-        return tools, beta_features
+        return tools, mcp_servers, beta_features
 
     async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[BetaMessageParam]]:  # noqa: C901
         """Just maps a `pydantic_ai.Message` to a `anthropic.types.MessageParam`."""
@@ -458,6 +490,8 @@ class AnthropicModel(Model):
                     | BetaCodeExecutionToolResultBlockParam
                     | BetaThinkingBlockParam
                     | BetaRedactedThinkingBlockParam
+                    | BetaMCPToolUseBlockParam
+                    | BetaMCPToolResultBlock
                 ] = []
                 for response_part in m.parts:
                     if isinstance(response_part, TextPart):
@@ -508,7 +542,7 @@ class AnthropicModel(Model):
                                     input=response_part.args_as_dict(),
                                 )
                                 assistant_content_params.append(server_tool_use_block_param)
-                            elif response_part.tool_name == CodeExecutionTool.kind:  # pragma: no branch
+                            elif response_part.tool_name == CodeExecutionTool.kind:
                                 server_tool_use_block_param = BetaServerToolUseBlockParam(
                                     id=tool_use_id,
                                     type='server_tool_use',
@@ -516,6 +550,21 @@ class AnthropicModel(Model):
                                     input=response_part.args_as_dict(),
                                 )
                                 assistant_content_params.append(server_tool_use_block_param)
+                            elif (
+                                response_part.tool_name.startswith(MCPServerTool.kind)
+                                and (server_id := response_part.tool_name.split(':', 1)[1])
+                                and (args := response_part.args_as_dict())
+                                and (tool_name := args.get('tool_name'))
+                                and (tool_args := args.get('tool_args'))
+                            ):  # pragma: no branch
+                                mcp_tool_use_block_param = BetaMCPToolUseBlockParam(
+                                    id=tool_use_id,
+                                    type='mcp_tool_use',
+                                    server_name=server_id,
+                                    name=tool_name,
+                                    input=tool_args,
+                                )
+                                assistant_content_params.append(mcp_tool_use_block_param)
                     elif isinstance(response_part, BuiltinToolReturnPart):
                         if response_part.provider_name == self.system:
                             tool_use_id = _guard_tool_call_id(t=response_part)
@@ -547,6 +596,16 @@ class AnthropicModel(Model):
                                         ),
                                     )
                                 )
+                            elif response_part.tool_name.startswith(MCPServerTool.kind) and isinstance(
+                                response_part.content, dict
+                            ):  # pragma: no branch
+                                assistant_content_params.append(
+                                    BetaMCPToolResultBlock(
+                                        tool_use_id=tool_use_id,
+                                        type='mcp_tool_result',
+                                        **cast(dict[str, Any], response_part.content),  # pyright: ignore[reportUnknownMemberType]
+                                    )
+                                )
                     elif isinstance(response_part, FilePart):  # pragma: no cover
                         # Files generated by models are not sent back to models that don't themselves generate files.
                         pass
@@ -661,6 +720,7 @@ class AnthropicStreamedResponse(StreamedResponse):
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         current_block: BetaContentBlock | None = None
 
+        builtin_tool_calls: dict[str, BuiltinToolCallPart] = {}
         async for event in self._response:
             if isinstance(event, BetaRawMessageStartEvent):
                 self._usage = _map_usage(event, self._provider_name, self._provider_url, self._model_name)
@@ -698,9 +758,11 @@ class AnthropicStreamedResponse(StreamedResponse):
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 elif isinstance(current_block, BetaServerToolUseBlock):
+                    call_part = _map_server_tool_use_block(current_block, self.provider_name)
+                    builtin_tool_calls[call_part.tool_call_id] = call_part
                     yield self._parts_manager.handle_part(
                         vendor_part_id=event.index,
-                        part=_map_server_tool_use_block(current_block, self.provider_name),
+                        part=call_part,
                     )
                 elif isinstance(current_block, BetaWebSearchToolResultBlock):
                     yield self._parts_manager.handle_part(
@@ -712,6 +774,32 @@ class AnthropicStreamedResponse(StreamedResponse):
                         vendor_part_id=event.index,
                         part=_map_code_execution_tool_result_block(current_block, self.provider_name),
                     )
+                elif isinstance(current_block, BetaMCPToolUseBlock):
+                    call_part = _map_mcp_server_use_block(current_block, self.provider_name)
+                    builtin_tool_calls[call_part.tool_call_id] = call_part
+
+                    args_json = call_part.args_as_json_str()
+                    # Drop the final `{}}` so that we can add tool args deltas
+                    args_json_delta = args_json[:-3]
+                    assert args_json_delta.endswith('"tool_args":'), (
+                        f'Expected {args_json_delta!r} to end in `"tool_args":`'
+                    )
+
+                    yield self._parts_manager.handle_part(
+                        vendor_part_id=event.index, part=replace(call_part, args=None)
+                    )
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=event.index,
+                        args=args_json_delta,
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
+                elif isinstance(current_block, BetaMCPToolResultBlock):
+                    call_part = builtin_tool_calls.get(current_block.tool_use_id)
+                    yield self._parts_manager.handle_part(
+                        vendor_part_id=event.index,
+                        part=_map_mcp_server_result_block(current_block, call_part, self.provider_name),
+                    )
 
             elif isinstance(event, BetaRawContentBlockDeltaEvent):
                 if isinstance(event.delta, BetaTextDelta):
@@ -749,7 +837,16 @@ class AnthropicStreamedResponse(StreamedResponse):
                     self.provider_details = {'finish_reason': raw_finish_reason}
                     self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
 
-            elif isinstance(event, BetaRawContentBlockStopEvent | BetaRawMessageStopEvent):  # pragma: no branch
+            elif isinstance(event, BetaRawContentBlockStopEvent):  # pragma: no branch
+                if isinstance(current_block, BetaMCPToolUseBlock):
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=event.index,
+                        args='}',
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
+                current_block = None
+            elif isinstance(event, BetaRawMessageStopEvent):  # pragma: no branch
                 current_block = None
 
     @property
@@ -817,3 +914,27 @@ def _map_code_execution_tool_result_block(
         content=code_execution_tool_result_content_ta.dump_python(item.content, mode='json'),
         tool_call_id=item.tool_use_id,
     )
+
+
+def _map_mcp_server_use_block(item: BetaMCPToolUseBlock, provider_name: str) -> BuiltinToolCallPart:
+    return BuiltinToolCallPart(
+        provider_name=provider_name,
+        tool_name=':'.join([MCPServerTool.kind, item.server_name]),
+        args={
+            'action': 'call_tool',
+            'tool_name': item.name,
+            'tool_args': cast(dict[str, Any], item.input),
+        },
+        tool_call_id=item.id,
+    )
+
+
+def _map_mcp_server_result_block(
+    item: BetaMCPToolResultBlock, call_part: BuiltinToolCallPart | None, provider_name: str
+) -> BuiltinToolReturnPart:
+    return BuiltinToolReturnPart(
+        provider_name=provider_name,
+        tool_name=call_part.tool_name if call_part else MCPServerTool.kind,
+        content=item.model_dump(mode='json', include={'content', 'is_error'}),
+        tool_call_id=item.tool_use_id,
+    )

pydantic_ai/models/bedrock.py

@@ -207,7 +207,7 @@ class BedrockConverseModel(Model):
         self,
         model_name: BedrockModelName,
         *,
-        provider: Literal['bedrock'] | Provider[BaseClient] = 'bedrock',
+        provider: Literal['bedrock', 'gateway'] | Provider[BaseClient] = 'bedrock',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
@@ -226,7 +226,7 @@ class BedrockConverseModel(Model):
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            provider = infer_provider('gateway/bedrock' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = cast('BedrockRuntimeClient', provider.client)
 
@@ -701,8 +701,8 @@ class BedrockStreamedResponse(StreamedResponse):
                                 signature=signature,
                                 provider_name=self.provider_name if signature else None,
                             )
-                    if 'text' in delta:
-                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
+                    if text := delta.get('text'):
+                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=text)
                         if maybe_event is not None:  # pragma: no branch
                             yield maybe_event
                     if 'toolUse' in delta:

pydantic_ai/models/cohere.py

@@ -62,15 +62,8 @@ except ImportError as _import_error:
 LatestCohereModelNames = Literal[
     'c4ai-aya-expanse-32b',
     'c4ai-aya-expanse-8b',
-    'command',
-    'command-light',
-    'command-light-nightly',
     'command-nightly',
-    'command-r',
-    'command-r-03-2024',
     'command-r-08-2024',
-    'command-r-plus',
-    'command-r-plus-04-2024',
     'command-r-plus-08-2024',
     'command-r7b-12-2024',
 ]

pydantic_ai/models/gemini.py

@@ -38,7 +38,7 @@ from ..messages import (
     VideoUrl,
 )
 from ..profiles import ModelProfileSpec
-from ..providers import Provider, infer_provider
+from ..providers import Provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
@@ -131,7 +131,14 @@ class GeminiModel(Model):
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            if provider == 'google-gla':
+                from pydantic_ai.providers.google_gla import GoogleGLAProvider  # type: ignore[reportDeprecated]
+
+                provider = GoogleGLAProvider()  # type: ignore[reportDeprecated]
+            else:
+                from pydantic_ai.providers.google_vertex import GoogleVertexProvider  # type: ignore[reportDeprecated]
+
+                provider = GoogleVertexProvider()  # type: ignore[reportDeprecated]
         self._provider = provider
         self.client = provider.client
         self._url = str(self.client.base_url)