diff --git a/litellm/images/main.py b/litellm/images/main.py index a5ae154190a..feac4263bd3 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -407,6 +407,7 @@ def image_generation( # noqa: PLR0915 litellm.LlmProviders.RUNWAYML, litellm.LlmProviders.VERTEX_AI, litellm.LlmProviders.OPENROUTER, + litellm.LlmProviders.VOLCENGINE, ): if image_generation_config is None: raise ValueError( diff --git a/litellm/llms/volcengine/__init__.py b/litellm/llms/volcengine/__init__.py index fc0098e84d9..e70b9449873 100644 --- a/litellm/llms/volcengine/__init__.py +++ b/litellm/llms/volcengine/__init__.py @@ -1,6 +1,6 @@ """ Volcengine LLM Provider -Support for Volcengine (ByteDance) chat, embedding, and responses models. +Support for Volcengine (ByteDance) chat, embedding, image generation, and responses models. """ from .chat.transformation import VolcEngineChatConfig @@ -10,6 +10,7 @@ from .common_utils import ( get_volcengine_headers, ) from .embedding import VolcEngineEmbeddingConfig +from .image_generation import VolcEngineImageGenerationConfig from .responses.transformation import VolcEngineResponsesAPIConfig # For backward compatibility, keep the old class name @@ -19,6 +20,7 @@ __all__ = [ "VolcEngineChatConfig", "VolcEngineConfig", # backward compatibility "VolcEngineEmbeddingConfig", + "VolcEngineImageGenerationConfig", "VolcEngineResponsesAPIConfig", "VolcEngineError", "get_volcengine_base_url", diff --git a/litellm/llms/volcengine/image_generation/__init__.py b/litellm/llms/volcengine/image_generation/__init__.py new file mode 100644 index 00000000000..5dde52a220d --- /dev/null +++ b/litellm/llms/volcengine/image_generation/__init__.py @@ -0,0 +1,13 @@ +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) + +from .transformation import VolcEngineImageGenerationConfig + +__all__ = [ + "VolcEngineImageGenerationConfig", +] + + +def get_volcengine_image_generation_config(model: str) -> BaseImageGenerationConfig: + return VolcEngineImageGenerationConfig() diff --git a/litellm/llms/volcengine/image_generation/transformation.py b/litellm/llms/volcengine/image_generation/transformation.py new file mode 100644 index 00000000000..6fe6ad4acd0 --- /dev/null +++ b/litellm/llms/volcengine/image_generation/transformation.py @@ -0,0 +1,210 @@ +""" +VolcEngine (ByteDance Ark) Image Generation Transformation + +Supports Seedream (即梦) and other image generation models via the +Volcengine Ark API: https://www.volcengine.com/docs/6791/1397048 + +The Ark image generation endpoint is OpenAI-compatible at +POST {api_base}/api/v3/images/generations +""" + +from typing import TYPE_CHECKING, Any, List, Optional, Union + +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 + +from ..common_utils import VolcEngineError, get_volcengine_base_url, get_volcengine_headers + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class VolcEngineImageGenerationConfig(BaseImageGenerationConfig): + """ + Configuration for VolcEngine Ark image generation models (e.g. Seedream / 即梦). + + Reference: https://www.volcengine.com/docs/6791/1397048 + """ + + IMAGE_GENERATION_ENDPOINT: str = "api/v3/images/generations" + + # Volcengine-native params that are not in the OpenAI spec but should be + # forwarded when passed via extra_body or non_default_params. + VOLCENGINE_EXTRA_PARAMS = ( + "output_format", + "watermark", + "guidance_scale", + "seed", + "sequential_image_generation", + "stream", + ) + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + return ["n", "response_format", "size", "quality", "style", "user", "seed"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + for k in non_default_params.keys(): + if k not in optional_params.keys(): + if k in supported_params: + optional_params[k] = non_default_params[k] + elif drop_params: + pass + else: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + f"Set drop_params=True to drop unsupported parameters." + ) + return optional_params + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + base_url = ( + api_base + or get_secret_str("VOLCENGINE_API_BASE") + or get_secret_str("ARK_API_BASE") + or get_volcengine_base_url() + ) + base_url = base_url.rstrip("/") + + if base_url.endswith("/images/generations"): + return base_url + if base_url.endswith("/api/v3"): + return f"{base_url}/images/generations" + return f"{base_url}/{self.IMAGE_GENERATION_ENDPOINT}" + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + final_api_key: Optional[str] = ( + api_key + or litellm_params.get("api_key") + or get_secret_str("ARK_API_KEY") + or get_secret_str("VOLCENGINE_API_KEY") + ) + if not final_api_key: + raise ValueError( + "VolcEngine API key is required. " + "Set ARK_API_KEY / VOLCENGINE_API_KEY or pass api_key." + ) + return get_volcengine_headers(api_key=final_api_key, extra_headers=headers) + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + request_body: dict = { + "model": model, + "prompt": prompt, + } + # Pass through OpenAI-compatible params + for k in ("n", "size", "response_format", "quality", "style", "user", "seed"): + if k in optional_params: + request_body[k] = optional_params[k] + + # Pass through volcengine-native params from extra_body + extra_body = optional_params.get("extra_body") or {} + for k in self.VOLCENGINE_EXTRA_PARAMS: + if k in extra_body: + request_body[k] = extra_body[k] + + return request_body + + 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: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ImageResponse: + try: + response_data = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Error parsing VolcEngine image generation response: {e}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + ## LOGGING + logging_obj.post_call( + input=request_data.get("prompt", ""), + api_key=api_key, + additional_args={"complete_input_dict": request_data}, + original_response=response_data, + ) + + if not model_response.data: + model_response.data = [] + + for image_data in response_data.get("data", []): + model_response.data.append( + ImageObject( + url=image_data.get("url"), + b64_json=image_data.get("b64_json"), + ) + ) + + model_response.created = response_data.get("created", model_response.created) + + return model_response + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> VolcEngineError: + typed_headers: httpx.Headers = ( + headers + if isinstance(headers, httpx.Headers) + else httpx.Headers(headers or {}) + ) + return VolcEngineError( + status_code=status_code, + message=error_message, + headers=typed_headers, + ) diff --git a/litellm/utils.py b/litellm/utils.py index 088ee07d630..5ad30682013 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8874,6 +8874,12 @@ class ProviderConfigManager: ) return get_openrouter_image_generation_config(model) + elif LlmProviders.VOLCENGINE == provider: + from litellm.llms.volcengine.image_generation import ( + get_volcengine_image_generation_config, + ) + + return get_volcengine_image_generation_config(model) return None @staticmethod diff --git a/tests/test_litellm/llms/volcengine/image_generation/__init__.py b/tests/test_litellm/llms/volcengine/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/volcengine/image_generation/test_volcengine_image_generation.py b/tests/test_litellm/llms/volcengine/image_generation/test_volcengine_image_generation.py new file mode 100644 index 00000000000..895648b90e2 --- /dev/null +++ b/tests/test_litellm/llms/volcengine/image_generation/test_volcengine_image_generation.py @@ -0,0 +1,258 @@ +""" +Tests for VolcEngine (ByteDance Ark) Image Generation provider. +""" + +import json +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.llms.volcengine.image_generation.transformation import ( + VolcEngineImageGenerationConfig, +) +from litellm.types.utils import ImageResponse + + +class TestVolcEngineImageGenerationConfig: + def setup_method(self): + self.config = VolcEngineImageGenerationConfig() + + def test_get_supported_openai_params(self): + params = self.config.get_supported_openai_params(model="seedream-3.0") + assert "n" in params + assert "size" in params + assert "response_format" in params + assert "quality" in params + assert "style" in params + assert "user" in params + assert "seed" in params + + def test_map_openai_params_supported(self): + result = self.config.map_openai_params( + non_default_params={"n": 2, "size": "1024x1024"}, + optional_params={}, + model="seedream-3.0", + drop_params=False, + ) + assert result == {"n": 2, "size": "1024x1024"} + + def test_map_openai_params_unsupported_raises(self): + with pytest.raises(ValueError, match="not supported"): + self.config.map_openai_params( + non_default_params={"unsupported_param": "value"}, + optional_params={}, + model="seedream-3.0", + drop_params=False, + ) + + def test_map_openai_params_unsupported_drop(self): + result = self.config.map_openai_params( + non_default_params={"unsupported_param": "value", "n": 1}, + optional_params={}, + model="seedream-3.0", + drop_params=True, + ) + assert result == {"n": 1} + + def test_get_complete_url_default(self): + url = self.config.get_complete_url( + api_base=None, + api_key="test", + model="seedream-3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://ark.cn-beijing.volces.com/api/v3/images/generations" + + def test_get_complete_url_custom_base(self): + url = self.config.get_complete_url( + api_base="https://custom.api.com/api/v3", + api_key="test", + model="seedream-3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.com/api/v3/images/generations" + + def test_get_complete_url_already_has_endpoint(self): + url = self.config.get_complete_url( + api_base="https://custom.api.com/api/v3/images/generations", + api_key="test", + model="seedream-3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.com/api/v3/images/generations" + + def test_get_complete_url_trailing_slash(self): + url = self.config.get_complete_url( + api_base="https://ark.cn-beijing.volces.com/", + api_key="test", + model="seedream-3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://ark.cn-beijing.volces.com/api/v3/images/generations" + + def test_validate_environment(self): + headers = self.config.validate_environment( + headers={}, + model="seedream-3.0", + messages=[], + optional_params={}, + litellm_params={"api_key": "test-key-123"}, + ) + assert headers["Authorization"] == "Bearer test-key-123" + assert headers["Content-Type"] == "application/json" + + def test_validate_environment_no_key_raises(self): + with pytest.raises(ValueError, match="API key is required"): + self.config.validate_environment( + headers={}, + model="seedream-3.0", + messages=[], + optional_params={}, + litellm_params={}, + ) + + def test_transform_image_generation_request_basic(self): + result = self.config.transform_image_generation_request( + model="seedream-3.0", + prompt="a cat in a suit", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert result == { + "model": "seedream-3.0", + "prompt": "a cat in a suit", + } + + def test_transform_image_generation_request_with_params(self): + result = self.config.transform_image_generation_request( + model="seedream-3.0", + prompt="a cat in a suit", + optional_params={ + "n": 2, + "size": "1024x1024", + "response_format": "url", + "quality": "hd", + "seed": 42, + }, + litellm_params={}, + headers={}, + ) + assert result == { + "model": "seedream-3.0", + "prompt": "a cat in a suit", + "n": 2, + "size": "1024x1024", + "response_format": "url", + "quality": "hd", + "seed": 42, + } + + def test_transform_image_generation_response_url(self): + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = { + "created": 1700000000, + "data": [ + {"url": "https://example.com/image1.png"}, + {"url": "https://example.com/image2.png"}, + ], + } + mock_response.status_code = 200 + + logging_obj = MagicMock() + model_response = ImageResponse() + + result = self.config.transform_image_generation_response( + model="seedream-3.0", + raw_response=mock_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"prompt": "a cat"}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].url == "https://example.com/image1.png" + assert result.data[1].url == "https://example.com/image2.png" + assert result.created == 1700000000 + + def test_transform_image_generation_response_b64(self): + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = { + "created": 1700000000, + "data": [ + {"b64_json": "iVBORw0KGgoAAAANSUhEUg=="}, + ], + } + mock_response.status_code = 200 + + logging_obj = MagicMock() + model_response = ImageResponse() + + result = self.config.transform_image_generation_response( + model="seedream-3.0", + raw_response=mock_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"prompt": "a cat"}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "iVBORw0KGgoAAAANSUhEUg==" + + def test_transform_image_generation_response_invalid_json(self): + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.side_effect = ValueError("bad json") + mock_response.status_code = 500 + mock_response.headers = httpx.Headers({}) + + logging_obj = MagicMock() + model_response = ImageResponse() + + with pytest.raises(Exception): + self.config.transform_image_generation_response( + model="seedream-3.0", + raw_response=mock_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"prompt": "a cat"}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +class TestVolcEngineImageGenerationE2E: + """End-to-end tests using litellm.image_generation with mocked HTTP.""" + + def test_image_generation_routes_to_volcengine(self): + """Verify that volcengine/ prefix routes through the llm_http_handler path.""" + from litellm.utils import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="volcengine/seedream-3.0" + ) + assert provider == "volcengine" + assert model == "seedream-3.0" + + def test_provider_config_registered(self): + """Verify VolcEngine image generation config is registered in ProviderConfigManager.""" + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_image_generation_config( + model="seedream-3.0", + provider=LlmProviders.VOLCENGINE, + ) + assert config is not None + assert isinstance(config, VolcEngineImageGenerationConfig)