diff --git a/litellm/constants.py b/litellm/constants.py index 5f8fa203b37..d00a200bfd1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1273,6 +1273,8 @@ modelscope_models: Final[set] = set( "Qwen/QwQ-32B-Preview", "Qwen/QVQ-72B-Preview", "Qwen/Qwen-Image-Edit", + "Qwen/Qwen-Image-Edit-2511", + "Qwen/Qwen-Image-2512", # DeepSeek series models "deepseek-ai/DeepSeek-R1-0528", "deepseek-ai/DeepSeek-R1-Distill-Llama-70B", diff --git a/litellm/images/main.py b/litellm/images/main.py index 6a94e7c8df2..1154b165416 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, @@ -424,6 +427,25 @@ def image_generation( client=client, aimg_generation=aimg_generation, ) + elif custom_llm_provider == "modelscope": + if model is None: + raise Exception("Model needs to be set for modelscope") + if api_base is not None and api_key is None: + raise ValueError("api_key must be provided when api_base is overridden for modelscope") + litellm_params_dict["api_key"] = api_key + litellm_params_dict["api_base"] = api_base + return modelscope_image_generation.image_generation( + model=model, + prompt=prompt, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params_dict, + logging_obj=litellm_logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client, + aimg_generation=aimg_generation, + ) elif custom_llm_provider == "azure_ai": from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, 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..b04a4b32c37 --- /dev/null +++ b/litellm/llms/modelscope/image_generation/handler.py @@ -0,0 +1,411 @@ +""" +ModelScope Image Generation Handler + +ModelScope image generation only supports async mode: the submit call returns +a task_id, then the caller polls GET /v1/tasks/{task_id} until task_status is +SUCCEED (image URLs in output_images) or FAILED. + +API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro +""" + +import asyncio +import json +import time +from typing import Final + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage] # litellm internal client factory, same pattern as BFL handler + get_async_httpx_client, +) +from litellm.llms.modelscope.common_utils import ( + DEFAULT_MAX_POLLING_TIME, + DEFAULT_POLLING_INTERVAL, + TASK_STATUS_FAILED, + TASK_STATUS_SUCCEED, + ModelScopeError, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse + +from .transformation import ModelScopeImageGenerationConfig + + +class ModelScopeImageGeneration: + """ModelScope image generation handler.""" + + def __init__(self) -> None: + self.config: Final = ModelScopeImageGenerationConfig() + + def image_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature + logging_obj: LiteLLMLoggingObj, + timeout: float | httpx.Timeout | None, + extra_headers: dict[str, str] | None = None, # mutable-ok: litellm override signature + client: HTTPHandler | AsyncHTTPHandler | None = None, + aimg_generation: bool = False, + ) -> ImageResponse: + if aimg_generation: + return self.async_image_generation( + model=model, + prompt=prompt, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client if isinstance(client, AsyncHTTPHandler) else None, + ) + + api_key: Final = litellm_params.get("api_key") if isinstance(litellm_params, dict) else litellm_params.api_key + api_base: Final = ( + litellm_params.get("api_base") if isinstance(litellm_params, dict) else litellm_params.api_base + ) + litellm_params_dict: Final = ( + litellm_params + if isinstance(litellm_params, dict) + else dict(litellm_params) # mutable-ok: litellm passes a mutable dict + ) + + sync_client: Final = client if isinstance(client, HTTPHandler) else _get_httpx_client() + + headers: Final = self.config.validate_environment( + api_key=api_key, + headers={}, # mutable-ok: validate_environment returns a new dict + model=model, + messages=[], # mutable-ok: required by base class signature + optional_params=optional_params, + litellm_params=litellm_params_dict, + ) + if extra_headers: + headers.update(extra_headers) + + complete_url: Final = self.config.get_complete_url( + api_base=api_base, + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params_dict, + ) + + data: Final = self.config.transform_image_generation_request( + model=model, + prompt=prompt, + optional_params=optional_params, + litellm_params=litellm_params_dict, + headers=headers, + ) + + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ # mutable-ok: litellm logging contract + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response: Final = sync_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except httpx.HTTPError as e: + raise ModelScopeError( + status_code=500, + message=f"Request failed: {e}", + ) + + final_response: Final = self._poll_for_result_sync( + initial_response=response, # pyright: ignore[reportArgumentType] # post() returns Response | None; None raises HTTPError above + api_base=api_base, + headers=headers, + sync_client=sync_client, + timeout=timeout, + ) + + return self.config.transform_image_generation_response( + model=model, + raw_response=final_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=data, + optional_params=optional_params, + litellm_params=litellm_params_dict, + encoding=None, + ) + + async def async_image_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature + logging_obj: LiteLLMLoggingObj, + timeout: float | httpx.Timeout | None, + extra_headers: dict[str, str] | None = None, # mutable-ok: litellm override signature + client: AsyncHTTPHandler | None = None, + ) -> ImageResponse: + api_key: Final = litellm_params.get("api_key") if isinstance(litellm_params, dict) else litellm_params.api_key + api_base: Final = ( + litellm_params.get("api_base") if isinstance(litellm_params, dict) else litellm_params.api_base + ) + litellm_params_dict: Final = ( + litellm_params + if isinstance(litellm_params, dict) + else dict(litellm_params) # mutable-ok: litellm passes a mutable dict + ) + + async_client: Final = client or get_async_httpx_client( + llm_provider=litellm.LlmProviders.MODELSCOPE, + ) + + headers: Final = self.config.validate_environment( + api_key=api_key, + headers={}, # mutable-ok: validate_environment returns a new dict + model=model, + messages=[], # mutable-ok: required by base class signature + optional_params=optional_params, + litellm_params=litellm_params_dict, + ) + if extra_headers: + headers.update(extra_headers) + + complete_url: Final = self.config.get_complete_url( + api_base=api_base, + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params_dict, + ) + + data: Final = self.config.transform_image_generation_request( + model=model, + prompt=prompt, + optional_params=optional_params, + litellm_params=litellm_params_dict, + headers=headers, + ) + + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ # mutable-ok: litellm logging contract + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + try: + response: Final = await async_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except httpx.HTTPError as e: + raise ModelScopeError( + status_code=500, + message=f"Request failed: {e}", + ) + + final_response: Final = await self._poll_for_result_async( + initial_response=response, # pyright: ignore[reportArgumentType] # post() returns Response | None; None raises HTTPError above + api_base=api_base, + headers=headers, + async_client=async_client, + timeout=timeout, + ) + + return self.config.transform_image_generation_response( + model=model, + raw_response=final_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=data, + optional_params=optional_params, + litellm_params=litellm_params_dict, + encoding=None, + ) + + def _poll_for_result_sync( + self, + initial_response: httpx.Response, + api_base: str | None, + headers: dict, # mutable-ok: caller passes a mutable dict + sync_client: HTTPHandler, + max_wait: float = DEFAULT_MAX_POLLING_TIME, + interval: float = DEFAULT_POLLING_INTERVAL, + timeout: float | httpx.Timeout | None = None, + ) -> httpx.Response: + if initial_response.status_code >= 400: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"ModelScope submit failed: {initial_response.text}", + ) + + try: + response_data: Final = initial_response.json() + except json.JSONDecodeError as e: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"Error parsing submit response: {e}", + ) + + if "errors" in response_data: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"ModelScope error: {response_data['errors']}", + ) + + task_id: Final = response_data.get("task_id") + if not task_id: + raise ModelScopeError( + status_code=500, + message="No task_id in ModelScope submit response", + ) + + polling_url: Final = self.config.get_task_status_url(api_base, task_id) + polling_headers: Final = self.config.get_polling_headers(headers) + + start_time: Final = time.time() + verbose_logger.debug("ModelScope starting sync polling at %s", polling_url) + + while time.time() - start_time < max_wait: + response = sync_client.get( + url=polling_url, + headers=polling_headers, + timeout=timeout, + ) + + if response.status_code >= 400: + raise ModelScopeError( + status_code=response.status_code, + message=f"Polling failed: {response.text}", + ) + + try: + data = response.json() + except json.JSONDecodeError as e: + raise ModelScopeError( + status_code=response.status_code, + message=f"Error parsing poll response: {e}", + ) + status = data.get("task_status") + verbose_logger.debug("ModelScope poll status: %s", status) + + if status == TASK_STATUS_SUCCEED: + return response + elif status == TASK_STATUS_FAILED: + raise ModelScopeError( + status_code=400, + message="ModelScope image generation task failed", + ) + + time.sleep(interval) + + raise ModelScopeError( + status_code=408, + message=f"Polling timed out after {max_wait} seconds", + ) + + async def _poll_for_result_async( + self, + initial_response: httpx.Response, + api_base: str | None, + headers: dict, # mutable-ok: caller passes a mutable dict + async_client: AsyncHTTPHandler, + max_wait: float = DEFAULT_MAX_POLLING_TIME, + interval: float = DEFAULT_POLLING_INTERVAL, + timeout: float | httpx.Timeout | None = None, + ) -> httpx.Response: + if initial_response.status_code >= 400: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"ModelScope submit failed: {initial_response.text}", + ) + + try: + response_data: Final = initial_response.json() + except json.JSONDecodeError as e: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"Error parsing submit response: {e}", + ) + + if "errors" in response_data: + raise ModelScopeError( + status_code=initial_response.status_code, + message=f"ModelScope error: {response_data['errors']}", + ) + + task_id: Final = response_data.get("task_id") + if not task_id: + raise ModelScopeError( + status_code=500, + message="No task_id in ModelScope submit response", + ) + + polling_url: Final = self.config.get_task_status_url(api_base, task_id) + polling_headers: Final = self.config.get_polling_headers(headers) + + start_time: Final = time.time() + verbose_logger.debug("ModelScope starting async polling at %s", polling_url) + + while time.time() - start_time < max_wait: + response = await async_client.get( + url=polling_url, + headers=polling_headers, + timeout=timeout, + ) + + if response.status_code >= 400: + raise ModelScopeError( + status_code=response.status_code, + message=f"Polling failed: {response.text}", + ) + + try: + data = response.json() + except json.JSONDecodeError as e: + raise ModelScopeError( + status_code=response.status_code, + message=f"Error parsing poll response: {e}", + ) + status = data.get("task_status") + verbose_logger.debug("ModelScope poll status: %s", status) + + if status == TASK_STATUS_SUCCEED: + return response + elif status == TASK_STATUS_FAILED: + raise ModelScopeError( + status_code=400, + message="ModelScope image generation task failed", + ) + + await asyncio.sleep(interval) + + raise ModelScopeError( + status_code=408, + message=f"Polling timed out after {max_wait} seconds", + ) + + +modelscope_image_generation: Final = ModelScopeImageGeneration() diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 3a8a37307d6..570ce692dee 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, @@ -25,86 +33,68 @@ from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - LiteLLMLoggingObj = _LiteLLMLoggingObj + LiteLLMLoggingObj: Final = _LiteLLMLoggingObj else: - LiteLLMLoggingObj = object + LiteLLMLoggingObj: Final = object class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): - """ - Configuration for ModelScope image generation. + """Configuration for ModelScope image generation and editing models.""" - Supports text-to-image models like: - - Qwen/Qwen-Image-Edit - - And other ModelScope-hosted image generation models - """ + DEFAULT_BASE_URL: Final[str] = "https://api-inference.modelscope.cn/v1" - DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1" - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - """ - Return list of OpenAI params supported by ModelScope. - - ModelScope supports standard OpenAI image generation parameters. - """ - return [ - "n", # Number of images to generate - "size", # Size of the generated images - "response_format", # url or b64_json - "user", # User identifier - ] + @override + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: litellm override signature + return ["size"] # mutable-ok: small fixed list + @override def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: dict, # mutable-ok: litellm override signature + optional_params: dict, # mutable-ok: litellm override signature model: str, drop_params: bool, - ) -> dict: - """ - Map OpenAI parameters to ModelScope parameters. - - ModelScope uses the same parameter names as OpenAI. - """ + ) -> dict: # mutable-ok: litellm override signature supported_params: Final = self.get_supported_openai_params(model) if drop_params: - non_default_params = {k: v for k, v in non_default_params.items() if k in supported_params} + non_default_params = { + k: v for k, v in non_default_params.items() if k in supported_params + } # rebind-ok: filtering in place is the litellm pattern # mutable-ok: dict comprehension for filtering optional_params.update(non_default_params) return optional_params + def _get_base_url(self, api_base: str | None) -> str: + base_url: Final[str] = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL + return base_url.rstrip("/") + @override def get_complete_url( self, api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: dict, # mutable-ok: litellm override signature stream: bool | None = None, ) -> str: - """ - Get the complete URL for the ModelScope image generation API request. - """ - base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") + return f"{self._get_base_url(api_base)}/images/generations" - # Return the images endpoint - return f"{base_url}/images/generations" + def get_task_status_url(self, api_base: str | None, task_id: str) -> str: + return f"{self._get_base_url(api_base)}/tasks/{task_id}" @override def validate_environment( self, - headers: dict, + headers: dict, # mutable-ok: litellm override signature model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: list[AllMessageValues], # mutable-ok: litellm override signature + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: dict, # mutable-ok: litellm override signature api_key: str | None = None, api_base: str | None = None, - ) -> dict: - """ - Validate environment and set up headers for ModelScope. - """ + ) -> dict: # mutable-ok: litellm override signature final_api_key: Final[str | None] = api_key or get_secret_str("MODELSCOPE_API_KEY") if not final_api_key: @@ -112,36 +102,40 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): "MODELSCOPE_API_KEY is not set. Please set it via environment variable or pass api_key parameter." ) - default_headers: Final = { + default_headers: Final = { # mutable-ok: http header dict "Content-Type": "application/json", "Authorization": f"Bearer {final_api_key}", + ASYNC_MODE_HEADER: "true", } - headers = {**headers, **default_headers} - return headers + return {**headers, **default_headers} # mutable-ok: merged header dict + def get_polling_headers( + self, + headers: dict, # mutable-ok: caller passes a mutable dict + ) -> dict: # mutable-ok: builds a new dict for polling + return { # mutable-ok: polling header dict + "Authorization": headers.get("Authorization", ""), + TASK_TYPE_HEADER: IMAGE_GENERATION_TASK_TYPE, + } + + @override def transform_image_generation_request( self, model: str, prompt: str, - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - """ - Transform OpenAI-style request to ModelScope request format. - - ModelScope uses the same format as OpenAI for image generation. - """ - # Build the request body (same as OpenAI) - request_data: Final[dict] = { + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: dict, # mutable-ok: litellm override signature + headers: dict, # mutable-ok: litellm override signature + ) -> dict: # mutable-ok: litellm override signature + request_data: Final = { # mutable-ok: request body dict "model": model, "prompt": prompt, } - # Add optional params - for key, value in optional_params.items(): - if key.startswith("_"): + extra_body: Final = optional_params.get("extra_body") or {} # mutable-ok: provider fields dict + for key, value in {**optional_params, **extra_body}.items(): # mutable-ok: merged params dict + if key in ("extra_body", "extra_headers", "extra_query", "model", "prompt") or key.startswith("_"): continue request_data[key] = value @@ -154,19 +148,13 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): raw_response: httpx.Response, model_response: ImageResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - optional_params: dict, - litellm_params: dict, + request_data: dict, # mutable-ok: litellm override signature + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: dict, # mutable-ok: litellm override signature encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - """ - Transform ModelScope response to OpenAI-compatible ImageResponse. - - ModelScope returns the same format as OpenAI: - {"created": timestamp, "data": [{"url": "..."}]} - """ try: response_data: Final = raw_response.json() except Exception as e: @@ -176,64 +164,52 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): headers=raw_response.headers, ) - # Check for errors in response - if "error" in response_data: - error_msg: Final = response_data["error"].get("message", str(response_data["error"])) + if "errors" in response_data: + errors: Final = response_data["errors"] + error_msg: Final = errors.get("message", str(errors)) if isinstance(errors, dict) else str(errors) raise self.get_error_class( error_message=f"ModelScope error: {error_msg}", status_code=raw_response.status_code, headers=raw_response.headers, ) - # Extract images from response - data_list: Final = response_data.get("data", []) - if not model_response.data: - model_response.data = [] - - for item in data_list: - image_obj = ImageObject( - url=item.get("url"), - b64_json=item.get("b64_json"), - revised_prompt=item.get("revised_prompt"), + task_status: Final = response_data.get("task_status") + if task_status == TASK_STATUS_FAILED: + raise self.get_error_class( + error_message="ModelScope image generation task failed", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + if task_status != TASK_STATUS_SUCCEED: + raise self.get_error_class( + error_message=f"ModelScope task did not succeed: status={task_status}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + output_images: Final = response_data.get("output_images", []) or [] # mutable-ok: API response list + + for image_url in output_images: + model_response.data.append(ImageObject(url=image_url)) + + if not model_response.data: + raise self.get_error_class( + error_message="ModelScope task SUCCEED but no output_images", + status_code=raw_response.status_code, + headers=raw_response.headers, ) - model_response.data.append(image_obj) return model_response + @override def get_error_class( self, error_message: str, status_code: int, - headers: dict | httpx.Headers, + headers: dict | httpx.Headers, # mutable-ok: litellm override signature ) -> BaseLLMException: - """Return the appropriate error class for ModelScope.""" - from litellm.exceptions import ( - AuthenticationError, - BadRequestError, - InternalServerError, + return ModelScopeError( + status_code=status_code, + message=error_message, + headers=headers, ) - - if status_code == 400: - return BadRequestError( - message=error_message, - model="", - llm_provider="modelscope", - ) - elif status_code == 401: - return AuthenticationError( - message=error_message, - model="", - llm_provider="modelscope", - ) - elif status_code >= 500: - return InternalServerError( - message=error_message, - model="", - llm_provider="modelscope", - ) - else: - return BadRequestError( - message=error_message, - model="", - llm_provider="modelscope", - ) 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..29246fb150c 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.""" @@ -149,7 +144,7 @@ class TestModelScopeImageGenerationTransformation: mock_get_secret.return_value = None headers = {} - with pytest.raises(ValueError, match='MODELSCOPE_API_KEY is not set\\. Please set it via') as exc_info: + with pytest.raises(ValueError, match="MODELSCOPE_API_KEY is not set\\. Please set it via") as exc_info: self.config.validate_environment( headers=headers, model=self.model, @@ -178,12 +173,10 @@ class TestModelScopeImageGenerationTransformation: assert result["prompt"] == prompt def test_transform_image_generation_request_with_optional_params(self): - """Test that transform_image_generation_request includes optional params.""" + """Supported OpenAI params (size) are included in the request body.""" prompt = "A beautiful sunset" optional_params = { - "n": 2, "size": "1024x1024", - "response_format": "b64_json", } result = self.config.transform_image_generation_request( @@ -196,15 +189,13 @@ class TestModelScopeImageGenerationTransformation: assert result["model"] == self.model assert result["prompt"] == prompt - assert result["n"] == 2 assert result["size"] == "1024x1024" - assert result["response_format"] == "b64_json" def test_transform_image_generation_request_ignores_internal_params(self): - """Test that transform_image_generation_request ignores params starting with _.""" + """Params starting with _ are dropped from the request body.""" prompt = "A beautiful sunset" optional_params = { - "n": 2, + "size": "1024x1024", "_internal_param": "should_be_ignored", } @@ -217,16 +208,103 @@ class TestModelScopeImageGenerationTransformation: ) assert result["model"] == self.model - assert result["n"] == 2 + assert result["size"] == "1024x1024" assert "_internal_param" not in result + def test_transform_image_generation_request_merges_extra_body(self): + """extra_body fields (e.g. image_url, negative_prompt) are merged into + the request body as top-level params, alongside supported OpenAI params + like size. ModelScope-specific fields are not OpenAI params, so they + arrive via extra_body and must be flattened here.""" + prompt = "add a birthday hat to the dog" + optional_params = { + "size": "1024x1024", + "extra_body": { + "image_url": ["https://example.com/dog.png"], + "negative_prompt": "lowres, blurry", + }, + } + + result = self.config.transform_image_generation_request( + model=self.model, + prompt=prompt, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["model"] == self.model + assert result["prompt"] == prompt + assert result["size"] == "1024x1024" + # extra_body contents are merged in as top-level params; the extra_body + # key itself is not + assert result["image_url"] == ["https://example.com/dog.png"] + assert result["negative_prompt"] == "lowres, blurry" + assert "extra_body" not in result + + def test_transform_image_generation_request_excludes_internal_extra_body_keys(self): + """extra_headers/extra_query are litellm control params, not ModelScope + body fields; they must not leak into the request body even though + get_optional_params_image_gen sweeps them into extra_body.""" + prompt = "add a birthday hat to the dog" + optional_params = { + "size": "1024x1024", + "extra_body": { + "image_url": ["https://example.com/dog.png"], + "extra_headers": {"X-Custom-Auth": "secret-value"}, + "extra_query": {"ref": "abc"}, + }, + } + + result = self.config.transform_image_generation_request( + model=self.model, + prompt=prompt, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["image_url"] == ["https://example.com/dog.png"] + assert "extra_body" not in result + assert "extra_headers" not in result + assert "extra_query" not in result + assert "secret-value" not in str(result) + + def test_transform_image_generation_request_rejects_extra_body_model_override(self): + """A client-controlled extra_body must not override the routed model or + prompt; otherwise a caller authorized for one model could make the + provider run a different, unauthorized one.""" + prompt = "a cute baby sea otter" + optional_params = { + "size": "1024x1024", + "extra_body": { + "model": "Qwen/Different-Model", + "prompt": "override prompt", + "negative_prompt": "lowres", + }, + } + + result = self.config.transform_image_generation_request( + model="modelscope/Qwen/Qwen-Image-2512", + prompt=prompt, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "modelscope/Qwen/Qwen-Image-2512" + assert result["prompt"] == prompt + assert result["negative_prompt"] == "lowres" + assert "extra_body" not in result + def test_transform_image_generation_response_with_url_images(self): - """Test that transform_image_generation_response correctly extracts URL images.""" + """Test that transform_image_generation_response extracts output_images URLs.""" response_data = { - "created": 1234567890, - "data": [ - {"url": "https://example.com/image1.png"}, - {"url": "https://example.com/image2.png"}, + "task_status": "SUCCEED", + "task_id": "abc-123", + "output_images": [ + "https://example.com/image1.png", + "https://example.com/image2.png", ], } @@ -252,13 +330,12 @@ class TestModelScopeImageGenerationTransformation: assert result.data[0].url == "https://example.com/image1.png" assert result.data[1].url == "https://example.com/image2.png" - def test_transform_image_generation_response_with_b64_json(self): - """Test that transform_image_generation_response correctly extracts base64 images.""" + def test_transform_image_generation_response_failed_status_raises(self): + """Test that a FAILED task_status raises an error.""" response_data = { - "created": 1234567890, - "data": [ - {"b64_json": "iVBORw0KGgoAAAANS"}, - ], + "task_status": "FAILED", + "task_id": "abc-123", + "output_images": [], } mock_response = MagicMock() @@ -268,62 +345,26 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) + with pytest.raises(ModelScopeError, match="ModelScope image generation task failed") as exc_info: + self.config.transform_image_generation_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=self.logging_obj, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) - assert len(result.data) == 1 - assert result.data[0].b64_json == "iVBORw0KGgoAAAANS" - assert result.data[0].url is None - - def test_transform_image_generation_response_with_revised_prompt(self): - """Test that transform_image_generation_response extracts revised_prompt.""" - response_data = { - "created": 1234567890, - "data": [ - { - "url": "https://example.com/image.png", - "revised_prompt": "A detailed description of a beautiful sunset", - }, - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert ( - result.data[0].revised_prompt - == "A detailed description of a beautiful sunset" - ) + assert "failed" in str(exc_info.value).lower() def test_transform_image_generation_response_empty_data(self): - """Test that transform_image_generation_response handles empty data array.""" + """SUCCEED with no output_images is an error, not a silent empty success.""" response_data = { - "created": 1234567890, - "data": [], + "task_status": "SUCCEED", + "task_id": "abc-123", + "output_images": [], } mock_response = MagicMock() @@ -333,23 +374,24 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) + with pytest.raises(ModelScopeError) as exc_info: + self.config.transform_image_generation_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=self.logging_obj, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) - assert len(result.data) == 0 + assert "no output_images" in str(exc_info.value).lower() def test_transform_image_generation_response_error_handling(self): """Test that transform_image_generation_response raises error on API error.""" response_data = { - "error": { + "errors": { "message": "Invalid prompt provided", "type": "invalid_request_error", } @@ -362,7 +404,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception, match='litellm\\.BadRequestError: ModelScope error: Invalid prompt') as exc_info: + with pytest.raises(ModelScopeError, match="ModelScope error: Invalid prompt") as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, @@ -388,7 +430,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception, match='litellm\\.InternalServerError: Error parsing ModelScope') as exc_info: + with pytest.raises(ModelScopeError, match="Error parsing ModelScope") as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, @@ -402,50 +444,402 @@ class TestModelScopeImageGenerationTransformation: assert "Error parsing ModelScope response" in str(exc_info.value) - def test_get_error_class_bad_request(self): - """Test that get_error_class returns BadRequestError for 400 status.""" - from litellm.exceptions import BadRequestError + def test_get_error_class_returns_modelscope_error(self): + """get_error_class always returns ModelScopeError and preserves status_code.""" + for status_code in (400, 401, 404, 429, 500, 502): + error = self.config.get_error_class( + error_message=f"err {status_code}", + status_code=status_code, + headers={"Content-Type": "application/json"}, + ) + assert isinstance(error, ModelScopeError), status_code + assert error.status_code == status_code + assert f"err {status_code}" in error.message - error = self.config.get_error_class( - error_message="Bad request", - status_code=400, - headers={"Content-Type": "application/json"}, + +class TestModelScopeImageGenerationHandler: + """Tests for the async submit + poll flow in ModelScopeImageGeneration.""" + + def setup_method(self): + from litellm.llms.modelscope.image_generation.handler import ( + ModelScopeImageGeneration, ) - assert isinstance(error, BadRequestError) + self.handler = ModelScopeImageGeneration() + self.model = "Qwen/Qwen-Image-2512" + self.prompt = "a cute baby sea otter" + self.logging_obj = MagicMock() - def test_get_error_class_authentication_error(self): - """Test that get_error_class returns AuthenticationError for 401 status.""" - from litellm.exceptions import AuthenticationError + @staticmethod + def _mock_response(json_data, status_code=200): + mock_response = MagicMock() + mock_response.json.return_value = json_data + mock_response.status_code = status_code + mock_response.text = str(json_data) + mock_response.headers = {} + return mock_response - error = self.config.get_error_class( - error_message="Invalid API key", - status_code=401, - headers={"Content-Type": "application/json"}, + def test_sync_submit_and_poll_succeed(self): + """Submit returns task_id, poll returns SUCCEED with output_images.""" + from litellm.types.utils import ImageResponse + + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = self._mock_response( + { + "task_status": "SUCCEED", + "task_id": "task-123", + "output_images": ["https://example.com/otter.png"], + } ) - assert isinstance(error, AuthenticationError) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + mock_client.get.return_value = poll_resp - def test_get_error_class_internal_server_error(self): - """Test that get_error_class returns InternalServerError for 500+ status.""" - from litellm.exceptions import InternalServerError - - error = self.config.get_error_class( - error_message="Internal server error", - status_code=500, - headers={"Content-Type": "application/json"}, + result = self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=ImageResponse(data=[]), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, ) - assert isinstance(error, InternalServerError) + assert len(result.data) == 1 + assert result.data[0].url == "https://example.com/otter.png" + # Submit POST hits /images/generations, poll GET hits /tasks/{task_id} + assert mock_client.post.call_count == 1 + assert "/images/generations" in mock_client.post.call_args.kwargs["url"] + assert mock_client.get.call_count == 1 + assert "/tasks/task-123" in mock_client.get.call_args.kwargs["url"] - def test_get_error_class_default(self): - """Test that get_error_class returns BadRequestError for other status codes.""" - from litellm.exceptions import BadRequestError + def test_sync_submit_returns_error_status_raises(self): + """A 4xx/5xx on the submit call raises immediately.""" + submit_resp = self._mock_response({"errors": {"message": "bad model"}}, status_code=400) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp - error = self.config.get_error_class( - error_message="Some error", - status_code=404, - headers={"Content-Type": "application/json"}, + with pytest.raises(ModelScopeError): + self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=MagicMock(), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, + ) + mock_client.get.assert_not_called() + + def test_sync_poll_failed_status_raises(self): + """A FAILED task_status during polling raises.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = self._mock_response({"task_status": "FAILED", "task_id": "task-123", "output_images": []}) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + mock_client.get.return_value = poll_resp + + with pytest.raises(ModelScopeError): + self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=MagicMock(), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, + ) + + def test_sync_submit_missing_task_id_raises(self): + """A submit response without task_id raises (cannot poll).""" + submit_resp = self._mock_response({"task_status": "SUCCEED"}) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + + with pytest.raises(ModelScopeError): + self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=MagicMock(), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, + ) + + def test_submit_sends_async_mode_header(self): + """The submit POST must include X-ModelScope-Async-Mode: true.""" + from litellm.types.utils import ImageResponse + + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = self._mock_response( + { + "task_status": "SUCCEED", + "task_id": "task-123", + "output_images": ["https://example.com/x.png"], + } + ) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + mock_client.get.return_value = poll_resp + + self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=ImageResponse(data=[]), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, ) - assert isinstance(error, BadRequestError) + submit_headers = mock_client.post.call_args.kwargs["headers"] + assert submit_headers.get("X-ModelScope-Async-Mode") == "true" + # Polling must carry the task-type header, not the async-mode header + poll_headers = mock_client.get.call_args.kwargs["headers"] + assert poll_headers.get("X-ModelScope-Task-Type") == "image_generation" + assert "X-ModelScope-Async-Mode" not in poll_headers + + def test_sync_poll_running_then_succeed(self): + """The poll loop must keep polling while RUNNING and stop on SUCCEED.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []}) + succeed_resp = self._mock_response( + { + "task_status": "SUCCEED", + "task_id": "task-123", + "output_images": ["https://example.com/otter.png"], + } + ) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + mock_client.get.side_effect = [running_resp, succeed_resp] + + with patch( # test-quality-ok: avoid real sleep delay in sync polling test + "litellm.llms.modelscope.image_generation.handler.time.sleep" + ): + result = self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=ImageResponse(data=[]), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, + ) + + assert len(result.data) == 1 + assert result.data[0].url == "https://example.com/otter.png" + # RUNNING did not terminate the loop; SUCCEED did on the second GET. + assert mock_client.get.call_count == 2 + + def test_sync_poll_timeout_raises(self): + """A task that never leaves RUNNING must time out with a 408.""" + submit_resp = self._mock_response({"task_id": "task-123"}) + running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []}) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.get.return_value = running_resp + + with pytest.raises(ModelScopeError) as exc_info: + self.handler._poll_for_result_sync( + initial_response=submit_resp, + api_base=None, + headers={"Authorization": "Bearer ms-test"}, + sync_client=mock_client, + max_wait=0.05, + interval=0.01, + timeout=30, + ) + + assert exc_info.value.status_code == 408 + assert "timed out" in str(exc_info.value).lower() + + def test_sync_submit_200_with_errors_raises(self): + """A 200 submit response carrying an `errors` key must raise.""" + submit_resp = self._mock_response({"errors": {"message": "rate limited"}}, status_code=200) + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = submit_resp + + with pytest.raises(ModelScopeError) as exc_info: + self.handler.image_generation( + model=self.model, + prompt=self.prompt, + model_response=ImageResponse(data=[]), + optional_params={}, + litellm_params={"api_key": "ms-test", "api_base": None}, + logging_obj=self.logging_obj, + timeout=30, + client=mock_client, + aimg_generation=False, + ) + + assert "rate limited" in str(exc_info.value) + mock_client.get.assert_not_called() + + def test_sync_poll_non_json_response_raises(self): + """A non-JSON poll response (e.g. gateway HTML) raises ModelScopeError, + not a raw JSONDecodeError that bypasses error normalization.""" + import json as _json + + submit_resp = self._mock_response({"task_id": "task-123"}) + poll_resp = MagicMock() + poll_resp.status_code = 200 + poll_resp.text = "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 + + +class TestModelScopeImageGenerationDispatch: + """Tests for the modelscope dispatch branch in images/main.py.""" + + def test_api_base_override_without_api_key_raises(self): + """Security: caller-controlled api_base must not receive deployment key.""" + import litellm + + with pytest.raises(Exception, match="api_key must be provided when api_base is overridden"): + litellm.image_generation( + prompt="test", + model="modelscope/Qwen/Qwen-Image-2512", + api_base="https://evil.example.com/v1", + api_key=None, + ) + + def test_dispatch_passes_api_key_and_api_base(self): + """Deployment credentials are forwarded to the handler.""" + import litellm.images.main as img_main + from litellm.types.utils import ImageResponse + + captured: dict = {} + + original = img_main.modelscope_image_generation.image_generation + + def spy(**kwargs): + captured.update(kwargs) + return ImageResponse(data=[]) + + with patch.object( # test-quality-ok: verify dispatch wiring without real HTTP + img_main.modelscope_image_generation, "image_generation", spy + ): + img_main.image_generation( + prompt="test", + model="modelscope/Qwen/Qwen-Image-2512", + api_key="ms-deploy-key", + api_base="https://api-inference.modelscope.cn/v1", + custom_llm_provider="modelscope", + ) + + assert captured["litellm_params"]["api_key"] == "ms-deploy-key" + assert captured["litellm_params"]["api_base"] == "https://api-inference.modelscope.cn/v1"