mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge 3121c97ebe into 252c71c0b2
This commit is contained in:
commit
a5e07be000
6 changed files with 1093 additions and 264 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
24
litellm/llms/modelscope/common_utils.py
Normal file
24
litellm/llms/modelscope/common_utils.py
Normal file
|
|
@ -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"
|
||||
411
litellm/llms/modelscope/image_generation/handler.py
Normal file
411
litellm/llms/modelscope/image_generation/handler.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 = "<html>502 Bad Gateway</html>"
|
||||
poll_resp.headers = {}
|
||||
poll_resp.json.side_effect = _json.JSONDecodeError("Expecting value", "<html>", 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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue