pydantic-ai-slim 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

pydantic_ai/models/huggingface.py

@@ -0,0 +1,463 @@
+from __future__ import annotations as _annotations
+
+import base64
+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_extensions import assert_never
+
+from pydantic_ai._thinking_part import split_content_into_text_and_thinking
+from pydantic_ai.providers import Provider, infer_provider
+
+from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
+from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc
+from ..messages import (
+    AudioUrl,
+    BinaryContent,
+    DocumentUrl,
+    ImageUrl,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    ModelResponseStreamEvent,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+    VideoUrl,
+)
+from ..settings import ModelSettings
+from ..tools import ToolDefinition
+from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests
+
+try:
+    import aiohttp
+    from huggingface_hub import (
+        AsyncInferenceClient,
+        ChatCompletionInputMessage,
+        ChatCompletionInputMessageChunk,
+        ChatCompletionInputTool,
+        ChatCompletionInputToolCall,
+        ChatCompletionInputURL,
+        ChatCompletionOutput,
+        ChatCompletionOutputMessage,
+        ChatCompletionStreamOutput,
+    )
+    from huggingface_hub.errors import HfHubHTTPError
+
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `huggingface_hub` to use Hugging Face Inference Providers, '
+        'you can use the `huggingface` optional group — `pip install "pydantic-ai-slim[huggingface]"`'
+    ) from _import_error
+
+__all__ = (
+    'HuggingFaceModel',
+    'HuggingFaceModelSettings',
+)
+
+
+HFSystemPromptRole = Literal['system', 'user']
+
+LatestHuggingFaceModelNames = Literal[
+    'deepseek-ai/DeepSeek-R1',
+    'meta-llama/Llama-3.3-70B-Instruct',
+    'meta-llama/Llama-4-Maverick-17B-128E-Instruct',
+    'meta-llama/Llama-4-Scout-17B-16E-Instruct',
+    'Qwen/QwQ-32B',
+    'Qwen/Qwen2.5-72B-Instruct',
+    'Qwen/Qwen3-235B-A22B',
+    'Qwen/Qwen3-32B',
+]
+"""Latest Hugging Face models."""
+
+
+HuggingFaceModelName = Union[str, LatestHuggingFaceModelNames]
+"""Possible Hugging Face model names.
+
+You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending).
+"""
+
+
+class HuggingFaceModelSettings(ModelSettings, total=False):
+    """Settings used for a Hugging Face model request."""
+
+    # ALL FIELDS MUST BE `huggingface_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    # This class is a placeholder for any future huggingface-specific settings
+
+
+@dataclass(init=False)
+class HuggingFaceModel(Model):
+    """A model that uses Hugging Face Inference Providers.
+
+    Internally, this uses the [HF Python client](https://github.com/huggingface/huggingface_hub) to interact with the API.
+
+    Apart from `__init__`, all methods are private or match those of the base class.
+    """
+
+    client: AsyncInferenceClient = field(repr=False)
+
+    _model_name: str = field(repr=False)
+    _system: str = field(default='huggingface', repr=False)
+
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        provider: Literal['huggingface'] | Provider[AsyncInferenceClient] = 'huggingface',
+    ):
+        """Initialize a Hugging Face model.
+
+        Args:
+            model_name: The name of the Model to use. You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending).
+            provider: The provider to use for Hugging Face Inference Providers. Can be either the string 'huggingface' or an
+                instance of `Provider[AsyncInferenceClient]`. If not provided, the other parameters will be used.
+        """
+        self._model_name = model_name
+        self._provider = provider
+        if isinstance(provider, str):
+            provider = infer_provider(provider)
+        self.client = provider.client
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        check_allow_model_requests()
+        response = await self._completions_create(
+            messages, False, cast(HuggingFaceModelSettings, model_settings or {}), model_request_parameters
+        )
+        model_response = self._process_response(response)
+        model_response.usage.requests = 1
+        return model_response
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> 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)
+
+    @property
+    def model_name(self) -> HuggingFaceModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The system / model provider."""
+        return self._system
+
+    @overload
+    async def _completions_create(
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[True],
+        model_settings: HuggingFaceModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterable[ChatCompletionStreamOutput]: ...
+
+    @overload
+    async def _completions_create(
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[False],
+        model_settings: HuggingFaceModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ChatCompletionOutput: ...
+
+    async def _completions_create(
+        self,
+        messages: list[ModelMessage],
+        stream: bool,
+        model_settings: HuggingFaceModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ChatCompletionOutput | AsyncIterable[ChatCompletionStreamOutput]:
+        tools = self._get_tools(model_request_parameters)
+
+        if not tools:
+            tool_choice: Literal['none', 'required', 'auto'] | None = None
+        elif not model_request_parameters.allow_text_output:
+            tool_choice = 'required'
+        else:
+            tool_choice = 'auto'
+
+        hf_messages = await self._map_messages(messages)
+
+        try:
+            return await self.client.chat.completions.create(  # type: ignore
+                model=self._model_name,
+                messages=hf_messages,  # type: ignore
+                tools=tools,
+                tool_choice=tool_choice or None,
+                stream=stream,
+                stop=model_settings.get('stop_sequences', None),
+                temperature=model_settings.get('temperature', None),
+                top_p=model_settings.get('top_p', None),
+                seed=model_settings.get('seed', None),
+                presence_penalty=model_settings.get('presence_penalty', None),
+                frequency_penalty=model_settings.get('frequency_penalty', None),
+                logit_bias=model_settings.get('logit_bias', None),  # type: ignore
+                logprobs=model_settings.get('logprobs', None),
+                top_logprobs=model_settings.get('top_logprobs', None),
+                extra_body=model_settings.get('extra_body'),  # type: ignore
+            )
+        except aiohttp.ClientResponseError as e:
+            raise ModelHTTPError(
+                status_code=e.status,
+                model_name=self.model_name,
+                body=e.response_error_payload,  # type: ignore
+            ) from e
+        except HfHubHTTPError as e:
+            raise ModelHTTPError(
+                status_code=e.response.status_code,
+                model_name=self.model_name,
+                body=e.response.content,
+            ) from e
+
+    def _process_response(self, response: ChatCompletionOutput) -> ModelResponse:
+        """Process a non-streamed response, and prepare a message to return."""
+        if response.created:
+            timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
+        else:
+            timestamp = _now_utc()
+
+        choice = response.choices[0]
+        content = choice.message.content
+        tool_calls = choice.message.tool_calls
+
+        items: list[ModelResponsePart] = []
+
+        if content is not None:
+            items.extend(split_content_into_text_and_thinking(content))
+        if tool_calls is not None:
+            for c in tool_calls:
+                items.append(ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id))
+        return ModelResponse(
+            items,
+            usage=_map_usage(response),
+            model_name=response.model,
+            timestamp=timestamp,
+            vendor_id=response.id,
+        )
+
+    async def _process_streamed_response(self, response: AsyncIterable[ChatCompletionStreamOutput]) -> StreamedResponse:
+        """Process a streamed response, and prepare a streaming response to return."""
+        peekable_response = _utils.PeekableAsyncStream(response)
+        first_chunk = await peekable_response.peek()
+        if isinstance(first_chunk, _utils.Unset):
+            raise UnexpectedModelBehavior(  # pragma: no cover
+                'Streamed response ended without content or tool calls'
+            )
+
+        return HuggingFaceStreamedResponse(
+            _model_name=self._model_name,
+            _response=peekable_response,
+            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+        )
+
+    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
+
+    async def _map_messages(
+        self, messages: list[ModelMessage]
+    ) -> list[ChatCompletionInputMessage | ChatCompletionOutputMessage]:
+        """Just maps a `pydantic_ai.Message` to a `huggingface_hub.ChatCompletionInputMessage`."""
+        hf_messages: list[ChatCompletionInputMessage | ChatCompletionOutputMessage] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                async for item in self._map_user_message(message):
+                    hf_messages.append(item)
+            elif isinstance(message, ModelResponse):
+                texts: list[str] = []
+                tool_calls: list[ChatCompletionInputToolCall] = []
+                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))
+                    elif isinstance(item, ThinkingPart):
+                        # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
+                        # please open an issue. The below code is the code to send thinking to the provider.
+                        # texts.append(f'<think>\n{item.content}\n</think>')
+                        pass
+                    else:
+                        assert_never(item)
+                message_param = ChatCompletionInputMessage(role='assistant')  # type: ignore
+                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
+                hf_messages.append(message_param)
+            else:
+                assert_never(message)
+        if instructions := self._get_instructions(messages):
+            hf_messages.insert(0, ChatCompletionInputMessage(content=instructions, role='system'))  # type: ignore
+        return hf_messages
+
+    @staticmethod
+    def _map_tool_call(t: ToolCallPart) -> ChatCompletionInputToolCall:
+        return ChatCompletionInputToolCall.parse_obj_as_instance(  # type: ignore
+            {
+                'id': _guard_tool_call_id(t=t),
+                'type': 'function',
+                'function': {
+                    'name': t.tool_name,
+                    'arguments': t.args_as_json_str(),
+                },
+            }
+        )
+
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> ChatCompletionInputTool:
+        tool_param: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj_as_instance(  # type: ignore
+            {
+                'type': 'function',
+                'function': {
+                    'name': f.name,
+                    'description': f.description,
+                    'parameters': f.parameters_json_schema,
+                },
+            }
+        )
+        if f.strict is not None:
+            tool_param['function']['strict'] = f.strict
+        return tool_param
+
+    async def _map_user_message(
+        self, message: ModelRequest
+    ) -> AsyncIterable[ChatCompletionInputMessage | ChatCompletionOutputMessage]:
+        for part in message.parts:
+            if isinstance(part, SystemPromptPart):
+                yield ChatCompletionInputMessage.parse_obj_as_instance({'role': 'system', 'content': part.content})  # type: ignore
+            elif isinstance(part, UserPromptPart):
+                yield await self._map_user_prompt(part)
+            elif isinstance(part, ToolReturnPart):
+                yield ChatCompletionOutputMessage.parse_obj_as_instance(  # type: ignore
+                    {
+                        'role': 'tool',
+                        'tool_call_id': _guard_tool_call_id(t=part),
+                        'content': part.model_response_str(),
+                    }
+                )
+            elif isinstance(part, RetryPromptPart):
+                if part.tool_name is None:
+                    yield ChatCompletionInputMessage.parse_obj_as_instance(  # type: ignore
+                        {'role': 'user', 'content': part.model_response()}
+                    )
+                else:
+                    yield ChatCompletionInputMessage.parse_obj_as_instance(  # type: ignore
+                        {
+                            'role': 'tool',
+                            'tool_call_id': _guard_tool_call_id(t=part),
+                            'content': part.model_response(),
+                        }
+                    )
+            else:
+                assert_never(part)
+
+    @staticmethod
+    async def _map_user_prompt(part: UserPromptPart) -> ChatCompletionInputMessage:
+        content: str | list[ChatCompletionInputMessage]
+        if isinstance(part.content, str):
+            content = part.content
+        else:
+            content = []
+            for item in part.content:
+                if isinstance(item, str):
+                    content.append(ChatCompletionInputMessageChunk(type='text', text=item))  # type: ignore
+                elif isinstance(item, ImageUrl):
+                    url = ChatCompletionInputURL(url=item.url)  # type: ignore
+                    content.append(ChatCompletionInputMessageChunk(type='image_url', image_url=url))  # type: ignore
+                elif isinstance(item, BinaryContent):
+                    base64_encoded = base64.b64encode(item.data).decode('utf-8')
+                    if item.is_image:
+                        url = ChatCompletionInputURL(url=f'data:{item.media_type};base64,{base64_encoded}')  # type: ignore
+                        content.append(ChatCompletionInputMessageChunk(type='image_url', image_url=url))  # type: ignore
+                    else:  # pragma: no cover
+                        raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
+                elif isinstance(item, AudioUrl):
+                    raise NotImplementedError('AudioUrl is not supported for Hugging Face')
+                elif isinstance(item, DocumentUrl):
+                    raise NotImplementedError('DocumentUrl is not supported for Hugging Face')
+                elif isinstance(item, VideoUrl):
+                    raise NotImplementedError('VideoUrl is not supported for Hugging Face')
+                else:
+                    assert_never(item)
+        return ChatCompletionInputMessage(role='user', content=content)  # type: ignore
+
+
+@dataclass
+class HuggingFaceStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Hugging Face models."""
+
+    _model_name: str
+    _response: AsyncIterable[ChatCompletionStreamOutput]
+    _timestamp: datetime
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        async for chunk in self._response:
+            self._usage += _map_usage(chunk)
+
+            try:
+                choice = chunk.choices[0]
+            except IndexError:
+                continue
+
+            # Handle the text part of the response
+            content = choice.delta.content
+            if content is not None:
+                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)
+
+            for dtc in choice.delta.tool_calls or []:
+                maybe_event = self._parts_manager.handle_tool_call_delta(
+                    vendor_part_id=dtc.index,
+                    tool_name=dtc.function and dtc.function.name,  # type: ignore
+                    args=dtc.function and dtc.function.arguments,
+                    tool_call_id=dtc.id,
+                )
+                if maybe_event is not None:
+                    yield maybe_event
+
+    @property
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
+
+    @property
+    def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
+        return self._timestamp
+
+
+def _map_usage(response: ChatCompletionOutput | ChatCompletionStreamOutput) -> usage.Usage:
+    response_usage = response.usage
+    if response_usage is None:
+        return usage.Usage()
+
+    return usage.Usage(
+        request_tokens=response_usage.prompt_tokens,
+        response_tokens=response_usage.completion_tokens,
+        total_tokens=response_usage.total_tokens,
+        details=None,
+    )

