Add Nebius AI Studio provider support (#3124)

Co-authored-by: Anton Chaplygin <[email protected]>

Author: antoncp

README.md

@@ -39,7 +39,7 @@ We built Pydantic AI with one simple aim: to bring that FastAPI feeling to GenAI
 [Pydantic Validation](https://docs.pydantic.dev/latest/) is the validation layer of the OpenAI SDK, the Google ADK, the Anthropic SDK, LangChain, LlamaIndex, AutoGPT, Transformers, CrewAI, Instructor and many more. _Why use the derivative when you can go straight to the source?_ :smiley:
 
 2. **Model-agnostic**:
-Supports virtually every [model](https://ai.pydantic.dev/models/overview) and provider: OpenAI, Anthropic, Gemini, DeepSeek, Grok, Cohere, Mistral, and Perplexity; Azure AI Foundry, Amazon Bedrock, Google Vertex AI, Ollama, LiteLLM, Groq, OpenRouter, Together AI, Fireworks AI, Cerebras, Hugging Face, GitHub, Heroku, Vercel. If your favorite model or provider is not listed, you can easily implement a [custom model](https://ai.pydantic.dev/models/overview#custom-models).
+Supports virtually every [model](https://ai.pydantic.dev/models/overview) and provider: OpenAI, Anthropic, Gemini, DeepSeek, Grok, Cohere, Mistral, and Perplexity; Azure AI Foundry, Amazon Bedrock, Google Vertex AI, Ollama, LiteLLM, Groq, OpenRouter, Together AI, Fireworks AI, Cerebras, Hugging Face, GitHub, Heroku, Vercel, Nebius. If your favorite model or provider is not listed, you can easily implement a [custom model](https://ai.pydantic.dev/models/overview#custom-models).
 
 3. **Seamless Observability**:
 Tightly [integrates](https://ai.pydantic.dev/logfire) with [Pydantic Logfire](https://pydantic.dev/logfire), our general-purpose OpenTelemetry observability platform, for real-time debugging, evals-based performance monitoring, and behavior, tracing, and cost tracking. If you already have an observability platform that supports OTel, you can [use that too](https://ai.pydantic.dev/logfire#alternative-observability-backends).

docs/api/providers.md

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

docs/index.md

@@ -14,7 +14,7 @@ We built Pydantic AI with one simple aim: to bring that FastAPI feeling to GenAI
 [Pydantic Validation](https://docs.pydantic.dev/latest/) is the validation layer of the OpenAI SDK, the Google ADK, the Anthropic SDK, LangChain, LlamaIndex, AutoGPT, Transformers, CrewAI, Instructor and many more. _Why use the derivative when you can go straight to the source?_ :smiley:
 
 2. **Model-agnostic**:
-Supports virtually every [model](models/overview.md) and provider: OpenAI, Anthropic, Gemini, DeepSeek, Grok, Cohere, Mistral, and Perplexity; Azure AI Foundry, Amazon Bedrock, Google Vertex AI, Ollama, LiteLLM, Groq, OpenRouter, Together AI, Fireworks AI, Cerebras, Hugging Face, GitHub, Heroku, Vercel. If your favorite model or provider is not listed, you can easily implement a [custom model](models/overview.md#custom-models).
+Supports virtually every [model](models/overview.md) and provider: OpenAI, Anthropic, Gemini, DeepSeek, Grok, Cohere, Mistral, and Perplexity; Azure AI Foundry, Amazon Bedrock, Google Vertex AI, Ollama, LiteLLM, Groq, OpenRouter, Together AI, Fireworks AI, Cerebras, Hugging Face, GitHub, Heroku, Vercel, Nebius. If your favorite model or provider is not listed, you can easily implement a [custom model](models/overview.md#custom-models).
 
 3. **Seamless Observability**:
 Tightly [integrates](logfire.md) with [Pydantic Logfire](https://pydantic.dev/logfire), our general-purpose OpenTelemetry observability platform, for real-time debugging, evals-based performance monitoring, and behavior, tracing, and cost tracking. If you already have an observability platform that supports OTel, you can [use that too](logfire.md#alternative-observability-backends).

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,102 @@
+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.harmony import harmony_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
+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': harmony_model_profile,  # used for gpt-oss models on Nebius
+            '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,131 @@
+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.harmony import harmony_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
+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)
+    harmony_mock = mocker.patch(f'{ns}.harmony_model_profile', wraps=harmony_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 harmony (for openai gpt-oss) provider
+    profile = provider.model_profile('openai/gpt-oss-120b')
+    harmony_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')