pydantic-ai-slim 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -285,6 +285,8 @@ class OpenAIChatModel(Model):
             'vercel',
             'litellm',
             'nebius',
+            'ovhcloud',
+            'gateway',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -314,6 +316,8 @@ class OpenAIChatModel(Model):
             'vercel',
             'litellm',
             'nebius',
+            'ovhcloud',
+            'gateway',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -342,6 +346,8 @@ class OpenAIChatModel(Model):
             'vercel',
             'litellm',
             'nebius',
+            'ovhcloud',
+            'gateway',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -363,7 +369,7 @@ class OpenAIChatModel(Model):
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -559,24 +565,7 @@ class OpenAIChatModel(Model):
         # - https://openrouter.ai/docs/use-cases/reasoning-tokens#preserving-reasoning-blocks
         # If you need this, please file an issue.
 
-        vendor_details: dict[str, Any] = {}
-
-        # Add logprobs to vendor_details if available
-        if choice.logprobs is not None and choice.logprobs.content:
-            # Convert logprobs to a serializable format
-            vendor_details['logprobs'] = [
-                {
-                    'token': lp.token,
-                    'bytes': lp.bytes,
-                    'logprob': lp.logprob,
-                    'top_logprobs': [
-                        {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
-                    ],
-                }
-                for lp in choice.logprobs.content
-            ]
-
-        if choice.message.content is not None:
+        if choice.message.content:
             items.extend(
                 (replace(part, id='content', provider_name=self.system) if isinstance(part, ThinkingPart) else part)
                 for part in split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags)
@@ -594,6 +583,23 @@ class OpenAIChatModel(Model):
                 part.tool_call_id = _guard_tool_call_id(part)
                 items.append(part)
 
+        vendor_details: dict[str, Any] = {}
+
+        # Add logprobs to vendor_details if available
+        if choice.logprobs is not None and choice.logprobs.content:
+            # Convert logprobs to a serializable format
+            vendor_details['logprobs'] = [
+                {
+                    'token': lp.token,
+                    'bytes': lp.bytes,
+                    'logprob': lp.logprob,
+                    'top_logprobs': [
+                        {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
+                    ],
+                }
+                for lp in choice.logprobs.content
+            ]
+
         raw_finish_reason = choice.finish_reason
         vendor_details['finish_reason'] = raw_finish_reason
         finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
@@ -903,7 +909,18 @@ class OpenAIResponsesModel(Model):
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'nebius']
+        provider: Literal[
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'grok',
+            'fireworks',
+            'together',
+            'nebius',
+            'ovhcloud',
+            'gateway',
+        ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
@@ -919,7 +936,7 @@ class OpenAIResponsesModel(Model):
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider(provider)
+            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -1616,21 +1633,6 @@ class OpenAIStreamedResponse(StreamedResponse):
                 self.provider_details = {'finish_reason': raw_finish_reason}
                 self.finish_reason = _CHAT_FINISH_REASON_MAP.get(raw_finish_reason)
 
-            # Handle the text part of the response
-            content = choice.delta.content
-            if content is not None:
-                maybe_event = self._parts_manager.handle_text_delta(
-                    vendor_part_id='content',
-                    content=content,
-                    thinking_tags=self._model_profile.thinking_tags,
-                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
-                        maybe_event.part.id = 'content'
-                        maybe_event.part.provider_name = self.provider_name
-                    yield maybe_event
-
             # The `reasoning_content` field is only present in DeepSeek models.
             # https://api-docs.deepseek.com/guides/reasoning_model
             if reasoning_content := getattr(choice.delta, 'reasoning_content', None):
@@ -1652,6 +1654,21 @@ class OpenAIStreamedResponse(StreamedResponse):
                     provider_name=self.provider_name,
                 )
 
+            # Handle the text part of the response
+            content = choice.delta.content
+            if content:
+                maybe_event = self._parts_manager.handle_text_delta(
+                    vendor_part_id='content',
+                    content=content,
+                    thinking_tags=self._model_profile.thinking_tags,
+                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+                )
+                if maybe_event is not None:  # pragma: no branch
+                    if isinstance(maybe_event, PartStartEvent) and isinstance(maybe_event.part, ThinkingPart):
+                        maybe_event.part.id = 'content'
+                        maybe_event.part.provider_name = self.provider_name
+                    yield maybe_event
+
             for dtc in choice.delta.tool_calls or []:
                 maybe_event = self._parts_manager.handle_tool_call_delta(
                     vendor_part_id=dtc.index,

pydantic_ai/providers/__init__.py

@@ -8,7 +8,7 @@ from __future__ import annotations as _annotations
 from abc import ABC, abstractmethod
 from typing import Any, Generic, TypeVar
 
-from pydantic_ai import ModelProfile
+from ..profiles import ModelProfile
 
 InterfaceClient = TypeVar('InterfaceClient')
 
@@ -53,7 +53,7 @@ class Provider(ABC, Generic[InterfaceClient]):
 
 def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
     """Infers the provider class from the provider name."""
-    if provider == 'openai':
+    if provider in ('openai', 'openai-chat', 'openai-responses'):
         from .openai import OpenAIProvider
 
         return OpenAIProvider
@@ -73,15 +73,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .azure import AzureProvider
 
         return AzureProvider
-    elif provider == 'google-vertex':
-        from .google_vertex import GoogleVertexProvider  # type: ignore[reportDeprecated]
+    elif provider in ('google-vertex', 'google-gla'):
+        from .google import GoogleProvider
 
-        return GoogleVertexProvider  # type: ignore[reportDeprecated]
-    elif provider == 'google-gla':
-        from .google_gla import GoogleGLAProvider  # type: ignore[reportDeprecated]
-
-        return GoogleGLAProvider  # type: ignore[reportDeprecated]
-    # NOTE: We don't test because there are many ways the `boto3.client` can retrieve the credentials.
+        return GoogleProvider
     elif provider == 'bedrock':
         from .bedrock import BedrockProvider
 
@@ -146,11 +141,25 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .nebius import NebiusProvider
 
         return NebiusProvider
+    elif provider == 'ovhcloud':
+        from .ovhcloud import OVHcloudProvider
+
+        return OVHcloudProvider
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')
 
 
 def infer_provider(provider: str) -> Provider[Any]:
     """Infer the provider from the provider name."""
-    provider_class = infer_provider_class(provider)
-    return provider_class()
+    if provider.startswith('gateway/'):
+        from .gateway import gateway_provider
+
+        provider = provider.removeprefix('gateway/')
+        return gateway_provider(provider)
+    elif provider in ('google-vertex', 'google-gla'):
+        from .google import GoogleProvider
+
+        return GoogleProvider(vertexai=provider == 'google-vertex')
+    else:
+        provider_class = infer_provider_class(provider)
+        return provider_class()

pydantic_ai/providers/bedrock.py

@@ -4,7 +4,7 @@ import os
 import re
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Literal, overload
+from typing import Any, Literal, overload
 
 from pydantic_ai import ModelProfile
 from pydantic_ai.exceptions import UserError
@@ -21,6 +21,8 @@ try:
     from botocore.client import BaseClient
     from botocore.config import Config
     from botocore.exceptions import NoRegionError
+    from botocore.session import Session
+    from botocore.tokens import FrozenAuthToken
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `boto3` package to use the Bedrock provider, '
@@ -117,10 +119,23 @@ class BedrockProvider(Provider[BaseClient]):
     def __init__(
         self,
         *,
+        api_key: str,
+        base_url: str | None = None,
         region_name: str | None = None,
+        profile_name: str | None = None,
+        aws_read_timeout: float | None = None,
+        aws_connect_timeout: float | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
         aws_access_key_id: str | None = None,
         aws_secret_access_key: str | None = None,
         aws_session_token: str | None = None,
+        base_url: str | None = None,
+        region_name: str | None = None,
         profile_name: str | None = None,
         aws_read_timeout: float | None = None,
         aws_connect_timeout: float | None = None,
@@ -130,11 +145,13 @@ class BedrockProvider(Provider[BaseClient]):
         self,
         *,
         bedrock_client: BaseClient | None = None,
-        region_name: str | None = None,
         aws_access_key_id: str | None = None,
         aws_secret_access_key: str | None = None,
         aws_session_token: str | None = None,
+        base_url: str | None = None,
+        region_name: str | None = None,
         profile_name: str | None = None,
+        api_key: str | None = None,
         aws_read_timeout: float | None = None,
         aws_connect_timeout: float | None = None,
     ) -> None:
@@ -142,10 +159,12 @@ class BedrockProvider(Provider[BaseClient]):
 
         Args:
             bedrock_client: A boto3 client for Bedrock Runtime. If provided, other arguments are ignored.
-            region_name: The AWS region name.
-            aws_access_key_id: The AWS access key ID.
-            aws_secret_access_key: The AWS secret access key.
-            aws_session_token: The AWS session token.
+            aws_access_key_id: The AWS access key ID. If not set, the `AWS_ACCESS_KEY_ID` environment variable will be used if available.
+            aws_secret_access_key: The AWS secret access key. If not set, the `AWS_SECRET_ACCESS_KEY` environment variable will be used if available.
+            aws_session_token: The AWS session token. If not set, the `AWS_SESSION_TOKEN` environment variable will be used if available.
+            api_key: The API key for Bedrock client. Can be used instead of `aws_access_key_id`, `aws_secret_access_key`, and `aws_session_token`. If not set, the `AWS_BEARER_TOKEN_BEDROCK` environment variable will be used if available.
+            base_url: The base URL for the Bedrock client.
+            region_name: The AWS region name. If not set, the `AWS_DEFAULT_REGION` environment variable will be used if available.
             profile_name: The AWS profile name.
             aws_read_timeout: The read timeout for Bedrock client.
             aws_connect_timeout: The connect timeout for Bedrock client.
@@ -153,19 +172,44 @@ class BedrockProvider(Provider[BaseClient]):
         if bedrock_client is not None:
             self._client = bedrock_client
         else:
+            read_timeout = aws_read_timeout or float(os.getenv('AWS_READ_TIMEOUT', 300))
+            connect_timeout = aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60))
+            config: dict[str, Any] = {
+                'read_timeout': read_timeout,
+                'connect_timeout': connect_timeout,
+            }
             try:
-                read_timeout = aws_read_timeout or float(os.getenv('AWS_READ_TIMEOUT', 300))
-                connect_timeout = aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60))
-                session = boto3.Session(
-                    aws_access_key_id=aws_access_key_id,
-                    aws_secret_access_key=aws_secret_access_key,
-                    aws_session_token=aws_session_token,
-                    region_name=region_name,
-                    profile_name=profile_name,
-                )
+                if api_key is not None:
+                    session = boto3.Session(
+                        botocore_session=_BearerTokenSession(api_key),
+                        region_name=region_name,
+                        profile_name=profile_name,
+                    )
+                    config['signature_version'] = 'bearer'
+                else:
+                    session = boto3.Session(
+                        aws_access_key_id=aws_access_key_id,
+                        aws_secret_access_key=aws_secret_access_key,
+                        aws_session_token=aws_session_token,
+                        region_name=region_name,
+                        profile_name=profile_name,
+                    )
                 self._client = session.client(  # type: ignore[reportUnknownMemberType]
                     'bedrock-runtime',
-                    config=Config(read_timeout=read_timeout, connect_timeout=connect_timeout),
+                    config=Config(**config),
+                    endpoint_url=base_url,
                 )
             except NoRegionError as exc:  # pragma: no cover
                 raise UserError('You must provide a `region_name` or a boto3 client for Bedrock Runtime.') from exc
+
+
+class _BearerTokenSession(Session):
+    def __init__(self, token: str):
+        super().__init__()
+        self.token = token
+
+    def get_auth_token(self, **_kwargs: Any) -> FrozenAuthToken:
+        return FrozenAuthToken(self.token)
+
+    def get_credentials(self) -> None:  # type: ignore[reportIncompatibleMethodOverride]
+        return None

pydantic_ai/providers/gateway.py

@@ -3,14 +3,16 @@
 from __future__ import annotations as _annotations
 
 import os
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 import httpx
 
 from pydantic_ai.exceptions import UserError
-from pydantic_ai.models import Model, cached_async_http_client, get_user_agent
+from pydantic_ai.models import cached_async_http_client
 
 if TYPE_CHECKING:
+    from botocore.client import BaseClient
     from google.genai import Client as GoogleClient
     from groq import AsyncGroq
     from openai import AsyncOpenAI
@@ -18,6 +20,8 @@ if TYPE_CHECKING:
     from pydantic_ai.models.anthropic import AsyncAnthropicClient
     from pydantic_ai.providers import Provider
 
+GATEWAY_BASE_URL = 'https://gateway.pydantic.dev/proxy'
+
 
 @overload
 def gateway_provider(
@@ -57,13 +61,34 @@ def gateway_provider(
 ) -> Provider[AsyncAnthropicClient]: ...
 
 
+@overload
+def gateway_provider(
+    upstream_provider: Literal['bedrock'],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> Provider[BaseClient]: ...
+
+
+@overload
 def gateway_provider(
-    upstream_provider: Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex', 'anthropic'] | str,
+    upstream_provider: str,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> Provider[Any]: ...
+
+
+UpstreamProvider = Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex', 'anthropic', 'bedrock']
+
+
+def gateway_provider(
+    upstream_provider: UpstreamProvider | str,
     *,
     # Every provider
     api_key: str | None = None,
     base_url: str | None = None,
-    # OpenAI & Groq
+    # OpenAI, Groq & Anthropic
     http_client: httpx.AsyncClient | None = None,
 ) -> Provider[Any]:
     """Create a new Gateway provider.
@@ -73,25 +98,21 @@ def gateway_provider(
         api_key: The API key to use for authentication. If not provided, the `PYDANTIC_AI_GATEWAY_API_KEY`
             environment variable will be used if available.
         base_url: The base URL to use for the Gateway. If not provided, the `PYDANTIC_AI_GATEWAY_BASE_URL`
-            environment variable will be used if available. Otherwise, defaults to `http://localhost:8787/`.
+            environment variable will be used if available. Otherwise, defaults to `https://gateway.pydantic.dev/proxy`.
         http_client: The HTTP client to use for the Gateway.
     """
     api_key = api_key or os.getenv('PYDANTIC_AI_GATEWAY_API_KEY')
     if not api_key:
         raise UserError(
-            'Set the `PYDANTIC_AI_GATEWAY_API_KEY` environment variable or pass it via `gateway_provider(api_key=...)`'
+            'Set the `PYDANTIC_AI_GATEWAY_API_KEY` environment variable or pass it via `gateway_provider(..., api_key=...)`'
             ' to use the Pydantic AI Gateway provider.'
         )
 
-    base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', 'https://gateway.pydantic.dev/proxy')
-    http_client = http_client or cached_async_http_client(provider=f'gateway-{upstream_provider}')
-    http_client.event_hooks = {'request': [_request_hook]}
-
-    if upstream_provider in ('openai', 'openai-chat'):
-        from .openai import OpenAIProvider
+    base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', GATEWAY_BASE_URL)
+    http_client = http_client or cached_async_http_client(provider=f'gateway/{upstream_provider}')
+    http_client.event_hooks = {'request': [_request_hook(api_key)]}
 
-        return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'openai'), http_client=http_client)
-    elif upstream_provider == 'openai-responses':
+    if upstream_provider in ('openai', 'openai-chat', 'openai-responses'):
         from .openai import OpenAIProvider
 
         return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'openai'), http_client=http_client)
@@ -111,79 +132,46 @@ def gateway_provider(
                 http_client=http_client,
             )
         )
-    elif upstream_provider == 'google-vertex':
-        from google.genai import Client as GoogleClient
+    elif upstream_provider == 'bedrock':
+        from .bedrock import BedrockProvider
 
+        return BedrockProvider(
+            api_key=api_key,
+            base_url=_merge_url_path(base_url, 'bedrock'),
+            region_name='pydantic-ai-gateway',  # Fake region name to avoid NoRegionError
+        )
+    elif upstream_provider == 'google-vertex':
         from .google import GoogleProvider
 
         return GoogleProvider(
-            client=GoogleClient(
-                vertexai=True,
-                api_key='unset',
-                http_options={
-                    'base_url': _merge_url_path(base_url, 'google-vertex'),
-                    'headers': {'User-Agent': get_user_agent(), 'Authorization': api_key},
-                    # TODO(Marcelo): Until https://github.com/googleapis/python-genai/issues/1357 is solved.
-                    'async_client_args': {
-                        'transport': httpx.AsyncHTTPTransport(),
-                        'event_hooks': {'request': [_request_hook]},
-                    },
-                },
-            )
+            vertexai=True,
+            api_key=api_key,
+            base_url=_merge_url_path(base_url, 'google-vertex'),
+            http_client=http_client,
         )
-    else:  # pragma: no cover
-        raise UserError(f'Unknown provider: {upstream_provider}')
+    else:
+        raise UserError(f'Unknown upstream provider: {upstream_provider}')
 
 
-def infer_model(model_name: str) -> Model:
-    """Infer the model class that will be used to make requests to the gateway.
-
-    Args:
-        model_name: The name of the model to infer. Must be in the format "provider/model_name".
+def _request_hook(api_key: str) -> Callable[[httpx.Request], Awaitable[httpx.Request]]:
+    """Request hook for the gateway provider.
 
-    Returns:
-        The model class that will be used to make requests to the gateway.
+    It adds the `"traceparent"` and `"Authorization"` headers to the request.
     """
-    try:
-        upstream_provider, model_name = model_name.split('/', 1)
-    except ValueError:
-        raise UserError(f'The model name "{model_name}" is not in the format "provider/model_name".')
 
-    if upstream_provider in ('openai', 'openai-chat'):
-        from pydantic_ai.models.openai import OpenAIChatModel
+    async def _hook(request: httpx.Request) -> httpx.Request:
+        from opentelemetry.propagate import inject
 
-        return OpenAIChatModel(model_name, provider=gateway_provider('openai'))
-    elif upstream_provider == 'openai-responses':
-        from pydantic_ai.models.openai import OpenAIResponsesModel
-
-        return OpenAIResponsesModel(model_name, provider=gateway_provider('openai'))
-    elif upstream_provider == 'groq':
-        from pydantic_ai.models.groq import GroqModel
+        headers: dict[str, Any] = {}
+        inject(headers)
+        request.headers.update(headers)
 
-        return GroqModel(model_name, provider=gateway_provider('groq'))
-    elif upstream_provider == 'anthropic':
-        from pydantic_ai.models.anthropic import AnthropicModel
-
-        return AnthropicModel(model_name, provider=gateway_provider('anthropic'))
-    elif upstream_provider == 'google-vertex':
-        from pydantic_ai.models.google import GoogleModel
-
-        return GoogleModel(model_name, provider=gateway_provider('google-vertex'))
-    raise UserError(f'Unknown upstream provider: {upstream_provider}')
-
-
-async def _request_hook(request: httpx.Request) -> httpx.Request:
-    """Request hook for the gateway provider.
-
-    It adds the `"traceparent"` header to the request.
-    """
-    from opentelemetry.propagate import inject
+        if 'Authorization' not in request.headers:
+            request.headers['Authorization'] = f'Bearer {api_key}'
 
-    headers: dict[str, Any] = {}
-    inject(headers)
-    request.headers.update(headers)
+        return request
 
-    return request
+    return _hook
 
 
 def _merge_url_path(base_url: str, path: str) -> str: