scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

feat(novita): add seedream image generation support

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mubashir Osmani 2026-07-18 22:36:47 +00:00
parent 1ebf2a78a9
commit b02bee2e64
8 changed files with 442 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/images/main.py
View file

@ -385,6 +385,7 @@ def image_generation(
#########################################################
elif custom_llm_provider in (
litellm.LlmProviders.RECRAFT,
litellm.LlmProviders.NOVITA,
litellm.LlmProviders.AIML,
litellm.LlmProviders.GEMINI,
litellm.LlmProviders.FAL_AI,

13
litellm/llms/novita/image_generation/__init__.py Normal file
View file

@ -0,0 +1,13 @@
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from .transformation import NovitaImageGenerationConfig
__all__ = [
"NovitaImageGenerationConfig",
]
def get_novita_image_generation_config(model: str) -> BaseImageGenerationConfig:
return NovitaImageGenerationConfig()

140
litellm/llms/novita/image_generation/transformation.py Normal file
View file

@ -0,0 +1,140 @@
from typing import TYPE_CHECKING, Any
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
DEFAULT_API_BASE = "https://api.novita.ai"
CHAT_BASE_SUFFIXES = ("/v3/openai", "/openai")
class NovitaImageGenerationConfig(BaseImageGenerationConfig):
"""
Configuration for Novita AI Seedream image generation.
Synchronous per-model endpoint POST {api_base}/v3/{model} that returns
{"images": ["<url>", ...]}
https://novita.ai/docs/api-reference/model-apis-seedream-4-0
"""
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
return ["n", "size"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
passthrough = {k: v for k, v in non_default_params.items() if k not in optional_params and k != "n"}
n = non_default_params.get("n")
sequential = {"sequential_image_generation": "auto", "max_images": n} if isinstance(n, int) and n > 1 else {}
return {**optional_params, **passthrough, **sequential}
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
stream: bool | None = None,
) -> str:
base_url = (api_base or get_secret_str("NOVITA_API_BASE") or DEFAULT_API_BASE).rstrip("/")
for suffix in CHAT_BASE_SUFFIXES:
if base_url.endswith(suffix):
base_url = base_url[: -len(suffix)]
break
return f"{base_url}/v3/{model}"
def validate_environment(
self,
headers: dict,
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
final_api_key = api_key or get_secret_str("NOVITA_API_KEY")
if not final_api_key:
raise ValueError("NOVITA_API_KEY is not set")
headers["Authorization"] = f"Bearer {final_api_key}"
headers["Content-Type"] = "application/json"
headers["X-Novita-Source"] = "litellm"
return headers
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
return {"prompt": prompt, **optional_params}
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
if raw_response.status_code != 200:
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
try:
response_data = raw_response.json()
except ValueError as e:
raise self.get_error_class(
error_message=f"Failed to parse Novita image generation response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
images = response_data.get("images")
if not images:
raise self.get_error_class(
error_message=f"Novita image generation response missing images: {response_data}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
model_response.data = [ImageObject(url=self._extract_url(item)) for item in images]
return model_response
def _extract_url(self, item: object) -> str | None:
if isinstance(item, str):
return item
if isinstance(item, dict):
value = item.get("image_url") or item.get("url")
return value if isinstance(value, str) else None
return None

8
litellm/model_prices_and_context_window_backup.json
View file

@ -41849,6 +41849,14 @@
"output_cost_per_token": 9e-07,
"supports_function_calling": true
},
"novita/seedream-4.0": {
"litellm_provider": "novita",
"mode": "image_generation",
"source": "https://novita.ai/docs/api-reference/model-apis-seedream-4-0",
"supported_endpoints": [
"/v1/images/generations"
]
},
"novita/deepseek/deepseek-v3.2": {
"litellm_provider": "novita",
"mode": "chat",

6
litellm/utils.py
View file

@ -8579,6 +8579,12 @@ class ProviderConfigManager:
)
return get_recraft_image_generation_config(model)
elif LlmProviders.NOVITA == provider:
from litellm.llms.novita.image_generation import (
get_novita_image_generation_config,
)
return get_novita_image_generation_config(model)
elif LlmProviders.AIML == provider:
from litellm.llms.aiml.image_generation import (
get_aiml_image_generation_config,

8
model_prices_and_context_window.json
View file

@ -41970,6 +41970,14 @@
"output_cost_per_token": 9e-07,
"supports_function_calling": true
},
"novita/seedream-4.0": {
"litellm_provider": "novita",
"mode": "image_generation",
"source": "https://novita.ai/docs/api-reference/model-apis-seedream-4-0",
"supported_endpoints": [
"/v1/images/generations"
]
},
"novita/deepseek/deepseek-v3.2": {
"litellm_provider": "novita",
"mode": "chat",

2
provider_endpoints_support.json
View file

@ -1678,7 +1678,7 @@
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"image_generations": true,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,

265
tests/test_litellm/llms/novita/image_generation/test_novita_image_gen_transformation.py Normal file
View file

@ -0,0 +1,265 @@
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
import litellm
from litellm import get_llm_provider
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.novita.image_generation.transformation import (
DEFAULT_API_BASE,
NovitaImageGenerationConfig,
)
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
MODULE = "litellm.llms.novita.image_generation.transformation.get_secret_str"
class TestNovitaImageGenerationTransformation:
def setup_method(self):
self.config = NovitaImageGenerationConfig()
self.model = "seedream-4.0"
self.logging_obj = MagicMock()
def test_provider_routing(self):
model, provider, _, _ = get_llm_provider("novita/seedream-4.0")
assert provider == "novita"
assert model == "seedream-4.0"
def test_provider_config_registered(self):
config = ProviderConfigManager.get_provider_image_generation_config(
model=self.model,
provider=LlmProviders.NOVITA,
)
assert isinstance(config, NovitaImageGenerationConfig)
def test_supported_params(self):
assert self.config.get_supported_openai_params(self.model) == ["n", "size"]
def test_map_openai_params_passthrough(self):
result = self.config.map_openai_params(
non_default_params={"size": "2048x2048", "watermark": False},
optional_params={},
model=self.model,
drop_params=False,
)
assert result == {"size": "2048x2048", "watermark": False}
def test_map_openai_params_n_expands_to_sequential(self):
result = self.config.map_openai_params(
non_default_params={"n": 4, "size": "1024x1024"},
optional_params={},
model=self.model,
drop_params=False,
)
assert result["size"] == "1024x1024"
assert result["sequential_image_generation"] == "auto"
assert result["max_images"] == 4
assert "n" not in result
def test_map_openai_params_n_one_no_sequential(self):
result = self.config.map_openai_params(
non_default_params={"n": 1},
optional_params={},
model=self.model,
drop_params=False,
)
assert "sequential_image_generation" not in result
assert "n" not in result
@patch(MODULE)
def test_get_complete_url_default(self, mock_secret):
mock_secret.return_value = None
result = self.config.get_complete_url(
api_base=None,
api_key="k",
model=self.model,
optional_params={},
litellm_params={},
)
assert result == f"{DEFAULT_API_BASE}/v3/seedream-4.0"
@patch(MODULE)
def test_get_complete_url_strips_chat_suffix(self, mock_secret):
mock_secret.return_value = None
result = self.config.get_complete_url(
api_base="https://api.novita.ai/v3/openai",
api_key="k",
model=self.model,
optional_params={},
litellm_params={},
)
assert result == "https://api.novita.ai/v3/seedream-4.0"
@patch(MODULE)
def test_get_complete_url_custom_base(self, mock_secret):
mock_secret.return_value = None
result = self.config.get_complete_url(
api_base="https://proxy.example.com",
api_key="k",
model=self.model,
optional_params={},
litellm_params={},
)
assert result == "https://proxy.example.com/v3/seedream-4.0"
@patch(MODULE)
def test_validate_environment_with_api_key(self, mock_secret):
headers = self.config.validate_environment(
headers={},
model=self.model,
messages=[],
optional_params={},
litellm_params={},
api_key="my_key",
)
assert headers["Authorization"] == "Bearer my_key"
assert headers["X-Novita-Source"] == "litellm"
mock_secret.assert_not_called()
@patch(MODULE)
def test_validate_environment_env_fallback(self, mock_secret):
mock_secret.return_value = "env_key"
headers = self.config.validate_environment(
headers={},
model=self.model,
messages=[],
optional_params={},
litellm_params={},
)
assert headers["Authorization"] == "Bearer env_key"
@patch(MODULE)
def test_validate_environment_missing_key_raises(self, mock_secret):
mock_secret.return_value = None
with pytest.raises(ValueError, match="NOVITA_API_KEY is not set"):
self.config.validate_environment(
headers={},
model=self.model,
messages=[],
optional_params={},
litellm_params={},
)
def test_transform_request(self):
result = self.config.transform_image_generation_request(
model=self.model,
prompt="a cat surfing",
optional_params={"size": "1024x1024", "max_images": 3},
litellm_params={},
headers={},
)
assert result == {"prompt": "a cat surfing", "size": "1024x1024", "max_images": 3}
def test_transform_response_url_list(self):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"images": ["https://cdn/img1.png", "https://cdn/img2.png"]}
model_response = litellm.ImageResponse()
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
assert [img.url for img in result.data] == ["https://cdn/img1.png", "https://cdn/img2.png"]
def test_transform_response_image_url_objects(self):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"images": [{"image_url": "https://cdn/obj.png", "image_url_ttl": 3600}]}
model_response = litellm.ImageResponse()
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
assert [img.url for img in result.data] == ["https://cdn/obj.png"]
def test_transform_response_non_200_raises(self):
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.text = "unauthorized"
mock_response.headers = {}
with pytest.raises(BaseLLMException):
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=litellm.ImageResponse(),
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
def test_transform_response_missing_images_raises(self):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.json.return_value = {"task_id": "abc"}
with pytest.raises(BaseLLMException):
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=litellm.ImageResponse(),
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
def test_transform_response_json_parse_error_raises(self):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.json.side_effect = ValueError("bad json")
with pytest.raises(BaseLLMException):
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=litellm.ImageResponse(),
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
def test_image_generation_dispatches_to_novita_handler(self):
fake_response = MagicMock()
with patch.object(
litellm.images.main.llm_http_handler,
"image_generation_handler",
return_value=fake_response,
) as mock_handler:
result = litellm.image_generation(
model="novita/seedream-4.0",
prompt="a cat surfing a wave",
api_key="sk-test",
)
assert result is fake_response
mock_handler.assert_called_once()
kwargs = mock_handler.call_args.kwargs
assert kwargs["custom_llm_provider"] == "novita"
assert kwargs["model"] == "seedream-4.0"
assert kwargs["prompt"] == "a cat surfing a wave"
assert isinstance(kwargs["image_generation_provider_config"], NovitaImageGenerationConfig)