Add PydanticAI Gateway provider (#2816)

Co-authored-by: Samuel Colvin <[email protected]>

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -718,7 +718,11 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         )
         provider = 'google-vertex'
 
-    if provider == 'cohere':
+    if provider == 'gateway':
+        from ..providers.gateway import infer_model as infer_model_from_gateway
+
+        return infer_model_from_gateway(model_name)
+    elif provider == 'cohere':
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -82,7 +82,7 @@
     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:
+except ImportError as e:  # pragma: lax 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]"`'

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -47,6 +47,9 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
         """The model profile for the named model, if available."""
         return None  # pragma: no cover
 
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(name={self.name}, base_url={self.base_url})'
+
 
 def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
     """Infers the provider class from the provider name."""

pydantic_ai_slim/pydantic_ai/providers/anthropic.py

@@ -68,7 +68,7 @@ def __init__(
             assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`'
             self._client = anthropic_client
         else:
-            api_key = api_key or os.environ.get('ANTHROPIC_API_KEY')
+            api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
             if not api_key:
                 raise UserError(
                     'Set the `ANTHROPIC_API_KEY` environment variable or pass it via `AnthropicProvider(api_key=...)`'

pydantic_ai_slim/pydantic_ai/providers/cohere.py

@@ -60,14 +60,14 @@ def __init__(
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
             self._client = cohere_client
         else:
-            api_key = api_key or os.environ.get('CO_API_KEY')
+            api_key = api_key or os.getenv('CO_API_KEY')
             if not api_key:
                 raise UserError(
                     'Set the `CO_API_KEY` environment variable or pass it via `CohereProvider(api_key=...)`'
                     'to use the Cohere provider.'
                 )
 
-            base_url = os.environ.get('CO_BASE_URL')
+            base_url = os.getenv('CO_BASE_URL')
             if http_client is not None:
                 self._client = AsyncClientV2(api_key=api_key, httpx_client=http_client, base_url=base_url)
             else:

pydantic_ai_slim/pydantic_ai/providers/gateway.py

@@ -0,0 +1,161 @@
+"""This module implements the Pydantic AI Gateway provider."""
+
+from __future__ import annotations as _annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Literal, overload
+from urllib.parse import urljoin
+
+import httpx
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import Model, cached_async_http_client, get_user_agent
+
+if TYPE_CHECKING:
+    from google.genai import Client as GoogleClient
+    from groq import AsyncGroq
+    from openai import AsyncOpenAI
+
+    from pydantic_ai.providers import Provider
+
+
+@overload
+def gateway_provider(
+    upstream_provider: Literal['openai', 'openai-chat', 'openai-responses'],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    http_client: httpx.AsyncClient | None = None,
+) -> Provider[AsyncOpenAI]: ...
+
+
+@overload
+def gateway_provider(
+    upstream_provider: Literal['groq'],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    http_client: httpx.AsyncClient | None = None,
+) -> Provider[AsyncGroq]: ...
+
+
+@overload
+def gateway_provider(
+    upstream_provider: Literal['google-vertex'],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> Provider[GoogleClient]: ...
+
+
+def gateway_provider(
+    upstream_provider: Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex'] | str,
+    *,
+    # Every provider
+    api_key: str | None = None,
+    base_url: str | None = None,
+    # OpenAI & Groq
+    http_client: httpx.AsyncClient | None = None,
+) -> Provider[Any]:
+    """Create a new Gateway provider.
+
+    Args:
+        upstream_provider: The upstream provider to use.
+        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/`.
+        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=...)`'
+            ' to use the Pydantic AI Gateway provider.'
+        )
+
+    base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', 'http://localhost:8787')
+    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
+
+        return OpenAIProvider(api_key=api_key, base_url=urljoin(base_url, 'openai'), http_client=http_client)
+    elif upstream_provider == 'openai-responses':
+        from .openai import OpenAIProvider
+
+        return OpenAIProvider(api_key=api_key, base_url=urljoin(base_url, 'openai'), http_client=http_client)
+    elif upstream_provider == 'groq':
+        from .groq import GroqProvider
+
+        return GroqProvider(api_key=api_key, base_url=urljoin(base_url, 'groq'), http_client=http_client)
+    elif upstream_provider == 'google-vertex':
+        from google.genai import Client as GoogleClient
+
+        from .google import GoogleProvider
+
+        return GoogleProvider(
+            client=GoogleClient(
+                vertexai=True,
+                api_key='unset',
+                http_options={
+                    'base_url': f'{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]},
+                    },
+                },
+            )
+        )
+    else:  # pragma: no cover
+        raise UserError(f'Unknown 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".
+
+    Returns:
+        The model class that will be used to make requests to the gateway.
+    """
+    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
+
+        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
+
+        return GroqModel(model_name, provider=gateway_provider('groq'))
+    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
+
+    headers: dict[str, Any] = {}
+    inject(headers)
+    request.headers.update(headers)
+
+    return request

pydantic_ai_slim/pydantic_ai/providers/google.py

@@ -106,13 +106,13 @@ def __init__(
             else:
                 self._client = Client(
                     vertexai=vertexai,
-                    project=project or os.environ.get('GOOGLE_CLOUD_PROJECT'),
+                    project=project or os.getenv('GOOGLE_CLOUD_PROJECT'),
                     # From https://github.com/pydantic/pydantic-ai/pull/2031/files#r2169682149:
                     # Currently `us-central1` supports the most models by far of any region including `global`, but not
                     # all of them. `us-central1` has all google models but is missing some Anthropic partner models,
                     # which use `us-east5` instead. `global` has fewer models but higher availability.
                     # For more details, check: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#available-regions
-                    location=location or os.environ.get('GOOGLE_CLOUD_LOCATION') or 'us-central1',
+                    location=location or os.getenv('GOOGLE_CLOUD_LOCATION') or 'us-central1',
                     credentials=credentials,
                     http_options=http_options,
                 )

pydantic_ai_slim/pydantic_ai/providers/google_gla.py

@@ -39,7 +39,7 @@ def __init__(self, api_key: str | None = None, http_client: httpx.AsyncClient |
                 will be used if available.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
-        api_key = api_key or os.environ.get('GEMINI_API_KEY')
+        api_key = api_key or os.getenv('GEMINI_API_KEY')
         if not api_key:
             raise UserError(
                 'Set the `GEMINI_API_KEY` environment variable or pass it via `GoogleGLAProvider(api_key=...)`'

pydantic_ai_slim/pydantic_ai/providers/groq.py

@@ -53,7 +53,7 @@ def name(self) -> str:
 
     @property
     def base_url(self) -> str:
-        return os.environ.get('GROQ_BASE_URL', 'https://api.groq.com')
+        return str(self.client.base_url)
 
     @property
     def client(self) -> AsyncGroq:
@@ -85,12 +85,15 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
     def __init__(self, *, groq_client: AsyncGroq | None = None) -> None: ...
 
     @overload
-    def __init__(self, *, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None: ...
+    def __init__(
+        self, *, api_key: str | None = None, base_url: str | None = None, http_client: httpx.AsyncClient | None = None
+    ) -> None: ...
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
+        base_url: str | None = None,
         groq_client: AsyncGroq | None = None,
         http_client: httpx.AsyncClient | None = None,
     ) -> None:
@@ -99,6 +102,8 @@ def __init__(
         Args:
             api_key: The API key to use for authentication, if not provided, the `GROQ_API_KEY` environment variable
                 will be used if available.
+            base_url: The base url for the Groq requests. If not provided, the `GROQ_BASE_URL` environment variable
+                will be used if available. Otherwise, defaults to Groq's base url.
             groq_client: An existing
                 [`AsyncGroq`](https://github.com/groq/groq-python?tab=readme-ov-file#async-usage)
                 client to use. If provided, `api_key` and `http_client` must be `None`.
@@ -107,17 +112,19 @@ def __init__(
         if groq_client is not None:
             assert http_client is None, 'Cannot provide both `groq_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `groq_client` and `api_key`'
