From f466f109bc64c7f151b7f7d22354ba31268a0886 Mon Sep 17 00:00:00 2001 From: yrk <2493404415@qq.com> Date: Tue, 27 Oct 2026 15:40:12 +0800 Subject: [PATCH] fix(modelscope): use async task polling for image generation ModelScope image generation, added in PR #28460, was non-functional: it parsed a synchronous OpenAI-style data array, but ModelScope's image API is asynchronous. This fixes it to POST /v1/images/generations (returns a task_id), then poll GET /v1/tasks/{task_id} with the X-ModelScope-Task-Type: image_generation header until task_status is SUCCEED (image URLs in output_images) or FAILED, for both sync and async paths. Only size is advertised as a supported OpenAI param; n, response_format, and user are not honored by ModelScope and were removed. Provider-specific fields (negative_prompt, seed, steps, guidance, loras, image_url) pass through via extra_body, which the transform flattens into the body while skipping extra_headers/extra_query to avoid leaking secrets into the request. Errors normalize to ModelScopeError. api_key and api_base are forwarded from the resolved deployment credentials, and a request-supplied api_base without an api_key is rejected so a deployment key is never sent to a caller-controlled host. --- litellm/constants.py | 39 +- litellm/images/main.py | 18 + litellm/llms/modelscope/common_utils.py | 24 + .../modelscope/image_generation/handler.py | 415 ++++++ .../image_generation/transformation.py | 220 ++-- ...est_modelscope_image_gen_transformation.py | 1141 ++++++++++++++--- 6 files changed, 1512 insertions(+), 345 deletions(-) create mode 100644 litellm/llms/modelscope/common_utils.py create mode 100644 litellm/llms/modelscope/image_generation/handler.py diff --git a/litellm/constants.py b/litellm/constants.py index 8cc51d0ab5e..a17b8632f6e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1416,42 +1416,21 @@ WANDB_MODELS: Final[set] = set( modelscope_models: Final[set] = set( [ # Qwen series models - "Qwen/Qwen3-0.6B", - "Qwen/Qwen3-1.7B", - "Qwen/Qwen3-4B", - "Qwen/Qwen3-8B", - "Qwen/Qwen3-14B", - "Qwen/Qwen3-30B-A3B", - "Qwen/Qwen3-32B", - "Qwen/Qwen3-235B-A22B", - "Qwen/Qwen3-235B-A22B-Instruct-2507", - "Qwen/Qwen3-235B-A22B-Thinking-2507", - "Qwen/Qwen3-30B-A3B-Thinking-2507", - "Qwen/Qwen3-Coder-30B-A3B-Instruct", - "Qwen/Qwen3-Coder-480B-A35B-Instruct", - "Qwen/Qwen3-Next-80B-A3B-Instruct", - "Qwen/Qwen3-Next-80B-A3B-Thinking", - "Qwen/Qwen3-VL-235B-A22B-Instruct", - "Qwen/Qwen3-VL-8B-Instruct", - "Qwen/Qwen3-VL-8B-Thinking", "Qwen/Qwen3.5-122B-A10B", "Qwen/Qwen3.5-27B", "Qwen/Qwen3.5-35B-A3B", "Qwen/Qwen3.5-397B-A17B", - "Qwen/QwQ-32B", - "Qwen/QwQ-32B-Preview", - "Qwen/QVQ-72B-Preview", + "Qwen/Qwen3.8-27B", + "Qwen/Qwen3.8-Flash-Next", "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", - "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", - "deepseek-ai/DeepSeek-V3.2", - "deepseek-ai/DeepSeek-V4-Flash", + "deepseek-ai/DeepSeek-V4-Flash-0731", + "deepseek-ai/DeepSeek-V4-Pro", + "deepseek-ai/DeepSeek-V4-Pro-0813", + "deepseek-ai/DeepSeek-V4.1-Flash", + "deepseek-ai/DeepSeek-V4.1-Flash-1017", ] ) diff --git a/litellm/images/main.py b/litellm/images/main.py index 9c88e737a40..aef967c72b8 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -39,6 +39,7 @@ 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 +from litellm.llms.modelscope.image_generation.handler import modelscope_image_generation from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.main import ( azure_chat_completions, @@ -496,6 +497,23 @@ def image_generation( client=client, aimg_generation=aimg_generation, ) + elif custom_llm_provider == "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, # pyright: ignore[reportArgumentType] # get_optional_params_image_gen returns dict[str, object]; values are JSON + litellm_params=litellm_params_dict, # pyright: ignore[reportUnknownArgumentType] # kwargs plumbing is untyped upstream + logging_obj=litellm_logging_obj, + timeout=timeout, + extra_headers=extra_headers, # pyright: ignore[reportUnknownArgumentType] # kwargs plumbing is untyped upstream + client=client, # pyright: ignore[reportUnknownArgumentType] # kwargs plumbing is untyped upstream + aimg_generation=aimg_generation, # pyright: ignore[reportUnknownArgumentType] # kwargs plumbing is untyped upstream + ) elif custom_llm_provider == "azure_ai": from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, diff --git a/litellm/llms/modelscope/common_utils.py b/litellm/llms/modelscope/common_utils.py new file mode 100644 index 00000000000..6906277d280 --- /dev/null +++ b/litellm/llms/modelscope/common_utils.py @@ -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" diff --git a/litellm/llms/modelscope/image_generation/handler.py b/litellm/llms/modelscope/image_generation/handler.py new file mode 100644 index 00000000000..fa48d174d58 --- /dev/null +++ b/litellm/llms/modelscope/image_generation/handler.py @@ -0,0 +1,415 @@ +""" +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 collections.abc import Coroutine, Mapping +from typing import Final + +import httpx +from pydantic import ConfigDict, JsonValue, TypeAdapter + +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_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm client factory is untyped + get_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm client factory is untyped +) +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 + +_JSON_RESPONSE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(hide_input_in_errors=True)) + + +def _str_or_none(value: JsonValue) -> str | None: + return value if isinstance(value, str) else None + + +def _extract_params( + litellm_params: GenericLiteLLMParams | dict[str, JsonValue], # mutable-ok: litellm passes a mutable dict +) -> tuple[str | None, str | None, dict[str, JsonValue]]: # mutable-ok: returns the caller's dict unchanged + if isinstance(litellm_params, dict): + return ( + _str_or_none(litellm_params.get("api_key")), + _str_or_none(litellm_params.get("api_base")), + litellm_params, + ) + return litellm_params.api_key, litellm_params.api_base, dict(litellm_params) + + +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[str, JsonValue], # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict[str, JsonValue], # mutable-ok: litellm override signature + logging_obj: LiteLLMLoggingObj, + timeout: float | httpx.Timeout | None, + extra_headers: Mapping[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + aimg_generation: bool = False, + ) -> ImageResponse | Coroutine[object, object, 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, api_base, litellm_params_dict = _extract_params(litellm_params) + + sync_client: Final = client if isinstance(client, HTTPHandler) else get_httpx_client() # pyright: ignore[reportUnknownVariableType] # litellm client factory is untyped + + headers: Final = self.config.validate_environment( + api_key=api_key, + headers={}, + model=model, + messages=[], + 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( # pyright: ignore[reportUnknownMemberType] # litellm logging API is untyped + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response: Final = sync_client.post( # pyright: ignore[reportUnknownMemberType] # untyped 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[str, JsonValue], # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict[str, JsonValue], # mutable-ok: litellm override signature + logging_obj: LiteLLMLoggingObj, + timeout: float | httpx.Timeout | None, + extra_headers: Mapping[str, str] | None = None, + client: AsyncHTTPHandler | None = None, + ) -> ImageResponse: + api_key, api_base, litellm_params_dict = _extract_params(litellm_params) + + async_client: Final = client or get_async_httpx_client( # pyright: ignore[reportUnknownVariableType] # litellm client factory is untyped + llm_provider=litellm.LlmProviders.MODELSCOPE, + ) + + headers: Final = self.config.validate_environment( + api_key=api_key, + headers={}, + model=model, + messages=[], + 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( # pyright: ignore[reportUnknownMemberType] # litellm logging API is untyped + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response: Final = await async_client.post( # pyright: ignore[reportUnknownMemberType] # untyped 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[str, str], # 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 = _JSON_RESPONSE.validate_python(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[str | None] = _str_or_none(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( # pyright: ignore[reportUnknownMemberType] # untyped 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 = _JSON_RESPONSE.validate_python(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[str, str], # 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 = _JSON_RESPONSE.validate_python(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[str | None] = _str_or_none(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( # pyright: ignore[reportUnknownMemberType] # untyped client get + url=polling_url, + headers=polling_headers, # pyright: ignore[reportArgumentType] # untyped client headers param + timeout=timeout, + ) + + if response.status_code >= 400: + raise ModelScopeError( + status_code=response.status_code, + message=f"Polling failed: {response.text}", + ) + + try: + data = _JSON_RESPONSE.validate_python(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() diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index c308de92d1d..551a1f530ba 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -10,7 +10,7 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter from typing_extensions import override from litellm.litellm_core_utils.url_utils import join_api_base @@ -18,6 +18,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, @@ -35,79 +43,71 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _DECODED_JSON: Final = TypeAdapter(object) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_JSON_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) + + +def _as_json_dict(value: JsonValue) -> dict[str, JsonValue]: # mutable-ok: request body fields + return value if isinstance(value, dict) else {} class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): - """ - Configuration for ModelScope image generation. - - Supports text-to-image models like: - - Qwen/Qwen-Image-Edit - - And other ModelScope-hosted image generation models - """ + """Configuration for ModelScope image generation and editing models.""" DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1" + @override 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 - ] + # Only size is honored (verified via real API: n is silently ignored, + # response_format/user are not honored by ModelScope). + # Provider-specific fields (negative_prompt, seed, steps, guidance, + # loras, image_url) pass through via extra_body, not this list. + return ["size"] + @override def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: dict[str, JsonValue], + optional_params: dict[str, JsonValue], model: str, drop_params: bool, - ) -> dict: - """ - Map OpenAI parameters to ModelScope parameters. - - ModelScope uses the same parameter names as OpenAI. - """ + ) -> dict[str, JsonValue]: 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} 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[str, JsonValue], + litellm_params: dict[str, JsonValue], stream: bool | None = None, ) -> str: - """ - Get the complete URL for the ModelScope image generation API request. - """ - base_url: Final = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(self._get_base_url(api_base), "images/generations") - # Return the images endpoint - return join_api_base(base_url, "images/generations") + def get_task_status_url(self, api_base: str | None, task_id: str) -> str: + """Build the URL used to poll an async image generation task.""" + return join_api_base(self._get_base_url(api_base), f"tasks/{task_id}") @override def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + optional_params: dict[str, JsonValue], + litellm_params: dict[str, JsonValue], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: """ Validate environment and set up headers for ModelScope. """ @@ -121,33 +121,40 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): default_headers: Final = { "Content-Type": "application/json", "Authorization": f"Bearer {final_api_key}", + # ModelScope image-gen only supports async mode: the submit call + # returns a task_id and the caller must poll GET /v1/tasks/{task_id}. + ASYNC_MODE_HEADER: "true", } headers = {**headers, **default_headers} return headers + def get_polling_headers(self, headers: Mapping[str, str]) -> Mapping[str, str]: + return { + "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. + optional_params: dict[str, JsonValue], + litellm_params: dict[str, JsonValue], + headers: dict[str, str], + ) -> dict[str, JsonValue]: + request_data: Final[dict[str, JsonValue]] = {"model": model, "prompt": prompt} # mutable-ok: request body - ModelScope uses the same format as OpenAI for image generation. - """ - # Build the request body (same as OpenAI) - request_data: Final[dict] = { - "model": model, - "prompt": prompt, - } - - # Add optional params - for key, value in optional_params.items(): - if key.startswith("_"): + # litellm wraps provider fields (e.g. image_url) in extra_body; flatten + # them into the body. Skip extra_headers/extra_query (litellm control + # params swept in alongside, may carry secrets) and model/prompt (set + # above from the routed values; a client-controlled extra_body must not + # override the authorized model or prompt). + extra_body_raw: Final = optional_params.get("extra_body") or {} + extra_body: Final[dict[str, JsonValue]] = _as_json_dict(extra_body_raw) + for key, value in {**optional_params, **extra_body}.items(): + if key in ("extra_body", "extra_headers", "extra_query", "model", "prompt") or key.startswith("_"): continue request_data[key] = value @@ -160,21 +167,21 @@ 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[str, JsonValue], + optional_params: dict[str, JsonValue], + litellm_params: dict[str, JsonValue], encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: """ - Transform ModelScope response to OpenAI-compatible ImageResponse. + Transform a completed ModelScope task response into an ImageResponse. - ModelScope returns the same format as OpenAI: - {"created": timestamp, "data": [{"url": "..."}]} + The polled response looks like: + {"task_status": "SUCCEED", "output_images": ["https://..."], ...} """ try: - response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) + response_data: Final = _JSON_RESPONSE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing ModelScope response: {e}", @@ -182,66 +189,59 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): headers=raw_response.headers, ) - # Check for errors in response - response_object: Final = _JSON_OBJECT.validate_python(response_data) - if "error" in response_object: - error: Final = _JSON_OBJECT.validate_python(response_object["error"]) - error_msg: Final = error.get("message", str(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 = _JSON_OBJECTS.validate_python(response_object.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, + ) + + raw_output: Final = response_data.get("output_images") + output_images: Final[tuple[str, ...]] = ( + tuple(url for url in raw_output if isinstance(url, str)) if isinstance(raw_output, list) else () + ) + + if not model_response.data: + model_response.data = [] # rebind-ok: litellm response contract expects the passed-in object populated + + 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[str, str] | httpx.Headers, # mutable-ok: litellm override signature ) -> BaseLLMException: - """Return the appropriate error class for ModelScope.""" - from litellm.exceptions import ( - AuthenticationError, - BadRequestError, - InternalServerError, + """Return the ModelScope error class, preserving the real status code.""" + 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", - ) diff --git a/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py index 5b0ca09ed13..1f05db0b89a 100644 --- a/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py +++ b/tests/unit/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py @@ -5,14 +5,16 @@ 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 +import json +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from pydantic import ValidationError -import litellm +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, ) @@ -27,47 +29,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.""" @@ -152,7 +146,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, @@ -181,12 +175,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( @@ -199,15 +191,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", } @@ -220,16 +210,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", ], } @@ -255,13 +332,13 @@ 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_intermediate_status_raises(self): + """A task_status that is neither SUCCEED nor FAILED (e.g. a poll response + handed straight to the transform) is an error, not a silent empty success.""" response_data = { - "created": 1234567890, - "data": [ - {"b64_json": "iVBORw0KGgoAAAANS"}, - ], + "task_status": "RUNNING", + "task_id": "abc-123", + "output_images": [], } mock_response = MagicMock() @@ -271,31 +348,27 @@ 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) == 1 - assert result.data[0].b64_json == "iVBORw0KGgoAAAANS" - assert result.data[0].url is None + assert "did not succeed" in str(exc_info.value).lower() + assert "RUNNING" in str(exc_info.value) - def test_transform_image_generation_response_with_revised_prompt(self): - """Test that transform_image_generation_response extracts revised_prompt.""" + def test_transform_image_generation_response_failed_status_raises(self): + """Test that a FAILED task_status raises an error.""" response_data = { - "created": 1234567890, - "data": [ - { - "url": "https://example.com/image.png", - "revised_prompt": "A detailed description of a beautiful sunset", - }, - ], + "task_status": "FAILED", + "task_id": "abc-123", + "output_images": [], } mock_response = MagicMock() @@ -305,28 +378,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].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() @@ -336,23 +407,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", } @@ -365,7 +437,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, @@ -391,7 +463,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, @@ -405,130 +477,789 @@ 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) as exc_info: + 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, + ) + + assert exc_info.value.status_code == 400 + assert "submit failed" in str(exc_info.value).lower() + 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) as exc_info: + 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, + ) + + assert exc_info.value.status_code == 400 + assert "task failed" in str(exc_info.value).lower() + + 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 = "502 Bad Gateway" + poll_resp.headers = {} + poll_resp.json.side_effect = _json.JSONDecodeError("Expecting value", "", 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) as exc_info: + 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 exc_info.value.status_code == 400 + assert "task failed" in str(exc_info.value).lower() + + 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 + + def test_aimg_generation_flag_dispatches_to_async(self): + """aimg_generation=True hands the call to async_image_generation and + returns an awaitable, the form litellm's aimage_generation entrypoint uses.""" + 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.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=True, + ) + ) + + assert len(result.data) == 1 + assert result.data[0].url == "https://example.com/otter.png" + mock_client.post.assert_awaited_once() + + def test_litellm_params_as_generic_object(self): + """The router hands the handler a GenericLiteLLMParams, not a dict; the + submit must still reach the provider with its credentials.""" + from litellm.types.router import GenericLiteLLMParams + + 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=HTTPHandler) + mock_client.post.return_value = submit_resp + mock_client.get.return_value = poll_resp + + result = self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=ImageResponse(data=[]), + optional_params={}, + litellm_params=GenericLiteLLMParams(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 + submit_headers = mock_client.post.call_args.kwargs["headers"] + assert submit_headers["Authorization"] == "Bearer ms-test" + + def test_sync_extra_headers_reach_submit_call(self): + """Caller-supplied extra_headers must be merged into the submit headers.""" + 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=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, + extra_headers={"X-Custom-Trace": "trace-42"}, + client=mock_client, + aimg_generation=False, + ) + + submit_headers = mock_client.post.call_args.kwargs["headers"] + assert submit_headers["X-Custom-Trace"] == "trace-42" + assert submit_headers["Authorization"] == "Bearer ms-test" + + def test_async_extra_headers_reach_submit_call(self): + """Async path: caller-supplied extra_headers must be merged into the submit headers.""" + 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) + + 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, + extra_headers={"X-Custom-Trace": "trace-42"}, + client=mock_client, + ) + ) + + submit_headers = mock_client.post.call_args.kwargs["headers"] + assert submit_headers["X-Custom-Trace"] == "trace-42" + assert submit_headers["Authorization"] == "Bearer ms-test" + + def test_sync_submit_network_error_raises_modelscope_error(self): + """A transport failure on the submit call surfaces as ModelScopeError(500), not a raw httpx error.""" + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.side_effect = httpx.ConnectError("connection refused") + + 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 exc_info.value.status_code == 500 + assert "Request failed" in str(exc_info.value) + + def test_async_submit_network_error_raises_modelscope_error(self): + """Async path: a transport failure on submit surfaces as ModelScopeError(500).""" + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + + with pytest.raises(ModelScopeError) as exc_info: + 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 exc_info.value.status_code == 500 + assert "Request failed" in str(exc_info.value) + + def test_sync_submit_non_json_response_raises(self): + """A 200 submit response that is not JSON (e.g. gateway HTML) raises ModelScopeError.""" + submit_resp = MagicMock() + submit_resp.status_code = 200 + submit_resp.text = "502 Bad Gateway" + submit_resp.headers = {} + submit_resp.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + + 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 "Error parsing submit response" in str(exc_info.value) + mock_client.get.assert_not_called() + + def test_sync_poll_error_status_raises(self): + """A 4xx/5xx poll response raises with the poll status code.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = self._mock_response({"message": "boom"}, status_code=502) + + 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 exc_info.value.status_code == 502 + assert "Polling failed" in str(exc_info.value) + + def test_async_submit_error_status_raises(self): + """Async path: a 4xx submit response raises with the submit status code.""" + submit_resp = self._mock_response({"errors": {"message": "bad model"}}, status_code=400) + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=submit_resp) + + with pytest.raises(ModelScopeError) as exc_info: + 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 exc_info.value.status_code == 400 + mock_client.get.assert_not_awaited() + + def test_async_submit_non_json_response_raises(self): + """Async path: a non-JSON submit response raises ModelScopeError.""" + submit_resp = MagicMock() + submit_resp.status_code = 200 + submit_resp.text = "502 Bad Gateway" + submit_resp.headers = {} + submit_resp.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=submit_resp) + + with pytest.raises(ModelScopeError) as exc_info: + 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 "Error parsing submit response" in str(exc_info.value) + + def test_async_submit_200_with_errors_raises(self): + """Async path: a 200 submit response carrying an `errors` key raises.""" + submit_resp = self._mock_response({"errors": {"message": "quota exceeded"}}, status_code=200) + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=submit_resp) + + with pytest.raises(ModelScopeError) as exc_info: + 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 "quota exceeded" in str(exc_info.value) + mock_client.get.assert_not_awaited() + + def test_async_submit_missing_task_id_raises(self): + """Async path: a submit response without task_id raises (cannot poll).""" + submit_resp = self._mock_response({"task_status": "SUCCEED"}) + mock_client = MagicMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=submit_resp) + + with pytest.raises(ModelScopeError) as exc_info: + 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 "No task_id" in str(exc_info.value) + + def test_async_poll_error_status_raises(self): + """Async path: a 4xx/5xx poll response raises with the poll status code.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = self._mock_response({"message": "boom"}, status_code=500) + + 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) as exc_info: + 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 exc_info.value.status_code == 500 + assert "Polling failed" in str(exc_info.value) + + def test_async_poll_non_json_response_raises(self): + """Async path: a non-JSON poll response raises ModelScopeError, not a raw JSONDecodeError.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = MagicMock() + poll_resp.status_code = 200 + poll_resp.text = "502 Bad Gateway" + poll_resp.headers = {} + poll_resp.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + + 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) as exc_info: + 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 "Error parsing poll response" in str(exc_info.value) -def _transform_generation_response(payload: object, status_code: int = 200) -> ImageResponse: - return ModelScopeImageGenerationConfig().transform_image_generation_response( - model="Qwen/Qwen-Image", - raw_response=httpx.Response(status_code, json=payload), - model_response=ImageResponse(data=[]), - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) +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 -@pytest.mark.parametrize( - ("item", "expected"), - [ - ({"url": "https://a.example/i.png"}, ("https://a.example/i.png", None, None)), - ({"b64_json": "aGVsbG8="}, (None, "aGVsbG8=", None)), - ( - {"url": "https://a.example/i.png", "revised_prompt": "a calmer cat", "seed": 7}, - ("https://a.example/i.png", None, "a calmer cat"), - ), - ({}, (None, None, None)), - ], -) -def test_transform_image_generation_response_maps_one_image( - item: dict[str, object], expected: tuple[str | None, str | None, str | None] -) -> None: - response = _transform_generation_response({"created": 1, "data": [item]}) + 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, + ) - assert [(image.url, image.b64_json, image.revised_prompt) for image in response.data] == [expected] + 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 = {} -@pytest.mark.parametrize("payload", [{}, {"data": []}, {"data": ""}, {"data": {}}, {"created": 1}]) -def test_transform_image_generation_response_without_images_is_empty(payload: dict[str, object]) -> None: - assert _transform_generation_response(payload).data == [] + original = img_main.modelscope_image_generation.image_generation + def spy(**kwargs): + captured.update(kwargs) + return ImageResponse(data=[]) -@pytest.mark.parametrize( - ("error", "status_code", "expected_class", "expected_message"), - [ - ({"message": "Invalid prompt"}, 400, litellm.BadRequestError, "ModelScope error: Invalid prompt"), - ({"message": "Bad key"}, 401, litellm.AuthenticationError, "ModelScope error: Bad key"), - ({"code": "overloaded"}, 503, litellm.InternalServerError, "ModelScope error: {'code': 'overloaded'}"), - ({}, 200, litellm.BadRequestError, "ModelScope error: {}"), - ], -) -def test_transform_image_generation_response_reports_api_error_bodies( - error: dict[str, object], status_code: int, expected_class: type[Exception], expected_message: str -) -> None: - with pytest.raises(expected_class) as exc_info: - _transform_generation_response({"error": error, "data": [{"url": "ignored"}]}, status_code) + with patch.object( # test-quality-ok: verify dispatch wiring without real HTTP + img_main.modelscope_image_generation, "image_generation", spy + ): + img_main.image_generation( # pyright: ignore[reportUnknownMemberType] # @client decorator is untyped + 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 str(exc_info.value).endswith(expected_message) - - -@pytest.mark.parametrize( - "payload", - [ - ["not", "an", "object"], - "error", - 7, - {"error": "plain text error"}, - {"error": None}, - {"error": ["not", "an", "object"]}, - {"data": 7}, - {"data": None}, - {"data": ["not an object"]}, - {"data": [{"url": "https://a.example/i.png"}, None]}, - ], -) -def test_transform_image_generation_response_rejects_malformed_payloads_without_echoing_them( - payload: object, -) -> None: - with pytest.raises(ValidationError) as exc_info: - _transform_generation_response(payload) - - assert "input_value" not in str(exc_info.value) + assert captured["litellm_params"]["api_key"] == "ms-deploy-key" + assert captured["litellm_params"]["api_base"] == "https://api-inference.modelscope.cn/v1"