Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
2 changes: 2 additions & 0 deletions docs/api/providers.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,5 @@
::: pydantic_ai.providers.ollama.OllamaProvider

::: pydantic_ai.providers.litellm.LiteLLMProvider

::: pydantic_ai.providers.nebius.NebiusProvider
2 changes: 1 addition & 1 deletion docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
32 changes: 32 additions & 0 deletions docs/models/openai.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
```
1 change: 1 addition & 0 deletions docs/models/overview.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions pydantic_ai_slim/pydantic_ai/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model: # noqa: C901
'together',
'vercel',
'litellm',
'nebius',
):
from .openai import OpenAIChatModel

Expand Down
5 changes: 4 additions & 1 deletion pydantic_ai_slim/pydantic_ai/models/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,7 @@ def __init__(
'together',
'vercel',
'litellm',
'nebius',
]
| Provider[AsyncOpenAI] = 'openai',
profile: ModelProfileSpec | None = None,
Expand Down Expand Up @@ -312,6 +313,7 @@ def __init__(
'together',
'vercel',
'litellm',
'nebius',
]
| Provider[AsyncOpenAI] = 'openai',
profile: ModelProfileSpec | None = None,
Expand Down Expand Up @@ -339,6 +341,7 @@ def __init__(
'together',
'vercel',
'litellm',
'nebius',
]
| Provider[AsyncOpenAI] = 'openai',
profile: ModelProfileSpec | None = None,
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions pydantic_ai_slim/pydantic_ai/providers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}')

Expand Down
102 changes: 102 additions & 0 deletions pydantic_ai_slim/pydantic_ai/providers/nebius.py
Original file line number Diff line number Diff line change
@@ -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,
}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are these all of the providers Nebius supports?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, all currently supported Nebius providers are included here. They are listed at https://studio.nebius.com/


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)
131 changes: 131 additions & 0 deletions tests/providers/test_nebius.py
Original file line number Diff line number Diff line change
@@ -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')
2 changes: 2 additions & 0 deletions tests/providers/test_provider_names.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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():
Expand Down
1 change: 1 addition & 0 deletions tests/test_examples.py
Original file line number Diff line number Diff line change
Expand Up @@ -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', '')
Expand Down