Add Vercel AI Gateway provider (#2277)

Co-authored-by: Douwe Maan <[email protected]>

Author: joshualipman123

docs/api/providers.md

@@ -32,4 +32,6 @@
 
 ::: pydantic_ai.providers.openrouter.OpenRouterProvider
 
+::: pydantic_ai.providers.vercel.VercelProvider
+
 ::: pydantic_ai.providers.huggingface.HuggingFaceProvider

docs/models/index.md

@@ -19,6 +19,7 @@ In addition, many providers are compatible with the OpenAI API, and can be used
 - [Grok (xAI)](openai.md#grok-xai)
 - [Ollama](openai.md#ollama)
 - [OpenRouter](openai.md#openrouter)
+- [Vercel AI Gateway](openai.md#vercel-ai-gateway)
 - [Perplexity](openai.md#perplexity)
 - [Fireworks AI](openai.md#fireworks-ai)
 - [Together AI](openai.md#together-ai)

docs/models/openai.md

@@ -348,6 +348,41 @@ agent = Agent(model)
 ...
 ```
 
+### Vercel AI Gateway
+
+To use [Vercel's AI Gateway](https://vercel.com/docs/ai-gateway), first follow the [documentation](https://vercel.com/docs/ai-gateway) instructions on obtaining an API key or OIDC token.
+
+You can set your credentials using one of these environment variables:
+
+```bash
+export VERCEL_AI_GATEWAY_API_KEY='your-ai-gateway-api-key'
+# OR
+export VERCEL_OIDC_TOKEN='your-oidc-token'
+```
+
+Once you have set the environment variable, you can use it with the [`VercelProvider`][pydantic_ai.providers.vercel.VercelProvider]:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.vercel import VercelProvider
+
+# Uses environment variable automatically
+model = OpenAIModel(
+    'anthropic/claude-4-sonnet',
+    provider=VercelProvider(),
+)
+agent = Agent(model)
+
+# Or pass the API key directly
+model = OpenAIModel(
+    'anthropic/claude-4-sonnet',
+    provider=VercelProvider(api_key='your-vercel-ai-gateway-api-key'),
+)
+agent = Agent(model)
+...
+```
+
 ### Grok (xAI)
 
 Go to [xAI API Console](https://console.x.ai/) and create an API key.

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -615,6 +615,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'deepseek',
         'azure',
         'openrouter',
+        'vercel',
         'grok',
         'fireworks',
         'together',

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -191,7 +191,16 @@ def __init__(
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'heroku', 'github'
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'vercel',
+            'grok',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -62,6 +62,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .openrouter import OpenRouterProvider
 
         return OpenRouterProvider
+    elif provider == 'vercel':
+        from .vercel import VercelProvider
+
+        return VercelProvider
     elif provider == 'azure':
         from .azure import AzureProvider
 

pydantic_ai_slim/pydantic_ai/providers/vercel.py

@@ -0,0 +1,107 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the Vercel provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class VercelProvider(Provider[AsyncOpenAI]):
+    """Provider for Vercel AI Gateway API."""
+
+    @property
+    def name(self) -> str:
+        return 'vercel'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://ai-gateway.vercel.sh/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        provider_to_profile = {
+            'anthropic': anthropic_model_profile,
+            'bedrock': amazon_model_profile,
+            'cohere': cohere_model_profile,
+            'deepseek': deepseek_model_profile,
+            'mistral': mistral_model_profile,
+            'openai': openai_model_profile,
+            'vertex': google_model_profile,
+            'xai': grok_model_profile,
+        }
+
+        profile = None
+
+        try:
+            provider, model_name = model_name.split('/', 1)
+        except ValueError:
+            raise UserError(f"Model name must be in 'provider/model' format, got: {model_name!r}")
+
+        if provider in provider_to_profile:
+            profile = provider_to_profile[provider](model_name)
+
+        # As VercelProvider is always used with OpenAIModel, which used to unconditionally use OpenAIJsonSchemaTransformer,
+        # we need to maintain that behavior unless json_schema_transformer is set explicitly
+        return OpenAIModelProfile(
+            json_schema_transformer=OpenAIJsonSchemaTransformer,
+        ).update(profile)
+
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        # Support Vercel AI Gateway's standard environment variables
+        api_key = api_key or os.getenv('VERCEL_AI_GATEWAY_API_KEY') or os.getenv('VERCEL_OIDC_TOKEN')
+
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `VERCEL_AI_GATEWAY_API_KEY` or `VERCEL_OIDC_TOKEN` environment variable '
+                'or pass the API key via `VercelProvider(api_key=...)` to use the Vercel provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            http_client = cached_async_http_client(provider='vercel')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

tests/providers/test_provider_names.py

@@ -29,12 +29,14 @@
     from pydantic_ai.providers.openai import OpenAIProvider
     from pydantic_ai.providers.openrouter import OpenRouterProvider
     from pydantic_ai.providers.together import TogetherProvider
+    from pydantic_ai.providers.vercel import VercelProvider
 
     test_infer_provider_params = [
         ('anthropic', AnthropicProvider, 'ANTHROPIC_API_KEY'),
         ('cohere', CohereProvider, 'CO_API_KEY'),
         ('deepseek', DeepSeekProvider, 'DEEPSEEK_API_KEY'),
         ('openrouter', OpenRouterProvider, 'OPENROUTER_API_KEY'),
+        ('vercel', VercelProvider, 'VERCEL_AI_GATEWAY_API_KEY'),
         ('openai', OpenAIProvider, 'OPENAI_API_KEY'),
         ('azure', AzureProvider, 'AZURE_OPENAI'),
         ('google-vertex', GoogleVertexProvider, None),

tests/providers/test_vercel.py

@@ -0,0 +1,143 @@
+import re
+
+import httpx
+import pytest
+from pytest_mock import MockerFixture
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.profiles._json_schema import InlineDefsJsonSchemaTransformer
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, openai_model_profile
+
+from ..conftest import TestEnv, try_import
+
+with try_import() as imports_successful:
+    import openai
+
+    from pydantic_ai.providers.vercel import VercelProvider
+
+
+pytestmark = [
+    pytest.mark.skipif(not imports_successful(), reason='openai not installed'),
+    pytest.mark.vcr,
+    pytest.mark.anyio,
+]
+
+
+def test_vercel_provider():
+    provider = VercelProvider(api_key='api-key')
+    assert provider.name == 'vercel'
+    assert provider.base_url == 'https://ai-gateway.vercel.sh/v1'
+    assert isinstance(provider.client, openai.AsyncOpenAI)
+    assert provider.client.api_key == 'api-key'
+
+
+def test_vercel_provider_need_api_key(env: TestEnv) -> None:
+    env.remove('VERCEL_AI_GATEWAY_API_KEY')
+    env.remove('VERCEL_OIDC_TOKEN')
+    with pytest.raises(
+        UserError,
+        match=re.escape(
+            'Set the `VERCEL_AI_GATEWAY_API_KEY` or `VERCEL_OIDC_TOKEN` environment variable '
+            'or pass the API key via `VercelProvider(api_key=...)` to use the Vercel provider.'
+        ),
+    ):
+        VercelProvider()
+
+
+def test_vercel_pass_openai_client() -> None:
+    openai_client = openai.AsyncOpenAI(api_key='api-key')
+    provider = VercelProvider(openai_client=openai_client)
+    assert provider.client == openai_client
+
+
+def test_vercel_provider_model_profile(mocker: MockerFixture):
+    provider = VercelProvider(api_key='api-key')
+
+    ns = 'pydantic_ai.providers.vercel'
+
+    # Mock all profile functions
+    anthropic_mock = mocker.patch(f'{ns}.anthropic_model_profile', wraps=anthropic_model_profile)
+    amazon_mock = mocker.patch(f'{ns}.amazon_model_profile', wraps=amazon_model_profile)
+    cohere_mock = mocker.patch(f'{ns}.cohere_model_profile', wraps=cohere_model_profile)
+    deepseek_mock = mocker.patch(f'{ns}.deepseek_model_profile', wraps=deepseek_model_profile)
+    google_mock = mocker.patch(f'{ns}.google_model_profile', wraps=google_model_profile)
+    grok_mock = mocker.patch(f'{ns}.grok_model_profile', wraps=grok_model_profile)
+    mistral_mock = mocker.patch(f'{ns}.mistral_model_profile', wraps=mistral_model_profile)
+    openai_mock = mocker.patch(f'{ns}.openai_model_profile', wraps=openai_model_profile)
+
+    # Test openai provider
+    profile = provider.model_profile('openai/gpt-4o')
+    openai_mock.assert_called_with('gpt-4o')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test anthropic provider
+    profile = provider.model_profile('anthropic/claude-3-sonnet')
+    anthropic_mock.assert_called_with('claude-3-sonnet')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test bedrock provider
+    profile = provider.model_profile('bedrock/anthropic.claude-3-sonnet')
+    amazon_mock.assert_called_with('anthropic.claude-3-sonnet')
+    assert profile is not None
+    assert profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
+
+    # Test cohere provider
+    profile = provider.model_profile('cohere/command-r-plus')
+    cohere_mock.assert_called_with('command-r-plus')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test deepseek provider
+    profile = provider.model_profile('deepseek/deepseek-chat')
+    deepseek_mock.assert_called_with('deepseek-chat')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test mistral provider
+    profile = provider.model_profile('mistral/mistral-large')
+    mistral_mock.assert_called_with('mistral-large')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test vertex provider
+    profile = provider.model_profile('vertex/gemini-1.5-pro')
+    google_mock.assert_called_with('gemini-1.5-pro')
+    assert profile is not None
+    assert profile.json_schema_transformer == GoogleJsonSchemaTransformer
+
+    # Test xai provider
+    profile = provider.model_profile('xai/grok-beta')
+    grok_mock.assert_called_with('grok-beta')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+
+def test_vercel_with_http_client():
+    http_client = httpx.AsyncClient()
+    provider = VercelProvider(api_key='test-key', http_client=http_client)
+    assert provider.client.api_key == 'test-key'
+    assert str(provider.client.base_url) == 'https://ai-gateway.vercel.sh/v1/'
+
+
+def test_vercel_provider_invalid_model_name():
+    provider = VercelProvider(api_key='api-key')
+
+    with pytest.raises(UserError, match="Model name must be in 'provider/model' format"):
+        provider.model_profile('invalid-model-name')
+
+
+def test_vercel_provider_unknown_provider():
+    provider = VercelProvider(api_key='api-key')
+
+    profile = provider.model_profile('unknown/gpt-4')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer

tests/test_examples.py

@@ -156,6 +156,7 @@ def print(self, *args: Any, **kwargs: Any) -> None:
     env.set('AWS_ACCESS_KEY_ID', 'testing')
     env.set('AWS_SECRET_ACCESS_KEY', 'testing')
     env.set('AWS_DEFAULT_REGION', 'us-east-1')
+    env.set('VERCEL_AI_GATEWAY_API_KEY', 'testing')
 
     prefix_settings = example.prefix_settings()
     opt_test = prefix_settings.get('test', '')