+            assert base_url is None, 'Cannot provide both `groq_client` and `base_url`'
             self._client = groq_client
         else:
-            api_key = api_key or os.environ.get('GROQ_API_KEY')
+            api_key = api_key or os.getenv('GROQ_API_KEY')
+            base_url = base_url or os.getenv('GROQ_BASE_URL', 'https://api.groq.com')
 
             if not api_key:
                 raise UserError(
                     'Set the `GROQ_API_KEY` environment variable or pass it via `GroqProvider(api_key=...)`'
                     'to use the Groq provider.'
                 )
             elif http_client is not None:
-                self._client = AsyncGroq(base_url=self.base_url, api_key=api_key, http_client=http_client)
+                self._client = AsyncGroq(base_url=base_url, api_key=api_key, http_client=http_client)
             else:
                 http_client = cached_async_http_client(provider='groq')
-                self._client = AsyncGroq(base_url=self.base_url, api_key=api_key, http_client=http_client)
+                self._client = AsyncGroq(base_url=base_url, api_key=api_key, http_client=http_client)

pydantic_ai_slim/pydantic_ai/providers/heroku.py

@@ -65,14 +65,14 @@ def __init__(
             assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
             self._client = openai_client
         else:
-            api_key = api_key or os.environ.get('HEROKU_INFERENCE_KEY')
+            api_key = api_key or os.getenv('HEROKU_INFERENCE_KEY')
             if not api_key:
                 raise UserError(
                     'Set the `HEROKU_INFERENCE_KEY` environment variable or pass it via `HerokuProvider(api_key=...)`'
                     'to use the Heroku provider.'
                 )
 
-            base_url = base_url or os.environ.get('HEROKU_INFERENCE_URL', 'https://us.inference.heroku.com')
+            base_url = base_url or os.getenv('HEROKU_INFERENCE_URL', 'https://us.inference.heroku.com')
             base_url = base_url.rstrip('/') + '/v1'
 
             if http_client is not None: