Add Nebius AI Studio provider support

Author: Anton Chaplygin

docs/api/providers.md

@@ -41,3 +41,5 @@
 ::: pydantic_ai.providers.ollama.OllamaProvider
 
 ::: pydantic_ai.providers.litellm.LiteLLMProvider
+
+::: pydantic_ai.providers.nebius.NebiusProvider

docs/models/openai.md

@@ -608,3 +608,35 @@ print(result.output)
 #> The capital of France is Paris.
 ...
 ```
+
+### Nebius AI Studio
+
+Go to [Nebius AI Studio](https://studio.nebius.com/) and create an API key.
+
+Once you've set the `NEBIUS_API_KEY` environment variable, you can run the following:
+
+```python
+from pydantic_ai import Agent
+
+agent = Agent('nebius:Qwen/Qwen3-32B-fast')
+result = agent.run_sync('What is the capital of France?')
+print(result.output)
+#> The capital of France is Paris.
+```
+
+If you need to configure the provider, you can use the [`NebiusProvider`][pydantic_ai.providers.nebius.NebiusProvider] class:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.nebius import NebiusProvider
+
+model = OpenAIChatModel(
+    'Qwen/Qwen3-32B-fast',
+    provider=NebiusProvider(api_key='your-nebius-api-key'),
+)
+agent = Agent(model)
+result = agent.run_sync('What is the capital of France?')
+print(result.output)
+#> The capital of France is Paris.
+```

docs/models/overview.md

