fix(modelscope): prevent extra_body model override; satisfy ruff strict budget

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,

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"

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()

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,
-        )
+        )

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"