diff --git a/litellm/constants.py b/litellm/constants.py index cc6db6c10cc..3b1fe4b6d9f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1173,6 +1173,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", diff --git a/litellm/images/main.py b/litellm/images/main.py index 1688087c2da..3158939ed0b 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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, @@ -422,6 +425,23 @@ def image_generation( client=client, aimg_generation=aimg_generation, ) + elif custom_llm_provider == "modelscope": + # ModelScope image gen is async: submit returns a task_id that + # must be polled until SUCCEED/FAILED. + if model is None: + raise Exception("Model needs to be set for modelscope") + 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, diff --git a/litellm/llms/modelscope/common_utils.py b/litellm/llms/modelscope/common_utils.py new file mode 100644 index 00000000000..b30116dbf8c --- /dev/null +++ b/litellm/llms/modelscope/common_utils.py @@ -0,0 +1,27 @@ +""" +ModelScope Common Utilities + +Shared constants and error handling for ModelScope API integration. +""" + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + + +class ModelScopeError(BaseLLMException): + """Exception class for ModelScope API errors.""" + + pass + + +# Polling configuration for async image generation tasks. +# ModelScope returns a task_id from the submit call and the caller must poll +# GET /v1/tasks/{task_id} until task_status is SUCCEED or FAILED. +DEFAULT_POLLING_INTERVAL = 2.0 # seconds +DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes + +ASYNC_MODE_HEADER = "X-ModelScope-Async-Mode" +TASK_TYPE_HEADER = "X-ModelScope-Task-Type" +IMAGE_GENERATION_TASK_TYPE = "image_generation" + +TASK_STATUS_SUCCEED = "SUCCEED" +TASK_STATUS_FAILED = "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..68383df8f2c --- /dev/null +++ b/litellm/llms/modelscope/image_generation/handler.py @@ -0,0 +1,414 @@ +""" +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 time +from typing import Any, Dict, Optional, Union + +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, + 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. + + Owns the submit + poll HTTP flow; request/response shaping is delegated to + ModelScopeImageGenerationConfig. + """ + + def __init__(self): + self.config = ModelScopeImageGenerationConfig() + + def image_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + optional_params: Dict, + litellm_params: Union[GenericLiteLLMParams, Dict], + logging_obj: LiteLLMLoggingObj, + timeout: Optional[Union[float, httpx.Timeout]], + extra_headers: Optional[Dict[str, Any]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + aimg_generation: bool = False, + ) -> Union[ImageResponse, Any]: + 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, + ) + + if isinstance(litellm_params, dict): + api_key = litellm_params.get("api_key") + api_base = litellm_params.get("api_base") + litellm_params_dict = litellm_params + else: + api_key = litellm_params.api_key + api_base = litellm_params.api_base + litellm_params_dict = dict(litellm_params) + + sync_client = client if isinstance(client, HTTPHandler) else _get_httpx_client() + + headers = 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 = 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 = 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={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response = sync_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise ModelScopeError( + status_code=500, + message=f"Request failed: {str(e)}", + ) + + final_response = self._poll_for_result_sync( + initial_response=response, + 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, + litellm_params: Union[GenericLiteLLMParams, Dict], + logging_obj: LiteLLMLoggingObj, + timeout: Optional[Union[float, httpx.Timeout]], + extra_headers: Optional[Dict[str, Any]] = None, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + if isinstance(litellm_params, dict): + api_key = litellm_params.get("api_key") + api_base = litellm_params.get("api_base") + litellm_params_dict = litellm_params + else: + api_key = litellm_params.api_key + api_base = litellm_params.api_base + litellm_params_dict = dict(litellm_params) + + async_client = client or get_async_httpx_client( + llm_provider=litellm.LlmProviders.MODELSCOPE, + ) + + headers = 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 = 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 = 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={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response = await async_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise ModelScopeError( + status_code=500, + message=f"Request failed: {str(e)}", + ) + + final_response = await self._poll_for_result_async( + initial_response=response, + 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: Optional[str], + headers: dict, + sync_client: HTTPHandler, + max_wait: float = DEFAULT_MAX_POLLING_TIME, + interval: float = DEFAULT_POLLING_INTERVAL, + timeout: Optional[Union[float, httpx.Timeout]] = 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 = initial_response.json() + except Exception 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 = response_data.get("task_id") + if not task_id: + raise ModelScopeError( + status_code=500, + message="No task_id in ModelScope submit response", + ) + + polling_url = self.config.get_task_status_url(api_base, task_id) + polling_headers = self.config.get_polling_headers(headers) + + start_time = time.time() + verbose_logger.debug(f"ModelScope starting sync polling at {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 Exception as e: + raise ModelScopeError( + status_code=response.status_code, + message=f"Error parsing poll response: {e}", + ) + status = data.get("task_status") + verbose_logger.debug(f"ModelScope poll status: {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: Optional[str], + headers: dict, + async_client: AsyncHTTPHandler, + max_wait: float = DEFAULT_MAX_POLLING_TIME, + interval: float = DEFAULT_POLLING_INTERVAL, + timeout: Optional[Union[float, httpx.Timeout]] = 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 = initial_response.json() + except Exception 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 = response_data.get("task_id") + if not task_id: + raise ModelScopeError( + status_code=500, + message="No task_id in ModelScope submit response", + ) + + polling_url = self.config.get_task_status_url(api_base, task_id) + polling_headers = self.config.get_polling_headers(headers) + + start_time = time.time() + verbose_logger.debug(f"ModelScope starting async polling at {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 Exception as e: + raise ModelScopeError( + status_code=response.status_code, + message=f"Error parsing poll response: {e}", + ) + status = data.get("task_status") + verbose_logger.debug(f"ModelScope poll status: {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", + ) + + +# Singleton instance for use in images/main.py +modelscope_image_generation = ModelScopeImageGeneration() diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 3a8a37307d6..5d89ccc5c2b 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -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, @@ -31,29 +39,19 @@ 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, @@ -61,17 +59,16 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): model: str, drop_params: bool, ) -> dict: - """ - Map OpenAI parameters to ModelScope parameters. - - ModelScope uses the same parameter names as OpenAI. - """ 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: 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, @@ -82,14 +79,11 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): litellm_params: dict, 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: + """Build the URL used to poll an async image generation task.""" + return f"{self._get_base_url(api_base)}/tasks/{task_id}" @override def validate_environment( @@ -115,11 +109,21 @@ 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: dict) -> dict: + return { + "Authorization": headers.get("Authorization", ""), + TASK_TYPE_HEADER: IMAGE_GENERATION_TASK_TYPE, + } + + @override def transform_image_generation_request( self, model: str, @@ -128,20 +132,19 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): 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] = { "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 = optional_params.get("extra_body") or {} + 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 @@ -162,10 +165,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): 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 = raw_response.json() @@ -176,64 +179,53 @@ 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 = response_data["errors"] + error_msg = 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 = 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 = response_data.get("output_images", []) or [] + + 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, ) -> BaseLLMException: - """Return the appropriate error class for ModelScope.""" - from litellm.exceptions import ( - AuthenticationError, - BadRequestError, - InternalServerError, - ) - - 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", - ) + """Return the ModelScope error class, preserving the real status code.""" + return ModelScopeError( + status_code=status_code, + message=error_message, + headers=headers, + ) \ No newline at end of file diff --git a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py index 2ffe7c3e686..f22c911230d 100644 --- a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py +++ b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py @@ -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.""" @@ -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,76 @@ 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_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 +303,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 +318,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(Exception) 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 +347,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", } @@ -402,50 +417,356 @@ 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] + + 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): + 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