Skip to content

Commit 6b010ce

Browse files
author
Anton Chaplygin
committed
Fix: address review feedback
1 parent 6614ecd commit 6b010ce

File tree

2 files changed

+8
-6
lines changed

2 files changed

+8
-6
lines changed

pydantic_ai_slim/pydantic_ai/providers/nebius.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,11 @@
1010
from pydantic_ai.models import cached_async_http_client
1111
from pydantic_ai.profiles.deepseek import deepseek_model_profile
1212
from pydantic_ai.profiles.google import google_model_profile
13+
from pydantic_ai.profiles.harmony import harmony_model_profile
1314
from pydantic_ai.profiles.meta import meta_model_profile
1415
from pydantic_ai.profiles.mistral import mistral_model_profile
1516
from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
16-
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
17+
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
1718
from pydantic_ai.profiles.qwen import qwen_model_profile
1819
from pydantic_ai.providers import Provider
1920

@@ -47,7 +48,7 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
4748
'deepseek-ai': deepseek_model_profile,
4849
'qwen': qwen_model_profile,
4950
'google': google_model_profile,
50-
'openai': openai_model_profile,
51+
'openai': harmony_model_profile, # used for gpt-oss models on Nebius
5152
'mistralai': mistral_model_profile,
5253
'moonshotai': moonshotai_model_profile,
5354
}

tests/providers/test_nebius.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@
88
from pydantic_ai.exceptions import UserError
99
from pydantic_ai.profiles.deepseek import deepseek_model_profile
1010
from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile
11+
from pydantic_ai.profiles.harmony import harmony_model_profile
1112
from pydantic_ai.profiles.meta import meta_model_profile
1213
from pydantic_ai.profiles.mistral import mistral_model_profile
1314
from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
14-
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, openai_model_profile
15+
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer
1516
from pydantic_ai.profiles.qwen import qwen_model_profile
1617

1718
from ..conftest import TestEnv, try_import
@@ -71,7 +72,7 @@ def test_nebius_provider_model_profile(mocker: MockerFixture):
7172
deepseek_mock = mocker.patch(f'{ns}.deepseek_model_profile', wraps=deepseek_model_profile)
7273
qwen_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile)
7374
google_mock = mocker.patch(f'{ns}.google_model_profile', wraps=google_model_profile)
74-
openai_mock = mocker.patch(f'{ns}.openai_model_profile', wraps=openai_model_profile)
75+
harmony_mock = mocker.patch(f'{ns}.harmony_model_profile', wraps=harmony_model_profile)
7576
mistral_mock = mocker.patch(f'{ns}.mistral_model_profile', wraps=mistral_model_profile)
7677
moonshotai_mock = mocker.patch(f'{ns}.moonshotai_model_profile', wraps=moonshotai_model_profile)
7778

@@ -99,9 +100,9 @@ def test_nebius_provider_model_profile(mocker: MockerFixture):
99100
assert google_profile is not None
100101
assert google_profile.json_schema_transformer == GoogleJsonSchemaTransformer
101102

102-
# Test openai provider
103+
# Test harmony (for openai gpt-oss) provider
103104
profile = provider.model_profile('openai/gpt-oss-120b')
104-
openai_mock.assert_called_with('gpt-oss-120b')
105+
harmony_mock.assert_called_with('gpt-oss-120b')
105106
assert profile is not None
106107
assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
107108

0 commit comments

Comments
 (0)