@@ -28,6 +28,7 @@ In addition, many providers are compatible with the OpenAI API, and can be used
 - [GitHub Models](openai.md#github-models)
 - [Cerebras](openai.md#cerebras)
 - [LiteLLM](openai.md#litellm)
+- [Nebius AI Studio](openai.md#nebius-ai-studio)
 
 Pydantic AI also comes with [`TestModel`](../api/models/test.md) and [`FunctionModel`](../api/models/function.md)
 for testing and development.

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -691,6 +691,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'together',
         'vercel',
         'litellm',
+        'nebius',
     ):
         from .openai import OpenAIChatModel
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -284,6 +284,7 @@ def __init__(
             'together',
             'vercel',
             'litellm',
+            'nebius',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -312,6 +313,7 @@ def __init__(
             'together',
             'vercel',
             'litellm',
+            'nebius',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -339,6 +341,7 @@ def __init__(
             'together',
             'vercel',
             'litellm',
+            'nebius',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -899,7 +902,7 @@ def __init__(
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together']
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'nebius']
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -142,6 +142,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .litellm import LiteLLMProvider
 
         return LiteLLMProvider
+    elif provider == 'nebius':
+        from .nebius import NebiusProvider
+
+        return NebiusProvider
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')
 

pydantic_ai_slim/pydantic_ai/providers/nebius.py

@@ -0,0 +1,101 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+import httpx
+
+from pydantic_ai import ModelProfile
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.profiles.qwen import qwen_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 Nebius provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class NebiusProvider(Provider[AsyncOpenAI]):
+    """Provider for Nebius AI Studio API."""
+
+    @property
+    def name(self) -> str:
+        return 'nebius'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://api.studio.nebius.com/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        provider_to_profile = {
+            'meta-llama': meta_model_profile,
+            'deepseek-ai': deepseek_model_profile,
+            'qwen': qwen_model_profile,
+            'google': google_model_profile,
+            'openai': openai_model_profile,
+            'mistralai': mistral_model_profile,
+            'moonshotai': moonshotai_model_profile,
+        }
+
+        profile = None
+
+        try:
+            model_name = model_name.lower()
+            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 NebiusProvider is always used with OpenAIChatModel, 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: httpx.AsyncClient) -> 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: httpx.AsyncClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('NEBIUS_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `NEBIUS_API_KEY` environment variable or pass it via '
+                '`NebiusProvider(api_key=...)` to use the Nebius AI Studio 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='nebius')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

tests/providers/test_nebius.py

@@ -0,0 +1,130 @@
+import re
+
+import httpx
+import pytest
+from pytest_mock import MockerFixture
+
+from pydantic_ai._json_schema import InlineDefsJsonSchemaTransformer
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, openai_model_profile
+from pydantic_ai.profiles.qwen import qwen_model_profile
+
+from ..conftest import TestEnv, try_import
+
+with try_import() as imports_successful:
+    import openai
+
+    from pydantic_ai.providers.nebius import NebiusProvider
+
+
+pytestmark = [
+    pytest.mark.skipif(not imports_successful(), reason='openai not installed'),
+    pytest.mark.vcr,
+    pytest.mark.anyio,
+]
+
+
+def test_nebius_provider():
+    provider = NebiusProvider(api_key='api-key')
+    assert provider.name == 'nebius'
+    assert provider.base_url == 'https://api.studio.nebius.com/v1'
+    assert isinstance(provider.client, openai.AsyncOpenAI)
+    assert provider.client.api_key == 'api-key'
+
+
+def test_nebius_provider_need_api_key(env: TestEnv) -> None:
+    env.remove('NEBIUS_API_KEY')
+    with pytest.raises(
+        UserError,
+        match=re.escape(
+            'Set the `NEBIUS_API_KEY` environment variable or pass it via '
+            '`NebiusProvider(api_key=...)` to use the Nebius AI Studio provider.'
+        ),
+    ):
+        NebiusProvider()
+
+
+def test_nebius_pass_openai_client() -> None:
+    openai_client = openai.AsyncOpenAI(api_key='api-key')
+    provider = NebiusProvider(openai_client=openai_client)
+    assert provider.client == openai_client
+
+
+def test_nebius_provider_pass_http_client() -> None:
+    http_client = httpx.AsyncClient()
+    provider = NebiusProvider(http_client=http_client, api_key='api-key')
+    assert provider.client._client == http_client  # type: ignore[reportPrivateUsage]
+
+
+def test_nebius_provider_model_profile(mocker: MockerFixture):
+    provider = NebiusProvider(api_key='api-key')
+
+    ns = 'pydantic_ai.providers.nebius'
+
+    # Mock all profile functions
+    meta_mock = mocker.patch(f'{ns}.meta_model_profile', wraps=meta_model_profile)
+    deepseek_mock = mocker.patch(f'{ns}.deepseek_model_profile', wraps=deepseek_model_profile)
+    qwen_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile)
+    google_mock = mocker.patch(f'{ns}.google_model_profile', wraps=google_model_profile)
+    openai_mock = mocker.patch(f'{ns}.openai_model_profile', wraps=openai_model_profile)
+    mistral_mock = mocker.patch(f'{ns}.mistral_model_profile', wraps=mistral_model_profile)
+    moonshotai_mock = mocker.patch(f'{ns}.moonshotai_model_profile', wraps=moonshotai_model_profile)
+
+    # Test meta provider
+    meta_profile = provider.model_profile('meta-llama/Llama-3.3-70B-Instruct')
+    meta_mock.assert_called_with('llama-3.3-70b-instruct')
+    assert meta_profile is not None
+    assert meta_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
+
+    # Test deepseek provider
+    profile = provider.model_profile('deepseek-ai/DeepSeek-R1-0528')
+    deepseek_mock.assert_called_with('deepseek-r1-0528')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test qwen provider
+    qwen_profile = provider.model_profile('Qwen/Qwen3-30B-A3B')
+    qwen_mock.assert_called_with('qwen3-30b-a3b')
+    assert qwen_profile is not None
+    assert qwen_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
+
+    # Test google provider
+    google_profile = provider.model_profile('google/gemma-2-2b-it')
+    google_mock.assert_called_with('gemma-2-2b-it')
+    assert google_profile is not None
+    assert google_profile.json_schema_transformer == GoogleJsonSchemaTransformer
+
+    # Test openai provider
+    profile = provider.model_profile('openai/gpt-oss-120b')
+    openai_mock.assert_called_with('gpt-oss-120b')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test mistral provider
+    profile = provider.model_profile('mistralai/Devstral-Small-2505')
+    mistral_mock.assert_called_with('devstral-small-2505')
+    assert profile is not None
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test moonshotai provider
+    moonshotai_profile = provider.model_profile('moonshotai/Kimi-K2-Instruct')
+    moonshotai_mock.assert_called_with('kimi-k2-instruct')
+    assert moonshotai_profile is not None
+    assert moonshotai_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test unknown provider
+    unknown_profile = provider.model_profile('unknown-provider/unknown-model')
+    assert unknown_profile is not None
+    assert unknown_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+
+def test_nebius_provider_invalid_model_name():
+    provider = NebiusProvider(api_key='api-key')
+
+    with pytest.raises(UserError, match="Model name must be in 'provider/model' format"):
+        provider.model_profile('invalid-model-name')

tests/providers/test_provider_names.py

@@ -28,6 +28,7 @@
     from pydantic_ai.providers.litellm import LiteLLMProvider
     from pydantic_ai.providers.mistral import MistralProvider
     from pydantic_ai.providers.moonshotai import MoonshotAIProvider
+    from pydantic_ai.providers.nebius import NebiusProvider
     from pydantic_ai.providers.ollama import OllamaProvider
     from pydantic_ai.providers.openai import OpenAIProvider
     from pydantic_ai.providers.openrouter import OpenRouterProvider
@@ -54,6 +55,7 @@
         ('github', GitHubProvider, 'GITHUB_API_KEY'),
         ('ollama', OllamaProvider, 'OLLAMA_BASE_URL'),
         ('litellm', LiteLLMProvider, None),
+        ('nebius', NebiusProvider, 'NEBIUS_API_KEY'),
     ]
 
 if not imports_successful():

tests/test_examples.py

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