This commit is contained in:
Rongkun Yan 2026-09-16 18:02:37 +08:00 committed by GitHub
commit a5e07be000
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1093 additions and 264 deletions

View file

@ -1273,6 +1273,8 @@ modelscope_models: Final[set] = set(
"Qwen/QwQ-32B-Preview",
"Qwen/QVQ-72B-Preview",
"Qwen/Qwen-Image-Edit",
"Qwen/Qwen-Image-Edit-2511",
"Qwen/Qwen-Image-2512",
# DeepSeek series models
"deepseek-ai/DeepSeek-R1-0528",
"deepseek-ai/DeepSeek-R1-Distill-Llama-70B",

View file

@ -34,6 +34,9 @@ from openai.types.audio.transcription_create_params import FileTypes
# BFL handlers
from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
# ModelScope handler
from litellm.llms.modelscope.image_generation.handler import modelscope_image_generation
from litellm.main import (
azure_chat_completions,
base_llm_aiohttp_handler,
@ -424,6 +427,25 @@ def image_generation(
client=client,
aimg_generation=aimg_generation,
)
elif custom_llm_provider == "modelscope":
if model is None:
raise Exception("Model needs to be set for modelscope")
if api_base is not None and api_key is None:
raise ValueError("api_key must be provided when api_base is overridden for modelscope")
litellm_params_dict["api_key"] = api_key
litellm_params_dict["api_base"] = api_base
return modelscope_image_generation.image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params_dict,
logging_obj=litellm_logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
aimg_generation=aimg_generation,
)
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,

View file

@ -0,0 +1,24 @@
"""
ModelScope Common Utilities
Shared constants and error handling for ModelScope API integration.
"""
from typing import Final
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class ModelScopeError(BaseLLMException):
"""Exception class for ModelScope API errors."""
DEFAULT_POLLING_INTERVAL: Final = 2.0
DEFAULT_MAX_POLLING_TIME: Final = 300
ASYNC_MODE_HEADER: Final = "X-ModelScope-Async-Mode"
TASK_TYPE_HEADER: Final = "X-ModelScope-Task-Type"
IMAGE_GENERATION_TASK_TYPE: Final = "image_generation"
TASK_STATUS_SUCCEED: Final = "SUCCEED"
TASK_STATUS_FAILED: Final = "FAILED"

View file

@ -0,0 +1,411 @@
"""
ModelScope Image Generation Handler
ModelScope image generation only supports async mode: the submit call returns
a task_id, then the caller polls GET /v1/tasks/{task_id} until task_status is
SUCCEED (image URLs in output_images) or FAILED.
API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro
"""
import asyncio
import json
import time
from typing import Final
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client, # pyright: ignore[reportPrivateUsage] # litellm internal client factory, same pattern as BFL handler
get_async_httpx_client,
)
from litellm.llms.modelscope.common_utils import (
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
TASK_STATUS_FAILED,
TASK_STATUS_SUCCEED,
ModelScopeError,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
from .transformation import ModelScopeImageGenerationConfig
class ModelScopeImageGeneration:
"""ModelScope image generation handler."""
def __init__(self) -> None:
self.config: Final = ModelScopeImageGenerationConfig()
def image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature
logging_obj: LiteLLMLoggingObj,
timeout: float | httpx.Timeout | None,
extra_headers: dict[str, str] | None = None, # mutable-ok: litellm override signature
client: HTTPHandler | AsyncHTTPHandler | None = None,
aimg_generation: bool = False,
) -> ImageResponse:
if aimg_generation:
return self.async_image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client if isinstance(client, AsyncHTTPHandler) else None,
)
api_key: Final = litellm_params.get("api_key") if isinstance(litellm_params, dict) else litellm_params.api_key
api_base: Final = (
litellm_params.get("api_base") if isinstance(litellm_params, dict) else litellm_params.api_base
)
litellm_params_dict: Final = (
litellm_params
if isinstance(litellm_params, dict)
else dict(litellm_params) # mutable-ok: litellm passes a mutable dict
)
sync_client: Final = client if isinstance(client, HTTPHandler) else _get_httpx_client()
headers: Final = self.config.validate_environment(
api_key=api_key,
headers={}, # mutable-ok: validate_environment returns a new dict
model=model,
messages=[], # mutable-ok: required by base class signature
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
complete_url: Final = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
data: Final = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={ # mutable-ok: litellm logging contract
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
try:
response: Final = sync_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except httpx.HTTPError as e:
raise ModelScopeError(
status_code=500,
message=f"Request failed: {e}",
)
final_response: Final = self._poll_for_result_sync(
initial_response=response, # pyright: ignore[reportArgumentType] # post() returns Response | None; None raises HTTPError above
api_base=api_base,
headers=headers,
sync_client=sync_client,
timeout=timeout,
)
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=data,
optional_params=optional_params,
litellm_params=litellm_params_dict,
encoding=None,
)
async def async_image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature
logging_obj: LiteLLMLoggingObj,
timeout: float | httpx.Timeout | None,
extra_headers: dict[str, str] | None = None, # mutable-ok: litellm override signature
client: AsyncHTTPHandler | None = None,
) -> ImageResponse:
api_key: Final = litellm_params.get("api_key") if isinstance(litellm_params, dict) else litellm_params.api_key
api_base: Final = (
litellm_params.get("api_base") if isinstance(litellm_params, dict) else litellm_params.api_base
)
litellm_params_dict: Final = (
litellm_params
if isinstance(litellm_params, dict)
else dict(litellm_params) # mutable-ok: litellm passes a mutable dict
)
async_client: Final = client or get_async_httpx_client(
llm_provider=litellm.LlmProviders.MODELSCOPE,
)
headers: Final = self.config.validate_environment(
api_key=api_key,
headers={}, # mutable-ok: validate_environment returns a new dict
model=model,
messages=[], # mutable-ok: required by base class signature
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
complete_url: Final = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
data: Final = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={ # mutable-ok: litellm logging contract
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
try:
response: Final = await async_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except httpx.HTTPError as e:
raise ModelScopeError(
status_code=500,
message=f"Request failed: {e}",
)
final_response: Final = await self._poll_for_result_async(
initial_response=response, # pyright: ignore[reportArgumentType] # post() returns Response | None; None raises HTTPError above
api_base=api_base,
headers=headers,
async_client=async_client,
timeout=timeout,
)
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=data,
optional_params=optional_params,
litellm_params=litellm_params_dict,
encoding=None,
)
def _poll_for_result_sync(
self,
initial_response: httpx.Response,
api_base: str | None,
headers: dict, # mutable-ok: caller passes a mutable dict
sync_client: HTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: float | httpx.Timeout | None = None,
) -> httpx.Response:
if initial_response.status_code >= 400:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"ModelScope submit failed: {initial_response.text}",
)
try:
response_data: Final = initial_response.json()
except json.JSONDecodeError as e:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"Error parsing submit response: {e}",
)
if "errors" in response_data:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"ModelScope error: {response_data['errors']}",
)
task_id: Final = response_data.get("task_id")
if not task_id:
raise ModelScopeError(
status_code=500,
message="No task_id in ModelScope submit response",
)
polling_url: Final = self.config.get_task_status_url(api_base, task_id)
polling_headers: Final = self.config.get_polling_headers(headers)
start_time: Final = time.time()
verbose_logger.debug("ModelScope starting sync polling at %s", polling_url)
while time.time() - start_time < max_wait:
response = sync_client.get(
url=polling_url,
headers=polling_headers,
timeout=timeout,
)
if response.status_code >= 400:
raise ModelScopeError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
try:
data = response.json()
except json.JSONDecodeError as e:
raise ModelScopeError(
status_code=response.status_code,
message=f"Error parsing poll response: {e}",
)
status = data.get("task_status")
verbose_logger.debug("ModelScope poll status: %s", status)
if status == TASK_STATUS_SUCCEED:
return response
elif status == TASK_STATUS_FAILED:
raise ModelScopeError(
status_code=400,
message="ModelScope image generation task failed",
)
time.sleep(interval)
raise ModelScopeError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
initial_response: httpx.Response,
api_base: str | None,
headers: dict, # mutable-ok: caller passes a mutable dict
async_client: AsyncHTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: float | httpx.Timeout | None = None,
) -> httpx.Response:
if initial_response.status_code >= 400:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"ModelScope submit failed: {initial_response.text}",
)
try:
response_data: Final = initial_response.json()
except json.JSONDecodeError as e:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"Error parsing submit response: {e}",
)
if "errors" in response_data:
raise ModelScopeError(
status_code=initial_response.status_code,
message=f"ModelScope error: {response_data['errors']}",
)
task_id: Final = response_data.get("task_id")
if not task_id:
raise ModelScopeError(
status_code=500,
message="No task_id in ModelScope submit response",
)
polling_url: Final = self.config.get_task_status_url(api_base, task_id)
polling_headers: Final = self.config.get_polling_headers(headers)
start_time: Final = time.time()
verbose_logger.debug("ModelScope starting async polling at %s", polling_url)
while time.time() - start_time < max_wait:
response = await async_client.get(
url=polling_url,
headers=polling_headers,
timeout=timeout,
)
if response.status_code >= 400:
raise ModelScopeError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
try:
data = response.json()
except json.JSONDecodeError as e:
raise ModelScopeError(
status_code=response.status_code,
message=f"Error parsing poll response: {e}",
)
status = data.get("task_status")
verbose_logger.debug("ModelScope poll status: %s", status)
if status == TASK_STATUS_SUCCEED:
return response
elif status == TASK_STATUS_FAILED:
raise ModelScopeError(
status_code=400,
message="ModelScope image generation task failed",
)
await asyncio.sleep(interval)
raise ModelScopeError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
modelscope_image_generation: Final = ModelScopeImageGeneration()