pydantic_ai/models/instrumented.py

@@ -138,7 +138,7 @@ class InstrumentationSettings:
                 **tokens_histogram_kwargs,
                 explicit_bucket_boundaries_advisory=TOKEN_HISTOGRAM_BOUNDARIES,
             )
-        except TypeError:  # pragma: lax no cover
+        except TypeError:
             # Older OTel/logfire versions don't support explicit_bucket_boundaries_advisory
             self.tokens_histogram = self.meter.create_histogram(
                 **tokens_histogram_kwargs,  # pyright: ignore

pydantic_ai/models/mistral.py

@@ -75,7 +75,7 @@ try:
     from mistralai.models.usermessage import UserMessage as MistralUserMessage
     from mistralai.types.basemodel import Unset as MistralUnset
     from mistralai.utils.eventstreaming import EventStreamAsync as MistralEventStreamAsync
-except ImportError as e:  # pragma: lax no cover
+except ImportError as e:  # pragma: no cover
     raise ImportError(
         'Please install `mistral` to use the Mistral model, '
         'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
@@ -217,7 +217,7 @@ class MistralModel(Model):
         except SDKError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise  # pragma: lax no cover
+            raise  # pragma: no cover
 
         assert response, 'A unexpected empty response from Mistral.'
         return response
@@ -428,7 +428,7 @@ class MistralModel(Model):
         if value_type == 'object':
             additional_properties = value.get('additionalProperties', {})
             if isinstance(additional_properties, bool):
-                return 'bool'  # pragma: no cover
+                return 'bool'  # pragma: lax no cover
             additional_properties_type = additional_properties.get('type')
             if (
                 additional_properties_type in SIMPLE_JSON_TYPE_MAPPING

pydantic_ai/models/openai.py

@@ -50,7 +50,7 @@ from . import (
 
 try:
     from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream, NotGiven
-    from openai.types import ChatModel, chat, responses
+    from openai.types import AllModels, chat, responses
     from openai.types.chat import (
         ChatCompletionChunk,
         ChatCompletionContentPartImageParam,
@@ -80,7 +80,7 @@ __all__ = (
     'OpenAIModelName',
 )
 
-OpenAIModelName = Union[str, ChatModel]
+OpenAIModelName = Union[str, AllModels]
 """
 Possible OpenAI model names.
 
@@ -345,7 +345,7 @@ class OpenAIModel(Model):
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise  # pragma: lax no cover
+            raise  # pragma: no cover
 
     def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -781,7 +781,7 @@ class OpenAIResponsesModel(Model):
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise  # pragma: lax no cover
+            raise  # pragma: no cover
 
     def _get_reasoning(self, model_settings: OpenAIResponsesModelSettings) -> Reasoning | NotGiven:
         reasoning_effort = model_settings.get('openai_reasoning_effort', None)
@@ -1051,7 +1051,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     vendor_part_id=chunk.item_id,
                     tool_name=None,
                     args=chunk.delta,
-                    tool_call_id=chunk.item_id,
+                    tool_call_id=None,
                 )
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event

pydantic_ai/output.py

@@ -10,7 +10,8 @@ from pydantic_core import core_schema
 from typing_extensions import TypeAliasType, TypeVar
 
 from . import _utils
-from .tools import RunContext
+from .messages import ToolCallPart
+from .tools import RunContext, ToolDefinition
 
 __all__ = (
     # classes
@@ -330,15 +331,17 @@ def StructuredDict(
     return _StructuredDict
 
 
+_OutputSpecItem = TypeAliasType(
+    '_OutputSpecItem',
+    Union[OutputTypeOrFunction[T_co], ToolOutput[T_co], NativeOutput[T_co], PromptedOutput[T_co], TextOutput[T_co]],
+    type_params=(T_co,),
+)
+
 OutputSpec = TypeAliasType(
     'OutputSpec',
     Union[
-        OutputTypeOrFunction[T_co],
-        ToolOutput[T_co],
-        NativeOutput[T_co],
-        PromptedOutput[T_co],
-        TextOutput[T_co],
-        Sequence[Union[OutputTypeOrFunction[T_co], ToolOutput[T_co], TextOutput[T_co]]],
+        _OutputSpecItem[T_co],
+        Sequence['OutputSpec[T_co]'],
     ],
     type_params=(T_co,),
 )
@@ -354,3 +357,14 @@ You should not need to import or use this type directly.
 
 See [output docs](../output.md) for more information.
 """
+
+
+@dataclass
+class DeferredToolCalls:
+    """Container for calls of deferred tools. This can be used as an agent's `output_type` and will be used as the output of the agent run if the model called any deferred tools.
+
+    See [deferred toolset docs](../toolsets.md#deferred-toolset) for more information.
+    """
+
+    tool_calls: list[ToolCallPart]
+    tool_defs: dict[str, ToolDefinition]

pydantic_ai/profiles/google.py

@@ -43,7 +43,7 @@ class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
                 f' Full schema: {self.schema}\n\n'
                 f'Source of additionalProperties within the full schema: {original_schema}\n\n'
                 'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
-                "If Google's APIs are updated to support this properly, please create an issue on the PydanticAI GitHub"
+                "If Google's APIs are updated to support this properly, please create an issue on the Pydantic AI GitHub"
                 ' and we will fix this behavior.',
                 UserWarning,
             )

pydantic_ai/profiles/moonshotai.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def moonshotai_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a MoonshotAI model."""
+    return None

pydantic_ai/providers/__init__.py

@@ -111,6 +111,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .heroku import HerokuProvider
 
         return HerokuProvider
+    elif provider == 'huggingface':
+        from .huggingface import HuggingFaceProvider
+
+        return HuggingFaceProvider
     elif provider == 'github':
         from .github import GitHubProvider
 

pydantic_ai/providers/google.py

@@ -86,7 +86,7 @@ class GoogleProvider(Provider[genai.Client]):
             # NOTE: We are keeping GEMINI_API_KEY for backwards compatibility.
             api_key = api_key or os.getenv('GOOGLE_API_KEY') or os.getenv('GEMINI_API_KEY')
 
-            if vertexai is None:  # pragma: lax no cover
+            if vertexai is None:
                 vertexai = bool(location or project or credentials)
 
             if not vertexai:
@@ -114,7 +114,7 @@ class GoogleProvider(Provider[genai.Client]):
                     http_options={'headers': {'User-Agent': get_user_agent()}},
                 )
         else:
-            self._client = client  # pragma: lax no cover
+            self._client = client
 
 
 VertexAILocation = Literal[

pydantic_ai/providers/google_vertex.py

@@ -50,7 +50,7 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
-        return google_model_profile(model_name)  # pragma: lax no cover
+        return google_model_profile(model_name)
 
     @overload
     def __init__(
@@ -116,6 +116,8 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
 class _VertexAIAuth(httpx.Auth):
     """Auth class for Vertex AI API."""
 
+    _refresh_lock: anyio.Lock = anyio.Lock()
+
     credentials: BaseCredentials | ServiceAccountCredentials | None
 
     def __init__(
@@ -169,10 +171,13 @@ class _VertexAIAuth(httpx.Auth):
         return creds
 
     async def _refresh_token(self) -> str:  # pragma: no cover
-        assert self.credentials is not None
-        await anyio.to_thread.run_sync(self.credentials.refresh, Request())  # type: ignore[reportUnknownMemberType]
-        assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
-        return self.credentials.token
+        async with self._refresh_lock:
+            assert self.credentials is not None
+            await anyio.to_thread.run_sync(self.credentials.refresh, Request())  # type: ignore[reportUnknownMemberType]
+            assert isinstance(self.credentials.token, str), (  # type: ignore[reportUnknownMemberType]
+                f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
+            )
+            return self.credentials.token
 
 
 async def _async_google_auth() -> tuple[BaseCredentials, str | None]: