From 88125735118adcd6bfe9ffc1777c1e7001566536 Mon Sep 17 00:00:00 2001 From: demianarc <13074866+demianarc@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:19:09 +0200 Subject: [PATCH] fix(nebius): route requests to Token Factory --- README.md | 2 +- litellm/constants.py | 2 +- .../get_llm_provider_logic.py | 5 +- litellm/llms/nebius/chat/transformation.py | 6 +- .../llms/nebius/embedding/transformation.py | 6 +- litellm/main.py | 7 +- .../provider_endpoints_support_backup.json | 10 +- provider_endpoints_support.json | 10 +- .../nebius/test_nebius_chat_transformation.py | 180 +++++++++--------- .../test_nebius_embedding_transformation.py | 62 +++--- 10 files changed, 138 insertions(+), 152 deletions(-) diff --git a/README.md b/README.md index 32b0160dbaa..8458e88ea3c 100644 --- a/README.md +++ b/README.md @@ -336,7 +336,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ | [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | | | [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | | | [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | | -| [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | | +| [Nebius Token Factory (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | | | [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | | | [Nscale (`nscale`)](https://docs.litellm.ai/docs/providers/nscale) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/litellm/constants.py b/litellm/constants.py index 554165f5d39..c939770a3e5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -731,7 +731,7 @@ openai_compatible_endpoints: Final[list] = [ "api.llama.com/compat/v1/", "api.featherless.ai/v1", "inference.api.nscale.com/v1", - "api.studio.nebius.ai/v1", + "api.tokenfactory.nebius.com/v1", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "https://api-inference.modelscope.cn/v1", "https://api.moonshot.ai/v1", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index dbb40913e14..cd160c919a3 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -290,6 +290,9 @@ def get_llm_provider( elif endpoint == litellm.NscaleConfig.API_BASE_URL: custom_llm_provider = "nscale" dynamic_api_key = litellm.NscaleConfig.get_api_key() + elif endpoint == "api.tokenfactory.nebius.com/v1": + custom_llm_provider = "nebius" + dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif endpoint == "dashscope-intl.aliyuncs.com/compatible-mode/v1": custom_llm_provider = "dashscope" dynamic_api_key = get_secret_str("DASHSCOPE_API_KEY") @@ -605,7 +608,7 @@ def _get_openai_compatible_provider_info( api_base = api_base or get_secret("LLAMA_API_BASE") or "https://api.llama.com/compat/v1" dynamic_api_key = api_key or get_secret_str("LLAMA_API_KEY") elif custom_llm_provider == "nebius": - api_base = api_base or get_secret("NEBIUS_API_BASE") or "https://api.studio.nebius.ai/v1" + api_base = api_base or get_secret("NEBIUS_API_BASE") or "https://api.tokenfactory.nebius.com/v1" dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": api_base = api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" diff --git a/litellm/llms/nebius/chat/transformation.py b/litellm/llms/nebius/chat/transformation.py index cf65699d629..939d7dd5c16 100644 --- a/litellm/llms/nebius/chat/transformation.py +++ b/litellm/llms/nebius/chat/transformation.py @@ -1,8 +1,4 @@ -""" -Nebius AI Studio Chat Completions API - Transformation - -This is OpenAI compatible - no translation needed / occurs -""" +"""Nebius Token Factory Chat Completions API transformation.""" from typing import Final diff --git a/litellm/llms/nebius/embedding/transformation.py b/litellm/llms/nebius/embedding/transformation.py index d56b7def13c..f3214d7be79 100644 --- a/litellm/llms/nebius/embedding/transformation.py +++ b/litellm/llms/nebius/embedding/transformation.py @@ -1,5 +1 @@ -""" -Calls handled in openai/ - -as Nebius AI Studio is an openai-compatible endpoint. -""" +"""Nebius Token Factory embeddings use the OpenAI-compatible handler.""" diff --git a/litellm/main.py b/litellm/main.py index 16eff5a0f3e..e725397a081 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6680,7 +6680,12 @@ def embedding( ) elif custom_llm_provider == "nebius": api_key = api_key or litellm.api_key or get_secret_str("NEBIUS_API_KEY") - api_base = api_base or litellm.api_base or get_secret_str("NEBIUS_API_BASE") or "api.studio.nebius.ai/v1" + api_base = ( + api_base + or litellm.api_base + or get_secret_str("NEBIUS_API_BASE") + or "https://api.tokenfactory.nebius.com/v1" + ) response = openai_chat_completions.embedding( model=model, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..5f1efc8d0c2 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1485,12 +1485,12 @@ } }, "nebius": { - "display_name": "Nebius AI Studio (`nebius`)", + "display_name": "Nebius Token Factory (`nebius`)", "url": "https://docs.litellm.ai/docs/providers/nebius", "endpoints": { "chat_completions": true, - "messages": true, - "responses": true, + "messages": false, + "responses": false, "embeddings": true, "image_generations": false, "audio_transcriptions": false, @@ -1498,8 +1498,8 @@ "moderations": false, "batches": false, "rerank": false, - "a2a": true, - "interactions": true + "a2a": false, + "interactions": false } }, "nlp_cloud": { diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..152aca87418 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1618,12 +1618,12 @@ } }, "nebius": { - "display_name": "Nebius AI Studio (`nebius`)", + "display_name": "Nebius Token Factory (`nebius`)", "url": "https://docs.litellm.ai/docs/providers/nebius", "endpoints": { "chat_completions": true, - "messages": true, - "responses": true, + "messages": false, + "responses": false, "embeddings": true, "image_generations": false, "audio_transcriptions": false, @@ -1631,8 +1631,8 @@ "moderations": false, "batches": false, "rerank": false, - "a2a": true, - "interactions": true + "a2a": false, + "interactions": false } }, "neosantara": { diff --git a/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py b/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py index cb15dd3fa3e..e9d9f165ade 100644 --- a/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py +++ b/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py @@ -1,113 +1,105 @@ -""" -Unit tests for Nebius AI Studio configuration. - -These tests validate the NebiusConfig class which extends OpenAIGPTConfig. -Nebius AI Studio is an OpenAI-compatible provider with minor customizations. -""" - -import os -import sys - -sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the parent directory to the system path +import json import pytest import litellm -from litellm import completion from litellm.llms.nebius.chat.transformation import NebiusConfig +TOKEN_FACTORY_API_BASE = "https://api.tokenfactory.nebius.com/v1" -class TestNebiusConfig: - """Test class for Nebius AI Studio functionality""" - def test_default_api_base(self): - """Test that default API base is used when none is provided""" - config = NebiusConfig() - headers = {} - api_key = "fake-nebius-key" +def test_nebius_provider_uses_token_factory_api_base(monkeypatch): + monkeypatch.delenv("NEBIUS_API_BASE", raising=False) - # Call validate_environment without specifying api_base - result = config.validate_environment( - headers=headers, - model="nebius/Qwen/Qwen3-4B", - messages=[{"role": "user", "content": "Hey"}], - optional_params={}, - litellm_params={}, - api_key=api_key, - api_base=None, # Not providing api_base - ) + model, provider, _, api_base = litellm.get_llm_provider( + model="nebius/moonshotai/Kimi-K3", + api_key="test-key", + ) - # Verify headers are still set correctly - assert result["Authorization"] == f"Bearer {api_key}" - assert result["Content-Type"] == "application/json" + assert model == "moonshotai/Kimi-K3" + assert provider == "nebius" + assert api_base == TOKEN_FACTORY_API_BASE - # We can't directly test the api_base value here since validate_environment - # only returns the headers, but we can verify it doesn't raise an exception - # which would happen if api_base handling was incorrect - @pytest.mark.respx() - def test_nebius_completion_mock(self, respx_mock): - """ - Mock test for Nebius AI Studio completion using the model format from docs. - This test mocks the actual HTTP request to test the integration properly. - """ +def test_token_factory_api_base_resolves_nebius_provider(): + model, provider, _, api_base = litellm.get_llm_provider( + model="moonshotai/Kimi-K3", + api_base=TOKEN_FACTORY_API_BASE, + api_key="test-key", + ) - litellm.disable_aiohttp_transport = ( - True # since this uses respx, we need to set use_aiohttp_transport to False - ) + assert model == "moonshotai/Kimi-K3" + assert provider == "nebius" + assert api_base == TOKEN_FACTORY_API_BASE - # Set up environment variables for the test - api_key = "fake-nebius-key" - api_base = "https://api.studio.nebius.ai/v1" - model = "nebius/Qwen/Qwen3-4B" - model_name = "Qwen3-4B" # The actual model name without provider prefix - # Mock the HTTP request to the Nebius AI Studio API - respx_mock.post(f"{api_base}/chat/completions").respond( - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model_name, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": '```python\nprint("Hey from LiteLLM!")\n```\n\nThis simple Python code prints a greeting message from LiteLLM.', - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - status_code=200, - ) +def test_nebius_api_base_override(monkeypatch): + monkeypatch.setenv("NEBIUS_API_BASE", "https://custom.example/v1") - # Make the actual API call through LiteLLM - response = completion( - model=model, - messages=[ - {"role": "user", "content": "write code for saying hey from LiteLLM"} + _, _, _, api_base = litellm.get_llm_provider( + model="nebius/moonshotai/Kimi-K3", + api_key="test-key", + ) + + assert api_base == "https://custom.example/v1" + + +def test_nebius_maps_max_completion_tokens(): + optional_params = NebiusConfig().map_openai_params( + non_default_params={"max_completion_tokens": 256}, + optional_params={}, + model="moonshotai/Kimi-K3", + drop_params=False, + ) + + assert optional_params == {"max_tokens": 256} + + +def test_nebius_does_not_claim_unimplemented_provider_endpoints(): + from litellm.utils import ProviderConfigManager + + assert ( + ProviderConfigManager.get_provider_responses_api_config(provider="nebius") + is None + ) + + +@pytest.mark.respx() +def test_nebius_completion_targets_token_factory(respx_mock, monkeypatch): + monkeypatch.delenv("NEBIUS_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route = respx_mock.post(f"{TOKEN_FACTORY_API_BASE}/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "moonshotai/Kimi-K3", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from Token Factory", + }, + "finish_reason": "stop", + } ], - api_key=api_key, - api_base=api_base, - ) + "usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + }, + }, + status_code=200, + ) - # Verify response structure - assert response is not None - assert hasattr(response, "choices") - assert len(response.choices) > 0 - assert hasattr(response.choices[0], "message") - assert hasattr(response.choices[0].message, "content") - assert response.choices[0].message.content is not None + response = litellm.completion( + model="nebius/moonshotai/Kimi-K3", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-key", + ) - # Check for specific content in the response - assert "```python" in response.choices[0].message.content - assert "Hey from LiteLLM" in response.choices[0].message.content + assert route.called + assert response.choices[0].message.content == "Hello from Token Factory" + assert route.calls[0].request.headers["authorization"] == "Bearer test-key" + assert json.loads(route.calls[0].request.content)["model"] == "moonshotai/Kimi-K3" diff --git a/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py b/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py index 88f79abbde6..fe0f68a9534 100644 --- a/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py +++ b/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py @@ -1,41 +1,35 @@ -from unittest.mock import patch +import json + +import pytest import litellm - -def mock_embedding_response(*args, **kwargs): - """Mock response mimicking litellm.embedding output.""" - - class MockResponse: - def __init__(self): - self.data = [{"embedding": [0.1, 0.2, 0.3]}] # Example embedding vector - self.usage = litellm.Usage() # Mock Usage object - self.model = kwargs.get("model", "nebius/BAAI/bge-en-icl") - self.object = "embedding" - - def __getitem__(self, key): - return getattr(self, key) - - return MockResponse() +TOKEN_FACTORY_API_BASE = "https://api.tokenfactory.nebius.com/v1" -def test_nebius_embeddings(): - """Mocked test for Nebius embeddings using MagicMock.""" - with patch("litellm.embedding", side_effect=mock_embedding_response) as mock_embed: - response = litellm.embedding( - model="nebius/BAAI/bge-en-icl", - input=["good morning from litellm"], - ) +@pytest.mark.respx() +def test_nebius_embedding_targets_token_factory(respx_mock, monkeypatch): + monkeypatch.delenv("NEBIUS_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route = respx_mock.post(f"{TOKEN_FACTORY_API_BASE}/embeddings").respond( + json={ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "Qwen/Qwen3-Embedding-8B", + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + }, + status_code=200, + ) - # Assertions to verify that the mock was called correctly - mock_embed.assert_called_once_with( - model="nebius/BAAI/bge-en-icl", - input=["good morning from litellm"], - ) + response = litellm.embedding( + model="nebius/Qwen/Qwen3-Embedding-8B", + input=["hello from LiteLLM"], + api_key="test-key", + ) - # Assertions to check the structure of the mocked response - assert isinstance(response.data, list) - assert "embedding" in response.data[0] - assert isinstance(response.data[0]["embedding"], list) - assert response.model == "nebius/BAAI/bge-en-icl" - assert response.object == "embedding" + assert route.called + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert route.calls[0].request.headers["authorization"] == "Bearer test-key" + assert ( + json.loads(route.calls[0].request.content)["model"] == "Qwen/Qwen3-Embedding-8B" + )