View file

@ -15,6 +15,14 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.modelscope.common_utils import (
ASYNC_MODE_HEADER,
IMAGE_GENERATION_TASK_TYPE,
TASK_STATUS_FAILED,
TASK_STATUS_SUCCEED,
TASK_TYPE_HEADER,
ModelScopeError,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
@ -25,86 +33,68 @@ from litellm.types.utils import ImageObject, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
LiteLLMLoggingObj: Final = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = object
LiteLLMLoggingObj: Final = object
class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
"""
Configuration for ModelScope image generation.
"""Configuration for ModelScope image generation and editing models."""
Supports text-to-image models like:
- Qwen/Qwen-Image-Edit
- And other ModelScope-hosted image generation models
"""
DEFAULT_BASE_URL: Final[str] = "https://api-inference.modelscope.cn/v1"
DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1"
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
"""
Return list of OpenAI params supported by ModelScope.
ModelScope supports standard OpenAI image generation parameters.
"""
return [
"n", # Number of images to generate
"size", # Size of the generated images
"response_format", # url or b64_json
"user", # User identifier
]
@override
def get_supported_openai_params(
self, model: str
) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: litellm override signature
return ["size"] # mutable-ok: small fixed list
@override
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
non_default_params: dict, # mutable-ok: litellm override signature
optional_params: dict, # mutable-ok: litellm override signature
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to ModelScope parameters.
ModelScope uses the same parameter names as OpenAI.
"""
) -> dict: # mutable-ok: litellm override signature
supported_params: Final = self.get_supported_openai_params(model)
if drop_params:
non_default_params = {k: v for k, v in non_default_params.items() if k in supported_params}
non_default_params = {
k: v for k, v in non_default_params.items() if k in supported_params
} # rebind-ok: filtering in place is the litellm pattern # mutable-ok: dict comprehension for filtering
optional_params.update(non_default_params)
return optional_params
def _get_base_url(self, api_base: str | None) -> str:
base_url: Final[str] = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
return base_url.rstrip("/")
@override
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: dict, # mutable-ok: litellm override signature
stream: bool | None = None,
) -> str:
"""
Get the complete URL for the ModelScope image generation API request.
"""
base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
base_url = base_url.rstrip("/")
return f"{self._get_base_url(api_base)}/images/generations"
# Return the images endpoint
return f"{base_url}/images/generations"
def get_task_status_url(self, api_base: str | None, task_id: str) -> str:
return f"{self._get_base_url(api_base)}/tasks/{task_id}"
@override
def validate_environment(
self,
headers: dict,
headers: dict, # mutable-ok: litellm override signature
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
messages: list[AllMessageValues], # mutable-ok: litellm override signature
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: dict, # mutable-ok: litellm override signature
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
"""
Validate environment and set up headers for ModelScope.
"""
) -> dict: # mutable-ok: litellm override signature
final_api_key: Final[str | None] = api_key or get_secret_str("MODELSCOPE_API_KEY")
if not final_api_key:
@ -112,36 +102,40 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
"MODELSCOPE_API_KEY is not set. Please set it via environment variable or pass api_key parameter."
)
default_headers: Final = {
default_headers: Final = { # mutable-ok: http header dict
"Content-Type": "application/json",
"Authorization": f"Bearer {final_api_key}",
ASYNC_MODE_HEADER: "true",
}
headers = {**headers, **default_headers}
return headers
return {**headers, **default_headers} # mutable-ok: merged header dict
def get_polling_headers(
self,
headers: dict, # mutable-ok: caller passes a mutable dict
) -> dict: # mutable-ok: builds a new dict for polling
return { # mutable-ok: polling header dict
"Authorization": headers.get("Authorization", ""),
TASK_TYPE_HEADER: IMAGE_GENERATION_TASK_TYPE,
}
@override
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform OpenAI-style request to ModelScope request format.
ModelScope uses the same format as OpenAI for image generation.
"""
# Build the request body (same as OpenAI)
request_data: Final[dict] = {
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: dict, # mutable-ok: litellm override signature
headers: dict, # mutable-ok: litellm override signature
) -> dict: # mutable-ok: litellm override signature
request_data: Final = { # mutable-ok: request body dict
"model": model,
"prompt": prompt,
}
# Add optional params
for key, value in optional_params.items():
if key.startswith("_"):
extra_body: Final = optional_params.get("extra_body") or {} # mutable-ok: provider fields dict
for key, value in {**optional_params, **extra_body}.items(): # mutable-ok: merged params dict
if key in ("extra_body", "extra_headers", "extra_query", "model", "prompt") or key.startswith("_"):
continue
request_data[key] = value
@ -154,19 +148,13 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
request_data: dict, # mutable-ok: litellm override signature
optional_params: dict, # mutable-ok: litellm override signature
litellm_params: dict, # mutable-ok: litellm override signature
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
"""
Transform ModelScope response to OpenAI-compatible ImageResponse.
ModelScope returns the same format as OpenAI:
{"created": timestamp, "data": [{"url": "..."}]}
"""
try:
response_data: Final = raw_response.json()
except Exception as e:
@ -176,64 +164,52 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
headers=raw_response.headers,
)
# Check for errors in response
if "error" in response_data:
error_msg: Final = response_data["error"].get("message", str(response_data["error"]))
if "errors" in response_data:
errors: Final = response_data["errors"]
error_msg: Final = errors.get("message", str(errors)) if isinstance(errors, dict) else str(errors)
raise self.get_error_class(
error_message=f"ModelScope error: {error_msg}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# Extract images from response
data_list: Final = response_data.get("data", [])
if not model_response.data:
model_response.data = []
for item in data_list:
image_obj = ImageObject(
url=item.get("url"),
b64_json=item.get("b64_json"),
revised_prompt=item.get("revised_prompt"),
task_status: Final = response_data.get("task_status")
if task_status == TASK_STATUS_FAILED:
raise self.get_error_class(
error_message="ModelScope image generation task failed",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
if task_status != TASK_STATUS_SUCCEED:
raise self.get_error_class(
error_message=f"ModelScope task did not succeed: status={task_status}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
output_images: Final = response_data.get("output_images", []) or [] # mutable-ok: API response list
for image_url in output_images:
model_response.data.append(ImageObject(url=image_url))
if not model_response.data:
raise self.get_error_class(
error_message="ModelScope task SUCCEED but no output_images",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
model_response.data.append(image_obj)
return model_response
@override
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict | httpx.Headers,
headers: dict | httpx.Headers, # mutable-ok: litellm override signature
) -> BaseLLMException:
"""Return the appropriate error class for ModelScope."""
from litellm.exceptions import (
AuthenticationError,
BadRequestError,
InternalServerError,
return ModelScopeError(
status_code=status_code,
message=error_message,
headers=headers,
)
if status_code == 400:
return BadRequestError(
message=error_message,
model="",
llm_provider="modelscope",
)
elif status_code == 401:
return AuthenticationError(
message=error_message,
model="",
llm_provider="modelscope",
)
elif status_code >= 500:
return InternalServerError(
message=error_message,
model="",
llm_provider="modelscope",
)
else:
return BadRequestError(
message=error_message,
model="",
llm_provider="modelscope",
)

View file

@ -5,11 +5,14 @@ These tests validate the ModelScopeImageGenerationConfig class which handles
transformation between OpenAI-compatible format and ModelScope API format.
"""
from unittest.mock import MagicMock, patch
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.modelscope.common_utils import ModelScopeError
from litellm.llms.modelscope.image_generation.transformation import (
ModelScopeImageGenerationConfig,
)
@ -24,47 +27,39 @@ class TestModelScopeImageGenerationTransformation:
self.logging_obj = MagicMock()
def test_get_supported_openai_params(self):
"""Test that get_supported_openai_params returns correct parameters."""
"""Only size is a documented+honored OpenAI param for ModelScope."""
supported_params = self.config.get_supported_openai_params(self.model)
assert "n" in supported_params
assert "size" in supported_params
assert "response_format" in supported_params
assert "user" in supported_params
# n is silently ignored (n=2 returns 1 image); user/response_format are
# not honored by ModelScope.
assert "n" not in supported_params
assert "user" not in supported_params
assert "response_format" not in supported_params
def test_map_openai_params(self):
"""Test that map_openai_params correctly passes through parameters."""
non_default_params = {
"n": 2,
"size": "1024x1024",
"response_format": "url",
}
optional_params = {}
"""Supported OpenAI param (size) passes through to optional_params."""
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
non_default_params={"size": "1024x1024"},
optional_params={},
model=self.model,
drop_params=False,
)
assert result["n"] == 2
assert result["size"] == "1024x1024"
assert result["response_format"] == "url"
def test_map_openai_params_with_user(self):
"""Test that map_openai_params correctly passes through user parameter."""
non_default_params = {"user": "test-user-123"}
optional_params = {}
def test_map_openai_params_drops_unsupported(self):
"""Unsupported OpenAI params (e.g. quality) are dropped when
drop_params=True; ModelScope only honors `size`."""
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
non_default_params={"size": "1024x1024", "quality": "hd"},
optional_params={},
model=self.model,
drop_params=False,
drop_params=True,
)
assert result["user"] == "test-user-123"
assert result["size"] == "1024x1024"
assert "quality" not in result
def test_get_complete_url_default(self):
"""Test that get_complete_url returns default ModelScope URL."""
@ -149,7 +144,7 @@ class TestModelScopeImageGenerationTransformation:
mock_get_secret.return_value = None
headers = {}
with pytest.raises(ValueError, match='MODELSCOPE_API_KEY is not set\\. Please set it via') as exc_info:
with pytest.raises(ValueError, match="MODELSCOPE_API_KEY is not set\\. Please set it via") as exc_info:
self.config.validate_environment(
headers=headers,
model=self.model,
@ -178,12 +173,10 @@ class TestModelScopeImageGenerationTransformation:
assert result["prompt"] == prompt
def test_transform_image_generation_request_with_optional_params(self):
"""Test that transform_image_generation_request includes optional params."""
"""Supported OpenAI params (size) are included in the request body."""
prompt = "A beautiful sunset"
optional_params = {
"n": 2,
"size": "1024x1024",
"response_format": "b64_json",
}
result = self.config.transform_image_generation_request(
@ -196,15 +189,13 @@ class TestModelScopeImageGenerationTransformation:
assert result["model"] == self.model
assert result["prompt"] == prompt
assert result["n"] == 2
assert result["size"] == "1024x1024"
assert result["response_format"] == "b64_json"
def test_transform_image_generation_request_ignores_internal_params(self):
"""Test that transform_image_generation_request ignores params starting with _."""
"""Params starting with _ are dropped from the request body."""
prompt = "A beautiful sunset"
optional_params = {
"n": 2,
"size": "1024x1024",
"_internal_param": "should_be_ignored",
}
@ -217,16 +208,103 @@ class TestModelScopeImageGenerationTransformation:
)
assert result["model"] == self.model
assert result["n"] == 2
assert result["size"] == "1024x1024"
assert "_internal_param" not in result
def test_transform_image_generation_request_merges_extra_body(self):
"""extra_body fields (e.g. image_url, negative_prompt) are merged into
the request body as top-level params, alongside supported OpenAI params
like size. ModelScope-specific fields are not OpenAI params, so they
arrive via extra_body and must be flattened here."""
prompt = "add a birthday hat to the dog"
optional_params = {
"size": "1024x1024",
"extra_body": {
"image_url": ["https://example.com/dog.png"],
"negative_prompt": "lowres, blurry",
},
}
result = self.config.transform_image_generation_request(
model=self.model,
prompt=prompt,
optional_params=optional_params,
litellm_params={},
headers={},
)
assert result["model"] == self.model
assert result["prompt"] == prompt
assert result["size"] == "1024x1024"
# extra_body contents are merged in as top-level params; the extra_body
# key itself is not
assert result["image_url"] == ["https://example.com/dog.png"]
assert result["negative_prompt"] == "lowres, blurry"
assert "extra_body" not in result
def test_transform_image_generation_request_excludes_internal_extra_body_keys(self):
"""extra_headers/extra_query are litellm control params, not ModelScope
body fields; they must not leak into the request body even though
get_optional_params_image_gen sweeps them into extra_body."""
prompt = "add a birthday hat to the dog"
optional_params = {
"size": "1024x1024",
"extra_body": {
"image_url": ["https://example.com/dog.png"],
"extra_headers": {"X-Custom-Auth": "secret-value"},
"extra_query": {"ref": "abc"},
},
}
result = self.config.transform_image_generation_request(
model=self.model,
prompt=prompt,
optional_params=optional_params,
litellm_params={},
headers={},
)
assert result["image_url"] == ["https://example.com/dog.png"]
assert "extra_body" not in result
assert "extra_headers" not in result
assert "extra_query" not in result
assert "secret-value" not in str(result)
def test_transform_image_generation_request_rejects_extra_body_model_override(self):
"""A client-controlled extra_body must not override the routed model or
prompt; otherwise a caller authorized for one model could make the
provider run a different, unauthorized one."""
prompt = "a cute baby sea otter"
optional_params = {
"size": "1024x1024",
"extra_body": {
"model": "Qwen/Different-Model",
"prompt": "override prompt",
"negative_prompt": "lowres",
},
}
result = self.config.transform_image_generation_request(
model="modelscope/Qwen/Qwen-Image-2512",
prompt=prompt,
optional_params=optional_params,
litellm_params={},
headers={},
)
assert result["model"] == "modelscope/Qwen/Qwen-Image-2512"
assert result["prompt"] == prompt
assert result["negative_prompt"] == "lowres"
assert "extra_body" not in result
def test_transform_image_generation_response_with_url_images(self):
"""Test that transform_image_generation_response correctly extracts URL images."""
"""Test that transform_image_generation_response extracts output_images URLs."""
response_data = {
"created": 1234567890,
"data": [
{"url": "https://example.com/image1.png"},
{"url": "https://example.com/image2.png"},
"task_status": "SUCCEED",
"task_id": "abc-123",
"output_images": [
"https://example.com/image1.png",
"https://example.com/image2.png",
],
}
@ -252,13 +330,12 @@ class TestModelScopeImageGenerationTransformation:
assert result.data[0].url == "https://example.com/image1.png"
assert result.data[1].url == "https://example.com/image2.png"
def test_transform_image_generation_response_with_b64_json(self):
"""Test that transform_image_generation_response correctly extracts base64 images."""
def test_transform_image_generation_response_failed_status_raises(self):
"""Test that a FAILED task_status raises an error."""
response_data = {
"created": 1234567890,
"data": [
{"b64_json": "iVBORw0KGgoAAAANS"},
],
"task_status": "FAILED",
"task_id": "abc-123",
"output_images": [],
}
mock_response = MagicMock()
@ -268,62 +345,26 @@ class TestModelScopeImageGenerationTransformation:
model_response = ImageResponse(data=[])
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,
)
with pytest.raises(ModelScopeError, match="ModelScope image generation task failed") as exc_info:
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 len(result.data) == 1
assert result.data[0].b64_json == "iVBORw0KGgoAAAANS"
assert result.data[0].url is None
def test_transform_image_generation_response_with_revised_prompt(self):
"""Test that transform_image_generation_response extracts revised_prompt."""
response_data = {
"created": 1234567890,
"data": [
{
"url": "https://example.com/image.png",
"revised_prompt": "A detailed description of a beautiful sunset",
},
],
}
mock_response = MagicMock()
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
model_response = ImageResponse(data=[])
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 len(result.data) == 1
assert (
result.data[0].revised_prompt
== "A detailed description of a beautiful sunset"
)
assert "failed" in str(exc_info.value).lower()
def test_transform_image_generation_response_empty_data(self):
"""Test that transform_image_generation_response handles empty data array."""
"""SUCCEED with no output_images is an error, not a silent empty success."""
response_data = {
"created": 1234567890,
"data": [],
"task_status": "SUCCEED",
"task_id": "abc-123",
"output_images": [],
}
mock_response = MagicMock()
@ -333,23 +374,24 @@ class TestModelScopeImageGenerationTransformation:
model_response = ImageResponse(data=[])
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,
)
with pytest.raises(ModelScopeError) as exc_info:
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 len(result.data) == 0
assert "no output_images" in str(exc_info.value).lower()
def test_transform_image_generation_response_error_handling(self):
"""Test that transform_image_generation_response raises error on API error."""
response_data = {
"error": {
"errors": {
"message": "Invalid prompt provided",
"type": "invalid_request_error",
}
@ -362,7 +404,7 @@ class TestModelScopeImageGenerationTransformation:
model_response = ImageResponse(data=[])
with pytest.raises(Exception, match='litellm\\.BadRequestError: ModelScope error: Invalid prompt') as exc_info:
with pytest.raises(ModelScopeError, match="ModelScope error: Invalid prompt") as exc_info:
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
@ -388,7 +430,7 @@ class TestModelScopeImageGenerationTransformation:
model_response = ImageResponse(data=[])
with pytest.raises(Exception, match='litellm\\.InternalServerError: Error parsing ModelScope') as exc_info:
with pytest.raises(ModelScopeError, match="Error parsing ModelScope") as exc_info:
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
@ -402,50 +444,402 @@ class TestModelScopeImageGenerationTransformation:
assert "Error parsing ModelScope response" in str(exc_info.value)
def test_get_error_class_bad_request(self):
"""Test that get_error_class returns BadRequestError for 400 status."""
from litellm.exceptions import BadRequestError
def test_get_error_class_returns_modelscope_error(self):
"""get_error_class always returns ModelScopeError and preserves status_code."""
for status_code in (400, 401, 404, 429, 500, 502):
error = self.config.get_error_class(
error_message=f"err {status_code}",
status_code=status_code,
headers={"Content-Type": "application/json"},
)
assert isinstance(error, ModelScopeError), status_code
assert error.status_code == status_code
assert f"err {status_code}" in error.message
error = self.config.get_error_class(
error_message="Bad request",
status_code=400,
headers={"Content-Type": "application/json"},
class TestModelScopeImageGenerationHandler:
"""Tests for the async submit + poll flow in ModelScopeImageGeneration."""
def setup_method(self):
from litellm.llms.modelscope.image_generation.handler import (
ModelScopeImageGeneration,
)
assert isinstance(error, BadRequestError)
self.handler = ModelScopeImageGeneration()
self.model = "Qwen/Qwen-Image-2512"
self.prompt = "a cute baby sea otter"
self.logging_obj = MagicMock()
def test_get_error_class_authentication_error(self):
"""Test that get_error_class returns AuthenticationError for 401 status."""
from litellm.exceptions import AuthenticationError
@staticmethod
def _mock_response(json_data, status_code=200):
mock_response = MagicMock()
mock_response.json.return_value = json_data
mock_response.status_code = status_code
mock_response.text = str(json_data)
mock_response.headers = {}
return mock_response
error = self.config.get_error_class(
error_message="Invalid API key",
status_code=401,
headers={"Content-Type": "application/json"},
def test_sync_submit_and_poll_succeed(self):
"""Submit returns task_id, poll returns SUCCEED with output_images."""
from litellm.types.utils import ImageResponse
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = self._mock_response(
{
"task_status": "SUCCEED",
"task_id": "task-123",
"output_images": ["https://example.com/otter.png"],
}
)
assert isinstance(error, AuthenticationError)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
mock_client.get.return_value = poll_resp
def test_get_error_class_internal_server_error(self):
"""Test that get_error_class returns InternalServerError for 500+ status."""
from litellm.exceptions import InternalServerError
error = self.config.get_error_class(
error_message="Internal server error",
status_code=500,
headers={"Content-Type": "application/json"},
result = self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
assert isinstance(error, InternalServerError)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/otter.png"
# Submit POST hits /images/generations, poll GET hits /tasks/{task_id}
assert mock_client.post.call_count == 1
assert "/images/generations" in mock_client.post.call_args.kwargs["url"]
assert mock_client.get.call_count == 1
assert "/tasks/task-123" in mock_client.get.call_args.kwargs["url"]
def test_get_error_class_default(self):
"""Test that get_error_class returns BadRequestError for other status codes."""
from litellm.exceptions import BadRequestError
def test_sync_submit_returns_error_status_raises(self):
"""A 4xx/5xx on the submit call raises immediately."""
submit_resp = self._mock_response({"errors": {"message": "bad model"}}, status_code=400)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
error = self.config.get_error_class(
error_message="Some error",
status_code=404,
headers={"Content-Type": "application/json"},
with pytest.raises(ModelScopeError):
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=MagicMock(),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
mock_client.get.assert_not_called()
def test_sync_poll_failed_status_raises(self):
"""A FAILED task_status during polling raises."""
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = self._mock_response({"task_status": "FAILED", "task_id": "task-123", "output_images": []})
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
mock_client.get.return_value = poll_resp
with pytest.raises(ModelScopeError):
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=MagicMock(),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
def test_sync_submit_missing_task_id_raises(self):
"""A submit response without task_id raises (cannot poll)."""
submit_resp = self._mock_response({"task_status": "SUCCEED"})
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
with pytest.raises(ModelScopeError):
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=MagicMock(),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
def test_submit_sends_async_mode_header(self):
"""The submit POST must include X-ModelScope-Async-Mode: true."""
from litellm.types.utils import ImageResponse
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = self._mock_response(
{
"task_status": "SUCCEED",
"task_id": "task-123",
"output_images": ["https://example.com/x.png"],
}
)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
mock_client.get.return_value = poll_resp
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
assert isinstance(error, BadRequestError)
submit_headers = mock_client.post.call_args.kwargs["headers"]
assert submit_headers.get("X-ModelScope-Async-Mode") == "true"
# Polling must carry the task-type header, not the async-mode header
poll_headers = mock_client.get.call_args.kwargs["headers"]
assert poll_headers.get("X-ModelScope-Task-Type") == "image_generation"
assert "X-ModelScope-Async-Mode" not in poll_headers
def test_sync_poll_running_then_succeed(self):
"""The poll loop must keep polling while RUNNING and stop on SUCCEED."""
submit_resp = self._mock_response({"task_id": "task-123"})
running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []})
succeed_resp = self._mock_response(
{
"task_status": "SUCCEED",
"task_id": "task-123",
"output_images": ["https://example.com/otter.png"],
}
)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
mock_client.get.side_effect = [running_resp, succeed_resp]
with patch( # test-quality-ok: avoid real sleep delay in sync polling test
"litellm.llms.modelscope.image_generation.handler.time.sleep"
):
result = self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/otter.png"
# RUNNING did not terminate the loop; SUCCEED did on the second GET.
assert mock_client.get.call_count == 2
def test_sync_poll_timeout_raises(self):
"""A task that never leaves RUNNING must time out with a 408."""
submit_resp = self._mock_response({"task_id": "task-123"})
running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []})
mock_client = MagicMock(spec=HTTPHandler)
mock_client.get.return_value = running_resp
with pytest.raises(ModelScopeError) as exc_info:
self.handler._poll_for_result_sync(
initial_response=submit_resp,
api_base=None,
headers={"Authorization": "Bearer ms-test"},
sync_client=mock_client,
max_wait=0.05,
interval=0.01,
timeout=30,
)
assert exc_info.value.status_code == 408
assert "timed out" in str(exc_info.value).lower()
def test_sync_submit_200_with_errors_raises(self):
"""A 200 submit response carrying an `errors` key must raise."""
submit_resp = self._mock_response({"errors": {"message": "rate limited"}}, status_code=200)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
with pytest.raises(ModelScopeError) as exc_info:
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
assert "rate limited" in str(exc_info.value)
mock_client.get.assert_not_called()
def test_sync_poll_non_json_response_raises(self):
"""A non-JSON poll response (e.g. gateway HTML) raises ModelScopeError,
not a raw JSONDecodeError that bypasses error normalization."""
import json as _json
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = MagicMock()
poll_resp.status_code = 200
poll_resp.text = "<html>502 Bad Gateway</html>"
poll_resp.headers = {}
poll_resp.json.side_effect = _json.JSONDecodeError("Expecting value", "<html>", 0)
mock_client = MagicMock(spec=HTTPHandler)
mock_client.post.return_value = submit_resp
mock_client.get.return_value = poll_resp
with pytest.raises(ModelScopeError) as exc_info:
self.handler.image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
aimg_generation=False,
)
assert "Error parsing poll response" in str(exc_info.value)
def test_async_submit_and_poll_succeed(self):
"""Async path: submit returns task_id, poll returns SUCCEED."""
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = self._mock_response(
{
"task_status": "SUCCEED",
"task_id": "task-123",
"output_images": ["https://example.com/otter.png"],
}
)
mock_client = MagicMock(spec=AsyncHTTPHandler)
mock_client.post = AsyncMock(return_value=submit_resp)
mock_client.get = AsyncMock(return_value=poll_resp)
result = asyncio.run(
self.handler.async_image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
)
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/otter.png"
assert mock_client.post.call_count == 1
assert mock_client.get.call_count == 1
def test_async_poll_failed_status_raises(self):
"""Async path: a FAILED task_status during polling raises."""
submit_resp = self._mock_response({"task_id": "task-123"})
poll_resp = self._mock_response({"task_status": "FAILED", "task_id": "task-123", "output_images": []})
mock_client = MagicMock(spec=AsyncHTTPHandler)
mock_client.post = AsyncMock(return_value=submit_resp)
mock_client.get = AsyncMock(return_value=poll_resp)
with pytest.raises(ModelScopeError):
asyncio.run(
self.handler.async_image_generation(
model=self.model,
prompt=self.prompt,
model_response=ImageResponse(data=[]),
optional_params={},
litellm_params={"api_key": "ms-test", "api_base": None},
logging_obj=self.logging_obj,
timeout=30,
client=mock_client,
)
)
def test_async_poll_timeout_raises(self):
"""Async path: a task that never leaves RUNNING must time out with 408."""
submit_resp = self._mock_response({"task_id": "task-123"})
running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []})
mock_client = MagicMock(spec=AsyncHTTPHandler)
mock_client.get = AsyncMock(return_value=running_resp)
with pytest.raises(ModelScopeError) as exc_info:
asyncio.run(
self.handler._poll_for_result_async(
initial_response=submit_resp,
api_base=None,
headers={"Authorization": "Bearer ms-test"},
async_client=mock_client,
max_wait=0.05,
interval=0.01,
timeout=30,
)
)
assert exc_info.value.status_code == 408
class TestModelScopeImageGenerationDispatch:
"""Tests for the modelscope dispatch branch in images/main.py."""
def test_api_base_override_without_api_key_raises(self):
"""Security: caller-controlled api_base must not receive deployment key."""
import litellm
with pytest.raises(Exception, match="api_key must be provided when api_base is overridden"):
litellm.image_generation(
prompt="test",
model="modelscope/Qwen/Qwen-Image-2512",
api_base="https://evil.example.com/v1",
api_key=None,
)
def test_dispatch_passes_api_key_and_api_base(self):
"""Deployment credentials are forwarded to the handler."""
import litellm.images.main as img_main
from litellm.types.utils import ImageResponse
captured: dict = {}
original = img_main.modelscope_image_generation.image_generation
def spy(**kwargs):
captured.update(kwargs)
return ImageResponse(data=[])
with patch.object( # test-quality-ok: verify dispatch wiring without real HTTP
img_main.modelscope_image_generation, "image_generation", spy
):
img_main.image_generation(
prompt="test",
model="modelscope/Qwen/Qwen-Image-2512",
api_key="ms-deploy-key",
api_base="https://api-inference.modelscope.cn/v1",
custom_llm_provider="modelscope",
)
assert captured["litellm_params"]["api_key"] == "ms-deploy-key"
assert captured["litellm_params"]["api_base"] == "https://api-inference.modelscope.cn/v1"