From 3121c97ebe8ec26915d95ed3a76ff0926823025b Mon Sep 17 00:00:00 2001 From: yrk <2493404415@qq.com> Date: Tue, 7 Jul 2026 10:13:09 +0800 Subject: [PATCH] fix(modelscope): prevent extra_body model override; satisfy ruff strict budget --- litellm/images/main.py | 6 +- litellm/llms/modelscope/common_utils.py | 21 +-- .../modelscope/image_generation/handler.py | 171 +++++++++--------- .../image_generation/transformation.py | 106 +++++------ ...est_modelscope_image_gen_transformation.py | 103 +++++++++-- 5 files changed, 230 insertions(+), 177 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 3158939ed0b..617ea1ecc36 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -426,10 +426,12 @@ def image_generation( 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") + 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, diff --git a/litellm/llms/modelscope/common_utils.py b/litellm/llms/modelscope/common_utils.py index b30116dbf8c..6906277d280 100644 --- a/litellm/llms/modelscope/common_utils.py +++ b/litellm/llms/modelscope/common_utils.py @@ -4,24 +4,21 @@ 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.""" - pass +DEFAULT_POLLING_INTERVAL: Final = 2.0 +DEFAULT_MAX_POLLING_TIME: Final = 300 -# 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: Final = "X-ModelScope-Async-Mode" +TASK_TYPE_HEADER: Final = "X-ModelScope-Task-Type" +IMAGE_GENERATION_TASK_TYPE: Final = "image_generation" -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" +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 index 68383df8f2c..b04a4b32c37 100644 --- a/litellm/llms/modelscope/image_generation/handler.py +++ b/litellm/llms/modelscope/image_generation/handler.py @@ -9,8 +9,9 @@ API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro """ import asyncio +import json import time -from typing import Any, Dict, Optional, Union +from typing import Final import httpx @@ -20,7 +21,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, - _get_httpx_client, + _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 ( @@ -37,29 +38,24 @@ from .transformation import ModelScopeImageGenerationConfig class ModelScopeImageGeneration: - """ - ModelScope image generation handler. + """ModelScope image generation handler.""" - Owns the submit + poll HTTP flow; request/response shaping is delegated to - ModelScopeImageGenerationConfig. - """ - - def __init__(self): - self.config = ModelScopeImageGenerationConfig() + def __init__(self) -> None: + self.config: Final = ModelScopeImageGenerationConfig() def image_generation( self, model: str, prompt: str, model_response: ImageResponse, - optional_params: Dict, - litellm_params: Union[GenericLiteLLMParams, Dict], + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature logging_obj: LiteLLMLoggingObj, - timeout: Optional[Union[float, httpx.Timeout]], - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + 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, - ) -> Union[ImageResponse, Any]: + ) -> ImageResponse: if aimg_generation: return self.async_image_generation( model=model, @@ -73,29 +69,30 @@ class ModelScopeImageGeneration: 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) + 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 = client if isinstance(client, HTTPHandler) else _get_httpx_client() + sync_client: Final = client if isinstance(client, HTTPHandler) else _get_httpx_client() - headers = self.config.validate_environment( + headers: Final = self.config.validate_environment( api_key=api_key, - headers={}, + headers={}, # mutable-ok: validate_environment returns a new dict model=model, - messages=[], + 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 = self.config.get_complete_url( + complete_url: Final = self.config.get_complete_url( api_base=api_base, api_key=api_key, model=model, @@ -103,7 +100,7 @@ class ModelScopeImageGeneration: litellm_params=litellm_params_dict, ) - data = self.config.transform_image_generation_request( + data: Final = self.config.transform_image_generation_request( model=model, prompt=prompt, optional_params=optional_params, @@ -114,7 +111,7 @@ class ModelScopeImageGeneration: logging_obj.pre_call( input=prompt, api_key="", - additional_args={ + additional_args={ # mutable-ok: litellm logging contract "complete_input_dict": data, "api_base": complete_url, "headers": headers, @@ -122,20 +119,20 @@ class ModelScopeImageGeneration: ) try: - response = sync_client.post( + response: Final = sync_client.post( url=complete_url, headers=headers, json=data, timeout=timeout, ) - except Exception as e: + except httpx.HTTPError as e: raise ModelScopeError( status_code=500, - message=f"Request failed: {str(e)}", + message=f"Request failed: {e}", ) - final_response = self._poll_for_result_sync( - initial_response=response, + 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, @@ -158,38 +155,39 @@ class ModelScopeImageGeneration: model: str, prompt: str, model_response: ImageResponse, - optional_params: Dict, - litellm_params: Union[GenericLiteLLMParams, Dict], + optional_params: dict, # mutable-ok: litellm override signature + litellm_params: GenericLiteLLMParams | dict, # mutable-ok: litellm override signature logging_obj: LiteLLMLoggingObj, - timeout: Optional[Union[float, httpx.Timeout]], - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[AsyncHTTPHandler] = None, + timeout: float | httpx.Timeout | None, + extra_headers: dict[str, str] | None = None, # mutable-ok: litellm override signature + client: AsyncHTTPHandler | None = 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) + 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 = client or get_async_httpx_client( + async_client: Final = client or get_async_httpx_client( llm_provider=litellm.LlmProviders.MODELSCOPE, ) - headers = self.config.validate_environment( + headers: Final = self.config.validate_environment( api_key=api_key, - headers={}, + headers={}, # mutable-ok: validate_environment returns a new dict model=model, - messages=[], + 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 = self.config.get_complete_url( + complete_url: Final = self.config.get_complete_url( api_base=api_base, api_key=api_key, model=model, @@ -197,7 +195,7 @@ class ModelScopeImageGeneration: litellm_params=litellm_params_dict, ) - data = self.config.transform_image_generation_request( + data: Final = self.config.transform_image_generation_request( model=model, prompt=prompt, optional_params=optional_params, @@ -208,7 +206,7 @@ class ModelScopeImageGeneration: logging_obj.pre_call( input=prompt, api_key="", - additional_args={ + additional_args={ # mutable-ok: litellm logging contract "complete_input_dict": data, "api_base": complete_url, "headers": headers, @@ -216,20 +214,20 @@ class ModelScopeImageGeneration: ) try: - response = await async_client.post( + response: Final = await async_client.post( url=complete_url, headers=headers, json=data, timeout=timeout, ) - except Exception as e: + except httpx.HTTPError as e: raise ModelScopeError( status_code=500, - message=f"Request failed: {str(e)}", + message=f"Request failed: {e}", ) - final_response = await self._poll_for_result_async( - initial_response=response, + 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, @@ -250,12 +248,12 @@ class ModelScopeImageGeneration: def _poll_for_result_sync( self, initial_response: httpx.Response, - api_base: Optional[str], - headers: dict, + 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: Optional[Union[float, httpx.Timeout]] = None, + timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: if initial_response.status_code >= 400: raise ModelScopeError( @@ -264,8 +262,8 @@ class ModelScopeImageGeneration: ) try: - response_data = initial_response.json() - except Exception as e: + 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}", @@ -277,18 +275,18 @@ class ModelScopeImageGeneration: message=f"ModelScope error: {response_data['errors']}", ) - task_id = response_data.get("task_id") + 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 = self.config.get_task_status_url(api_base, task_id) - polling_headers = self.config.get_polling_headers(headers) + polling_url: Final = self.config.get_task_status_url(api_base, task_id) + polling_headers: Final = self.config.get_polling_headers(headers) - start_time = time.time() - verbose_logger.debug(f"ModelScope starting sync polling at {polling_url}") + 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( @@ -305,13 +303,13 @@ class ModelScopeImageGeneration: try: data = response.json() - except Exception as e: + 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(f"ModelScope poll status: {status}") + verbose_logger.debug("ModelScope poll status: %s", status) if status == TASK_STATUS_SUCCEED: return response @@ -331,12 +329,12 @@ class ModelScopeImageGeneration: async def _poll_for_result_async( self, initial_response: httpx.Response, - api_base: Optional[str], - headers: dict, + 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: Optional[Union[float, httpx.Timeout]] = None, + timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: if initial_response.status_code >= 400: raise ModelScopeError( @@ -345,8 +343,8 @@ class ModelScopeImageGeneration: ) try: - response_data = initial_response.json() - except Exception as e: + 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}", @@ -358,18 +356,18 @@ class ModelScopeImageGeneration: message=f"ModelScope error: {response_data['errors']}", ) - task_id = response_data.get("task_id") + 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 = self.config.get_task_status_url(api_base, task_id) - polling_headers = self.config.get_polling_headers(headers) + polling_url: Final = self.config.get_task_status_url(api_base, task_id) + polling_headers: Final = self.config.get_polling_headers(headers) - start_time = time.time() - verbose_logger.debug(f"ModelScope starting async polling at {polling_url}") + 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( @@ -386,13 +384,13 @@ class ModelScopeImageGeneration: try: data = response.json() - except Exception as e: + 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(f"ModelScope poll status: {status}") + verbose_logger.debug("ModelScope poll status: %s", status) if status == TASK_STATUS_SUCCEED: return response @@ -410,5 +408,4 @@ class ModelScopeImageGeneration: ) -# Singleton instance for use in images/main.py -modelscope_image_generation = ModelScopeImageGeneration() +modelscope_image_generation: Final = ModelScopeImageGeneration() diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 5d89ccc5c2b..570ce692dee 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -33,40 +33,40 @@ 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 and editing models.""" - DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1" + DEFAULT_BASE_URL: Final[str] = "https://api-inference.modelscope.cn/v1" @override - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - # 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"] + 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: + ) -> 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: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL + base_url: Final[str] = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL return base_url.rstrip("/") @override @@ -75,30 +75,26 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): 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: return f"{self._get_base_url(api_base)}/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( 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: @@ -106,19 +102,19 @@ 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}", - # 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 + return {**headers, **default_headers} # mutable-ok: merged header dict - def get_polling_headers(self, headers: dict) -> dict: - return { + 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, } @@ -128,22 +124,17 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): self, model: str, prompt: str, - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - 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, } - # 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(): + 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 @@ -157,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 a completed ModelScope task response into an ImageResponse. - - The polled response looks like: - {"task_status": "SUCCEED", "output_images": ["https://..."], ...} - """ try: response_data: Final = raw_response.json() except Exception as e: @@ -180,15 +165,15 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): ) if "errors" in response_data: - errors = response_data["errors"] - error_msg = errors.get("message", str(errors)) if isinstance(errors, dict) else str(errors) + 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, ) - task_status = response_data.get("task_status") + 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", @@ -197,12 +182,12 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): ) if task_status != TASK_STATUS_SUCCEED: raise self.get_error_class( - error_message=(f"ModelScope task did not succeed: status={task_status}"), + 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 [] + 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)) @@ -221,11 +206,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): self, error_message: str, status_code: int, - headers: dict | httpx.Headers, + headers: dict | httpx.Headers, # mutable-ok: litellm override signature ) -> BaseLLMException: - """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 f22c911230d..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 @@ -144,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, @@ -270,6 +270,33 @@ class TestModelScopeImageGenerationTransformation: 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 extracts output_images URLs.""" response_data = { @@ -318,7 +345,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + 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, @@ -377,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, @@ -403,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, @@ -600,17 +627,20 @@ class TestModelScopeImageGenerationHandler: 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, - ) + 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" @@ -770,3 +800,46 @@ class TestModelScopeImageGenerationHandler: ) 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"