From b02bee2e64545c7b2acd2921a32429b0a321b777 Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Sat, 18 Jul 2026 22:36:47 +0000 Subject: [PATCH] feat(novita): add seedream image generation support Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/images/main.py | 1 + .../llms/novita/image_generation/__init__.py | 13 + .../novita/image_generation/transformation.py | 140 +++++++++ ...odel_prices_and_context_window_backup.json | 8 + litellm/utils.py | 6 + model_prices_and_context_window.json | 8 + provider_endpoints_support.json | 2 +- .../test_novita_image_gen_transformation.py | 265 ++++++++++++++++++ 8 files changed, 442 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/novita/image_generation/__init__.py create mode 100644 litellm/llms/novita/image_generation/transformation.py create mode 100644 tests/test_litellm/llms/novita/image_generation/test_novita_image_gen_transformation.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 17ea9aa177b..f2f4709d42d 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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, diff --git a/litellm/llms/novita/image_generation/__init__.py b/litellm/llms/novita/image_generation/__init__.py new file mode 100644 index 00000000000..7a1bddca990 --- /dev/null +++ b/litellm/llms/novita/image_generation/__init__.py @@ -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() diff --git a/litellm/llms/novita/image_generation/transformation.py b/litellm/llms/novita/image_generation/transformation.py new file mode 100644 index 00000000000..e3b963a2bd9 --- /dev/null +++ b/litellm/llms/novita/image_generation/transformation.py @@ -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": ["", ...]} + 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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dedb9bbf40a..1ab7b7e2dc3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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", diff --git a/litellm/utils.py b/litellm/utils.py index 0636d3683b7..e7d468b72b4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e10dde793d1..5d4f28fcc20 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..06b11e294b9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1678,7 +1678,7 @@ "messages": true, "responses": true, "embeddings": false, - "image_generations": false, + "image_generations": true, "audio_transcriptions": false, "audio_speech": false, "moderations": false, diff --git a/tests/test_litellm/llms/novita/image_generation/test_novita_image_gen_transformation.py b/tests/test_litellm/llms/novita/image_generation/test_novita_image_gen_transformation.py new file mode 100644 index 00000000000..f6f71ca22cf --- /dev/null +++ b/tests/test_litellm/llms/novita/image_generation/test_novita_image_gen_transformation.py @@ -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)