scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(nebius): route requests to Token Factory

Browse source
This commit is contained in:
demianarc 2026-08-13 17:19:09 +02:00
parent 09889e1986
commit 8812573511
10 changed files with 138 additions and 152 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -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) | ✅ | ✅ | ✅ | | | | | | | |

2
litellm/constants.py
View file

@ -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",

5
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -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"

6
litellm/llms/nebius/chat/transformation.py
View file

@ -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

6
litellm/llms/nebius/embedding/transformation.py
View file

@ -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."""

7
litellm/main.py
View file

@ -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,

10
litellm/provider_endpoints_support_backup.json
View file

@ -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": {

10
provider_endpoints_support.json
View file

@ -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": {

180
tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py
View file

@ -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"

62
tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py
View file

@ -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"
)