scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

feat(volcengine): add image generation support for Ark/Seedream

Browse source 
Add first-class image generation provider for VolcEngine (ByteDance Ark),
supporting Seedream (即梦) models via the llm_http_handler path.

- VolcEngineImageGenerationConfig with request/response transformation
- Auth via ARK_API_KEY / VOLCENGINE_API_KEY env vars
- Supports OpenAI-compatible params: n, size, response_format, quality, style, user, seed
- Forwards volcengine-native params via extra_body: output_format, watermark,
  guidance_scale, sequential_image_generation, stream
- 17 unit tests covering config, routing, and response parsing
- Verified E2E against live Ark API with doubao-seedream-4.0/4.5/5.0
This commit is contained in:
yansaitao 2026-03-26 22:22:32 +08:00
parent bdf4acc472
commit 494f0ed604
7 changed files with 491 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

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

4
litellm/llms/volcengine/__init__.py
View file

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

13
litellm/llms/volcengine/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 VolcEngineImageGenerationConfig
__all__ = [
"VolcEngineImageGenerationConfig",
]
def get_volcengine_image_generation_config(model: str) -> BaseImageGenerationConfig:
return VolcEngineImageGenerationConfig()

210
litellm/llms/volcengine/image_generation/transformation.py Normal file
View file

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

6
litellm/utils.py
View file

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

0
tests/test_litellm/llms/volcengine/image_generation/__init__.py Normal file
View file

258
tests/test_litellm/llms/volcengine/image_generation/test_volcengine_image_generation.py Normal file
View file

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