mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(modelscope): use async task polling for image generation
This commit is contained in:
yrk
parent
10631eb834
commit
c55cae4642
6 changed files with 1005 additions and 229 deletions
|
|
@ -1173,6 +1173,8 @@ modelscope_models: Final[set] = set(
|
|||
"Qwen/QwQ-32B-Preview",
|
||||
"Qwen/QVQ-72B-Preview",
|
||||
"Qwen/Qwen-Image-Edit",
|
||||
"Qwen/Qwen-Image-Edit-2511",
|
||||
"Qwen/Qwen-Image-2512",
|
||||
# DeepSeek series models
|
||||
"deepseek-ai/DeepSeek-R1-0528",
|
||||
"deepseek-ai/DeepSeek-R1-Distill-Llama-70B",
|
||||
|
|
|
|||
|
|
@ -34,6 +34,9 @@ from openai.types.audio.transcription_create_params import FileTypes
|
|||
# BFL handlers
|
||||
from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
|
||||
from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
|
||||
|
||||
# ModelScope handler
|
||||
from litellm.llms.modelscope.image_generation.handler import modelscope_image_generation
|
||||
from litellm.main import (
|
||||
azure_chat_completions,
|
||||
base_llm_aiohttp_handler,
|
||||
|
|
@ -422,6 +425,23 @@ def image_generation(
|
|||
client=client,
|
||||
aimg_generation=aimg_generation,
|
||||
)
|
||||
elif custom_llm_provider == "modelscope":
|
||||
# ModelScope image gen is async: submit returns a task_id that
|
||||
# must be polled until SUCCEED/FAILED.
|
||||
if model is None:
|
||||
raise Exception("Model needs to be set for modelscope")
|
||||
return modelscope_image_generation.image_generation(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
model_response=model_response,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
logging_obj=litellm_logging_obj,
|
||||
timeout=timeout,
|
||||
extra_headers=extra_headers,
|
||||
client=client,
|
||||
aimg_generation=aimg_generation,
|
||||
)
|
||||
elif custom_llm_provider == "azure_ai":
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
|
|
|
|||
27
litellm/llms/modelscope/common_utils.py
Normal file
27
litellm/llms/modelscope/common_utils.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
"""
|
||||
ModelScope Common Utilities
|
||||
|
||||
Shared constants and error handling for ModelScope API integration.
|
||||
"""
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
|
||||
class ModelScopeError(BaseLLMException):
|
||||
"""Exception class for ModelScope API errors."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
# Polling configuration for async image generation tasks.
|
||||
# ModelScope returns a task_id from the submit call and the caller must poll
|
||||
# GET /v1/tasks/{task_id} until task_status is SUCCEED or FAILED.
|
||||
DEFAULT_POLLING_INTERVAL = 2.0 # seconds
|
||||
DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes
|
||||
|
||||
ASYNC_MODE_HEADER = "X-ModelScope-Async-Mode"
|
||||
TASK_TYPE_HEADER = "X-ModelScope-Task-Type"
|
||||
IMAGE_GENERATION_TASK_TYPE = "image_generation"
|
||||
|
||||
TASK_STATUS_SUCCEED = "SUCCEED"
|
||||
TASK_STATUS_FAILED = "FAILED"
|
||||
414
litellm/llms/modelscope/image_generation/handler.py
Normal file
414
litellm/llms/modelscope/image_generation/handler.py
Normal file
|
|
@ -0,0 +1,414 @@
|
|||
"""
|
||||
ModelScope Image Generation Handler
|
||||
|
||||
ModelScope image generation only supports async mode: the submit call returns
|
||||
a task_id, then the caller polls GET /v1/tasks/{task_id} until task_status is
|
||||
SUCCEED (image URLs in output_images) or FAILED.
|
||||
|
||||
API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.llms.modelscope.common_utils import (
|
||||
DEFAULT_MAX_POLLING_TIME,
|
||||
DEFAULT_POLLING_INTERVAL,
|
||||
TASK_STATUS_FAILED,
|
||||
TASK_STATUS_SUCCEED,
|
||||
ModelScopeError,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
from .transformation import ModelScopeImageGenerationConfig
|
||||
|
||||
|
||||
class ModelScopeImageGeneration:
|
||||
"""
|
||||
ModelScope image generation handler.
|
||||
|
||||
Owns the submit + poll HTTP flow; request/response shaping is delegated to
|
||||
ModelScopeImageGenerationConfig.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.config = ModelScopeImageGenerationConfig()
|
||||
|
||||
def image_generation(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
model_response: ImageResponse,
|
||||
optional_params: Dict,
|
||||
litellm_params: Union[GenericLiteLLMParams, Dict],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
extra_headers: Optional[Dict[str, Any]] = None,
|
||||
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
|
||||
aimg_generation: bool = False,
|
||||
) -> Union[ImageResponse, Any]:
|
||||
if aimg_generation:
|
||||
return self.async_image_generation(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
model_response=model_response,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
extra_headers=extra_headers,
|
||||
client=client if isinstance(client, AsyncHTTPHandler) else None,
|
||||
)
|
||||
|
||||
if isinstance(litellm_params, dict):
|
||||
api_key = litellm_params.get("api_key")
|
||||
api_base = litellm_params.get("api_base")
|
||||
litellm_params_dict = litellm_params
|
||||
else:
|
||||
api_key = litellm_params.api_key
|
||||
api_base = litellm_params.api_base
|
||||
litellm_params_dict = dict(litellm_params)
|
||||
|
||||
sync_client = client if isinstance(client, HTTPHandler) else _get_httpx_client()
|
||||
|
||||
headers = self.config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers={},
|
||||
model=model,
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
complete_url = self.config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
|
||||
data = self.config.transform_image_generation_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=prompt,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": complete_url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response = sync_client.post(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=500,
|
||||
message=f"Request failed: {str(e)}",
|
||||
)
|
||||
|
||||
final_response = self._poll_for_result_sync(
|
||||
initial_response=response,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
sync_client=sync_client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return self.config.transform_image_generation_response(
|
||||
model=model,
|
||||
raw_response=final_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=data,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
async def async_image_generation(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
model_response: ImageResponse,
|
||||
optional_params: Dict,
|
||||
litellm_params: Union[GenericLiteLLMParams, Dict],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
extra_headers: Optional[Dict[str, Any]] = None,
|
||||
client: Optional[AsyncHTTPHandler] = None,
|
||||
) -> ImageResponse:
|
||||
if isinstance(litellm_params, dict):
|
||||
api_key = litellm_params.get("api_key")
|
||||
api_base = litellm_params.get("api_base")
|
||||
litellm_params_dict = litellm_params
|
||||
else:
|
||||
api_key = litellm_params.api_key
|
||||
api_base = litellm_params.api_base
|
||||
litellm_params_dict = dict(litellm_params)
|
||||
|
||||
async_client = client or get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders.MODELSCOPE,
|
||||
)
|
||||
|
||||
headers = self.config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers={},
|
||||
model=model,
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
complete_url = self.config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
|
||||
data = self.config.transform_image_generation_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=prompt,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": complete_url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response = await async_client.post(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=500,
|
||||
message=f"Request failed: {str(e)}",
|
||||
)
|
||||
|
||||
final_response = await self._poll_for_result_async(
|
||||
initial_response=response,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
async_client=async_client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return self.config.transform_image_generation_response(
|
||||
model=model,
|
||||
raw_response=final_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=data,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
def _poll_for_result_sync(
|
||||
self,
|
||||
initial_response: httpx.Response,
|
||||
api_base: Optional[str],
|
||||
headers: dict,
|
||||
sync_client: HTTPHandler,
|
||||
max_wait: float = DEFAULT_MAX_POLLING_TIME,
|
||||
interval: float = DEFAULT_POLLING_INTERVAL,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> httpx.Response:
|
||||
if initial_response.status_code >= 400:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"ModelScope submit failed: {initial_response.text}",
|
||||
)
|
||||
|
||||
try:
|
||||
response_data = initial_response.json()
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"Error parsing submit response: {e}",
|
||||
)
|
||||
|
||||
if "errors" in response_data:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"ModelScope error: {response_data['errors']}",
|
||||
)
|
||||
|
||||
task_id = response_data.get("task_id")
|
||||
if not task_id:
|
||||
raise ModelScopeError(
|
||||
status_code=500,
|
||||
message="No task_id in ModelScope submit response",
|
||||
)
|
||||
|
||||
polling_url = self.config.get_task_status_url(api_base, task_id)
|
||||
polling_headers = self.config.get_polling_headers(headers)
|
||||
|
||||
start_time = time.time()
|
||||
verbose_logger.debug(f"ModelScope starting sync polling at {polling_url}")
|
||||
|
||||
while time.time() - start_time < max_wait:
|
||||
response = sync_client.get(
|
||||
url=polling_url,
|
||||
headers=polling_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise ModelScopeError(
|
||||
status_code=response.status_code,
|
||||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=response.status_code,
|
||||
message=f"Error parsing poll response: {e}",
|
||||
)
|
||||
status = data.get("task_status")
|
||||
verbose_logger.debug(f"ModelScope poll status: {status}")
|
||||
|
||||
if status == TASK_STATUS_SUCCEED:
|
||||
return response
|
||||
elif status == TASK_STATUS_FAILED:
|
||||
raise ModelScopeError(
|
||||
status_code=400,
|
||||
message="ModelScope image generation task failed",
|
||||
)
|
||||
|
||||
time.sleep(interval)
|
||||
|
||||
raise ModelScopeError(
|
||||
status_code=408,
|
||||
message=f"Polling timed out after {max_wait} seconds",
|
||||
)
|
||||
|
||||
async def _poll_for_result_async(
|
||||
self,
|
||||
initial_response: httpx.Response,
|
||||
api_base: Optional[str],
|
||||
headers: dict,
|
||||
async_client: AsyncHTTPHandler,
|
||||
max_wait: float = DEFAULT_MAX_POLLING_TIME,
|
||||
interval: float = DEFAULT_POLLING_INTERVAL,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> httpx.Response:
|
||||
if initial_response.status_code >= 400:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"ModelScope submit failed: {initial_response.text}",
|
||||
)
|
||||
|
||||
try:
|
||||
response_data = initial_response.json()
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"Error parsing submit response: {e}",
|
||||
)
|
||||
|
||||
if "errors" in response_data:
|
||||
raise ModelScopeError(
|
||||
status_code=initial_response.status_code,
|
||||
message=f"ModelScope error: {response_data['errors']}",
|
||||
)
|
||||
|
||||
task_id = response_data.get("task_id")
|
||||
if not task_id:
|
||||
raise ModelScopeError(
|
||||
status_code=500,
|
||||
message="No task_id in ModelScope submit response",
|
||||
)
|
||||
|
||||
polling_url = self.config.get_task_status_url(api_base, task_id)
|
||||
polling_headers = self.config.get_polling_headers(headers)
|
||||
|
||||
start_time = time.time()
|
||||
verbose_logger.debug(f"ModelScope starting async polling at {polling_url}")
|
||||
|
||||
while time.time() - start_time < max_wait:
|
||||
response = await async_client.get(
|
||||
url=polling_url,
|
||||
headers=polling_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise ModelScopeError(
|
||||
status_code=response.status_code,
|
||||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception as e:
|
||||
raise ModelScopeError(
|
||||
status_code=response.status_code,
|
||||
message=f"Error parsing poll response: {e}",
|
||||
)
|
||||
status = data.get("task_status")
|
||||
verbose_logger.debug(f"ModelScope poll status: {status}")
|
||||
|
||||
if status == TASK_STATUS_SUCCEED:
|
||||
return response
|
||||
elif status == TASK_STATUS_FAILED:
|
||||
raise ModelScopeError(
|
||||
status_code=400,
|
||||
message="ModelScope image generation task failed",
|
||||
)
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
raise ModelScopeError(
|
||||
status_code=408,
|
||||
message=f"Polling timed out after {max_wait} seconds",
|
||||
)
|
||||
|
||||
|
||||
# Singleton instance for use in images/main.py
|
||||
modelscope_image_generation = ModelScopeImageGeneration()
|
||||
|
|
@ -15,6 +15,14 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.modelscope.common_utils import (
|
||||
ASYNC_MODE_HEADER,
|
||||
IMAGE_GENERATION_TASK_TYPE,
|
||||
TASK_STATUS_FAILED,
|
||||
TASK_STATUS_SUCCEED,
|
||||
TASK_TYPE_HEADER,
|
||||
ModelScopeError,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -31,29 +39,19 @@ else:
|
|||
|
||||
|
||||
class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""
|
||||
Configuration for ModelScope image generation.
|
||||
|
||||
Supports text-to-image models like:
|
||||
- Qwen/Qwen-Image-Edit
|
||||
- And other ModelScope-hosted image generation models
|
||||
"""
|
||||
"""Configuration for ModelScope image generation and editing models."""
|
||||
|
||||
DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1"
|
||||
|
||||
@override
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
"""
|
||||
Return list of OpenAI params supported by ModelScope.
|
||||
|
||||
ModelScope supports standard OpenAI image generation parameters.
|
||||
"""
|
||||
return [
|
||||
"n", # Number of images to generate
|
||||
"size", # Size of the generated images
|
||||
"response_format", # url or b64_json
|
||||
"user", # User identifier
|
||||
]
|
||||
# Only size is honored (verified via real API: n is silently ignored,
|
||||
# response_format/user are not honored by ModelScope).
|
||||
# Provider-specific fields (negative_prompt, seed, steps, guidance,
|
||||
# loras, image_url) pass through via extra_body, not this list.
|
||||
return ["size"]
|
||||
|
||||
@override
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
|
|
@ -61,17 +59,16 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
"""
|
||||
Map OpenAI parameters to ModelScope parameters.
|
||||
|
||||
ModelScope uses the same parameter names as OpenAI.
|
||||
"""
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
if drop_params:
|
||||
non_default_params = {k: v for k, v in non_default_params.items() if k in supported_params}
|
||||
optional_params.update(non_default_params)
|
||||
return optional_params
|
||||
|
||||
def _get_base_url(self, api_base: str | None) -> str:
|
||||
base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
|
||||
return base_url.rstrip("/")
|
||||
|
||||
@override
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
@ -82,14 +79,11 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Get the complete URL for the ModelScope image generation API request.
|
||||
"""
|
||||
base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
|
||||
base_url = base_url.rstrip("/")
|
||||
return f"{self._get_base_url(api_base)}/images/generations"
|
||||
|
||||
# Return the images endpoint
|
||||
return f"{base_url}/images/generations"
|
||||
def get_task_status_url(self, api_base: str | None, task_id: str) -> str:
|
||||
"""Build the URL used to poll an async image generation task."""
|
||||
return f"{self._get_base_url(api_base)}/tasks/{task_id}"
|
||||
|
||||
@override
|
||||
def validate_environment(
|
||||
|
|
@ -115,11 +109,21 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
default_headers: Final = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {final_api_key}",
|
||||
# ModelScope image-gen only supports async mode: the submit call
|
||||
# returns a task_id and the caller must poll GET /v1/tasks/{task_id}.
|
||||
ASYNC_MODE_HEADER: "true",
|
||||
}
|
||||
|
||||
headers = {**headers, **default_headers}
|
||||
return headers
|
||||
|
||||
def get_polling_headers(self, headers: dict) -> dict:
|
||||
return {
|
||||
"Authorization": headers.get("Authorization", ""),
|
||||
TASK_TYPE_HEADER: IMAGE_GENERATION_TASK_TYPE,
|
||||
}
|
||||
|
||||
@override
|
||||
def transform_image_generation_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -128,20 +132,19 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform OpenAI-style request to ModelScope request format.
|
||||
|
||||
ModelScope uses the same format as OpenAI for image generation.
|
||||
"""
|
||||
# Build the request body (same as OpenAI)
|
||||
request_data: Final[dict] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
}
|
||||
|
||||
# Add optional params
|
||||
for key, value in optional_params.items():
|
||||
if key.startswith("_"):
|
||||
# litellm wraps provider fields (e.g. image_url) in extra_body; flatten
|
||||
# them into the body. Skip extra_headers/extra_query (litellm control
|
||||
# params swept in alongside, may carry secrets) and model/prompt (set
|
||||
# above from the routed values; a client-controlled extra_body must not
|
||||
# override the authorized model or prompt).
|
||||
extra_body = optional_params.get("extra_body") or {}
|
||||
for key, value in {**optional_params, **extra_body}.items():
|
||||
if key in ("extra_body", "extra_headers", "extra_query", "model", "prompt") or key.startswith("_"):
|
||||
continue
|
||||
request_data[key] = value
|
||||
|
||||
|
|
@ -162,10 +165,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
Transform ModelScope response to OpenAI-compatible ImageResponse.
|
||||
Transform a completed ModelScope task response into an ImageResponse.
|
||||
|
||||
ModelScope returns the same format as OpenAI:
|
||||
{"created": timestamp, "data": [{"url": "..."}]}
|
||||
The polled response looks like:
|
||||
{"task_status": "SUCCEED", "output_images": ["https://..."], ...}
|
||||
"""
|
||||
try:
|
||||
response_data: Final = raw_response.json()
|
||||
|
|
@ -176,64 +179,53 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
# Check for errors in response
|
||||
if "error" in response_data:
|
||||
error_msg: Final = response_data["error"].get("message", str(response_data["error"]))
|
||||
if "errors" in response_data:
|
||||
errors = response_data["errors"]
|
||||
error_msg = errors.get("message", str(errors)) if isinstance(errors, dict) else str(errors)
|
||||
raise self.get_error_class(
|
||||
error_message=f"ModelScope error: {error_msg}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
# Extract images from response
|
||||
data_list: Final = response_data.get("data", [])
|
||||
if not model_response.data:
|
||||
model_response.data = []
|
||||
|
||||
for item in data_list:
|
||||
image_obj = ImageObject(
|
||||
url=item.get("url"),
|
||||
b64_json=item.get("b64_json"),
|
||||
revised_prompt=item.get("revised_prompt"),
|
||||
task_status = response_data.get("task_status")
|
||||
if task_status == TASK_STATUS_FAILED:
|
||||
raise self.get_error_class(
|
||||
error_message="ModelScope image generation task failed",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
if task_status != TASK_STATUS_SUCCEED:
|
||||
raise self.get_error_class(
|
||||
error_message=(f"ModelScope task did not succeed: status={task_status}"),
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
output_images = response_data.get("output_images", []) or []
|
||||
|
||||
for image_url in output_images:
|
||||
model_response.data.append(ImageObject(url=image_url))
|
||||
|
||||
if not model_response.data:
|
||||
raise self.get_error_class(
|
||||
error_message="ModelScope task SUCCEED but no output_images",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
model_response.data.append(image_obj)
|
||||
|
||||
return model_response
|
||||
|
||||
@override
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: dict | httpx.Headers,
|
||||
) -> BaseLLMException:
|
||||
"""Return the appropriate error class for ModelScope."""
|
||||
from litellm.exceptions import (
|
||||
AuthenticationError,
|
||||
BadRequestError,
|
||||
InternalServerError,
|
||||
)
|
||||
|
||||
if status_code == 400:
|
||||
return BadRequestError(
|
||||
message=error_message,
|
||||
model="",
|
||||
llm_provider="modelscope",
|
||||
)
|
||||
elif status_code == 401:
|
||||
return AuthenticationError(
|
||||
message=error_message,
|
||||
model="",
|
||||
llm_provider="modelscope",
|
||||
)
|
||||
elif status_code >= 500:
|
||||
return InternalServerError(
|
||||
message=error_message,
|
||||
model="",
|
||||
llm_provider="modelscope",
|
||||
)
|
||||
else:
|
||||
return BadRequestError(
|
||||
message=error_message,
|
||||
model="",
|
||||
llm_provider="modelscope",
|
||||
)
|
||||
"""Return the ModelScope error class, preserving the real status code."""
|
||||
return ModelScopeError(
|
||||
status_code=status_code,
|
||||
message=error_message,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
@ -5,11 +5,14 @@ These tests validate the ModelScopeImageGenerationConfig class which handles
|
|||
transformation between OpenAI-compatible format and ModelScope API format.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.modelscope.common_utils import ModelScopeError
|
||||
from litellm.llms.modelscope.image_generation.transformation import (
|
||||
ModelScopeImageGenerationConfig,
|
||||
)
|
||||
|
|
@ -24,47 +27,39 @@ class TestModelScopeImageGenerationTransformation:
|
|||
self.logging_obj = MagicMock()
|
||||
|
||||
def test_get_supported_openai_params(self):
|
||||
"""Test that get_supported_openai_params returns correct parameters."""
|
||||
"""Only size is a documented+honored OpenAI param for ModelScope."""
|
||||
supported_params = self.config.get_supported_openai_params(self.model)
|
||||
|
||||
assert "n" in supported_params
|
||||
assert "size" in supported_params
|
||||
assert "response_format" in supported_params
|
||||
assert "user" in supported_params
|
||||
# n is silently ignored (n=2 returns 1 image); user/response_format are
|
||||
# not honored by ModelScope.
|
||||
assert "n" not in supported_params
|
||||
assert "user" not in supported_params
|
||||
assert "response_format" not in supported_params
|
||||
|
||||
def test_map_openai_params(self):
|
||||
"""Test that map_openai_params correctly passes through parameters."""
|
||||
non_default_params = {
|
||||
"n": 2,
|
||||
"size": "1024x1024",
|
||||
"response_format": "url",
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
"""Supported OpenAI param (size) passes through to optional_params."""
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
non_default_params={"size": "1024x1024"},
|
||||
optional_params={},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["n"] == 2
|
||||
assert result["size"] == "1024x1024"
|
||||
assert result["response_format"] == "url"
|
||||
|
||||
def test_map_openai_params_with_user(self):
|
||||
"""Test that map_openai_params correctly passes through user parameter."""
|
||||
non_default_params = {"user": "test-user-123"}
|
||||
optional_params = {}
|
||||
|
||||
def test_map_openai_params_drops_unsupported(self):
|
||||
"""Unsupported OpenAI params (e.g. quality) are dropped when
|
||||
drop_params=True; ModelScope only honors `size`."""
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
non_default_params={"size": "1024x1024", "quality": "hd"},
|
||||
optional_params={},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert result["user"] == "test-user-123"
|
||||
assert result["size"] == "1024x1024"
|
||||
assert "quality" not in result
|
||||
|
||||
def test_get_complete_url_default(self):
|
||||
"""Test that get_complete_url returns default ModelScope URL."""
|
||||
|
|
@ -178,12 +173,10 @@ class TestModelScopeImageGenerationTransformation:
|
|||
assert result["prompt"] == prompt
|
||||
|
||||
def test_transform_image_generation_request_with_optional_params(self):
|
||||
"""Test that transform_image_generation_request includes optional params."""
|
||||
"""Supported OpenAI params (size) are included in the request body."""
|
||||
prompt = "A beautiful sunset"
|
||||
optional_params = {
|
||||
"n": 2,
|
||||
"size": "1024x1024",
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
|
||||
result = self.config.transform_image_generation_request(
|
||||
|
|
@ -196,15 +189,13 @@ class TestModelScopeImageGenerationTransformation:
|
|||
|
||||
assert result["model"] == self.model
|
||||
assert result["prompt"] == prompt
|
||||
assert result["n"] == 2
|
||||
assert result["size"] == "1024x1024"
|
||||
assert result["response_format"] == "b64_json"
|
||||
|
||||
def test_transform_image_generation_request_ignores_internal_params(self):
|
||||
"""Test that transform_image_generation_request ignores params starting with _."""
|
||||
"""Params starting with _ are dropped from the request body."""
|
||||
prompt = "A beautiful sunset"
|
||||
optional_params = {
|
||||
"n": 2,
|
||||
"size": "1024x1024",
|
||||
"_internal_param": "should_be_ignored",
|
||||
}
|
||||
|
||||
|
|
@ -217,16 +208,76 @@ class TestModelScopeImageGenerationTransformation:
|
|||
)
|
||||
|
||||
assert result["model"] == self.model
|
||||
assert result["n"] == 2
|
||||
assert result["size"] == "1024x1024"
|
||||
assert "_internal_param" not in result
|
||||
|
||||
def test_transform_image_generation_request_merges_extra_body(self):
|
||||
"""extra_body fields (e.g. image_url, negative_prompt) are merged into
|
||||
the request body as top-level params, alongside supported OpenAI params
|
||||
like size. ModelScope-specific fields are not OpenAI params, so they
|
||||
arrive via extra_body and must be flattened here."""
|
||||
prompt = "add a birthday hat to the dog"
|
||||
optional_params = {
|
||||
"size": "1024x1024",
|
||||
"extra_body": {
|
||||
"image_url": ["https://example.com/dog.png"],
|
||||
"negative_prompt": "lowres, blurry",
|
||||
},
|
||||
}
|
||||
|
||||
result = self.config.transform_image_generation_request(
|
||||
model=self.model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result["model"] == self.model
|
||||
assert result["prompt"] == prompt
|
||||
assert result["size"] == "1024x1024"
|
||||
# extra_body contents are merged in as top-level params; the extra_body
|
||||
# key itself is not
|
||||
assert result["image_url"] == ["https://example.com/dog.png"]
|
||||
assert result["negative_prompt"] == "lowres, blurry"
|
||||
assert "extra_body" not in result
|
||||
|
||||
def test_transform_image_generation_request_excludes_internal_extra_body_keys(self):
|
||||
"""extra_headers/extra_query are litellm control params, not ModelScope
|
||||
body fields; they must not leak into the request body even though
|
||||
get_optional_params_image_gen sweeps them into extra_body."""
|
||||
prompt = "add a birthday hat to the dog"
|
||||
optional_params = {
|
||||
"size": "1024x1024",
|
||||
"extra_body": {
|
||||
"image_url": ["https://example.com/dog.png"],
|
||||
"extra_headers": {"X-Custom-Auth": "secret-value"},
|
||||
"extra_query": {"ref": "abc"},
|
||||
},
|
||||
}
|
||||
|
||||
result = self.config.transform_image_generation_request(
|
||||
model=self.model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result["image_url"] == ["https://example.com/dog.png"]
|
||||
assert "extra_body" not in result
|
||||
assert "extra_headers" not in result
|
||||
assert "extra_query" not in result
|
||||
assert "secret-value" not in str(result)
|
||||
|
||||
def test_transform_image_generation_response_with_url_images(self):
|
||||
"""Test that transform_image_generation_response correctly extracts URL images."""
|
||||
"""Test that transform_image_generation_response extracts output_images URLs."""
|
||||
response_data = {
|
||||
"created": 1234567890,
|
||||
"data": [
|
||||
{"url": "https://example.com/image1.png"},
|
||||
{"url": "https://example.com/image2.png"},
|
||||
"task_status": "SUCCEED",
|
||||
"task_id": "abc-123",
|
||||
"output_images": [
|
||||
"https://example.com/image1.png",
|
||||
"https://example.com/image2.png",
|
||||
],
|
||||
}
|
||||
|
||||
|
|
@ -252,13 +303,12 @@ class TestModelScopeImageGenerationTransformation:
|
|||
assert result.data[0].url == "https://example.com/image1.png"
|
||||
assert result.data[1].url == "https://example.com/image2.png"
|
||||
|
||||
def test_transform_image_generation_response_with_b64_json(self):
|
||||
"""Test that transform_image_generation_response correctly extracts base64 images."""
|
||||
def test_transform_image_generation_response_failed_status_raises(self):
|
||||
"""Test that a FAILED task_status raises an error."""
|
||||
response_data = {
|
||||
"created": 1234567890,
|
||||
"data": [
|
||||
{"b64_json": "iVBORw0KGgoAAAANS"},
|
||||
],
|
||||
"task_status": "FAILED",
|
||||
"task_id": "abc-123",
|
||||
"output_images": [],
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -268,62 +318,26 @@ class TestModelScopeImageGenerationTransformation:
|
|||
|
||||
model_response = ImageResponse(data=[])
|
||||
|
||||
result = self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].b64_json == "iVBORw0KGgoAAAANS"
|
||||
assert result.data[0].url is None
|
||||
|
||||
def test_transform_image_generation_response_with_revised_prompt(self):
|
||||
"""Test that transform_image_generation_response extracts revised_prompt."""
|
||||
response_data = {
|
||||
"created": 1234567890,
|
||||
"data": [
|
||||
{
|
||||
"url": "https://example.com/image.png",
|
||||
"revised_prompt": "A detailed description of a beautiful sunset",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = response_data
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
|
||||
model_response = ImageResponse(data=[])
|
||||
|
||||
result = self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert (
|
||||
result.data[0].revised_prompt
|
||||
== "A detailed description of a beautiful sunset"
|
||||
)
|
||||
assert "failed" in str(exc_info.value).lower()
|
||||
|
||||
def test_transform_image_generation_response_empty_data(self):
|
||||
"""Test that transform_image_generation_response handles empty data array."""
|
||||
"""SUCCEED with no output_images is an error, not a silent empty success."""
|
||||
response_data = {
|
||||
"created": 1234567890,
|
||||
"data": [],
|
||||
"task_status": "SUCCEED",
|
||||
"task_id": "abc-123",
|
||||
"output_images": [],
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -333,23 +347,24 @@ class TestModelScopeImageGenerationTransformation:
|
|||
|
||||
model_response = ImageResponse(data=[])
|
||||
|
||||
result = self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
with pytest.raises(ModelScopeError) as exc_info:
|
||||
self.config.transform_image_generation_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 0
|
||||
assert "no output_images" in str(exc_info.value).lower()
|
||||
|
||||
def test_transform_image_generation_response_error_handling(self):
|
||||
"""Test that transform_image_generation_response raises error on API error."""
|
||||
response_data = {
|
||||
"error": {
|
||||
"errors": {
|
||||
"message": "Invalid prompt provided",
|
||||
"type": "invalid_request_error",
|
||||
}
|
||||
|
|
@ -402,50 +417,356 @@ class TestModelScopeImageGenerationTransformation:
|
|||
|
||||
assert "Error parsing ModelScope response" in str(exc_info.value)
|
||||
|
||||
def test_get_error_class_bad_request(self):
|
||||
"""Test that get_error_class returns BadRequestError for 400 status."""
|
||||
from litellm.exceptions import BadRequestError
|
||||
def test_get_error_class_returns_modelscope_error(self):
|
||||
"""get_error_class always returns ModelScopeError and preserves status_code."""
|
||||
for status_code in (400, 401, 404, 429, 500, 502):
|
||||
error = self.config.get_error_class(
|
||||
error_message=f"err {status_code}",
|
||||
status_code=status_code,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
assert isinstance(error, ModelScopeError), status_code
|
||||
assert error.status_code == status_code
|
||||
assert f"err {status_code}" in error.message
|
||||
|
||||
error = self.config.get_error_class(
|
||||
error_message="Bad request",
|
||||
status_code=400,
|
||||
headers={"Content-Type": "application/json"},
|
||||
|
||||
class TestModelScopeImageGenerationHandler:
|
||||
"""Tests for the async submit + poll flow in ModelScopeImageGeneration."""
|
||||
|
||||
def setup_method(self):
|
||||
from litellm.llms.modelscope.image_generation.handler import (
|
||||
ModelScopeImageGeneration,
|
||||
)
|
||||
|
||||
assert isinstance(error, BadRequestError)
|
||||
self.handler = ModelScopeImageGeneration()
|
||||
self.model = "Qwen/Qwen-Image-2512"
|
||||
self.prompt = "a cute baby sea otter"
|
||||
self.logging_obj = MagicMock()
|
||||
|
||||
def test_get_error_class_authentication_error(self):
|
||||
"""Test that get_error_class returns AuthenticationError for 401 status."""
|
||||
from litellm.exceptions import AuthenticationError
|
||||
@staticmethod
|
||||
def _mock_response(json_data, status_code=200):
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = json_data
|
||||
mock_response.status_code = status_code
|
||||
mock_response.text = str(json_data)
|
||||
mock_response.headers = {}
|
||||
return mock_response
|
||||
|
||||
error = self.config.get_error_class(
|
||||
error_message="Invalid API key",
|
||||
status_code=401,
|
||||
headers={"Content-Type": "application/json"},
|
||||
def test_sync_submit_and_poll_succeed(self):
|
||||
"""Submit returns task_id, poll returns SUCCEED with output_images."""
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
poll_resp = self._mock_response(
|
||||
{
|
||||
"task_status": "SUCCEED",
|
||||
"task_id": "task-123",
|
||||
"output_images": ["https://example.com/otter.png"],
|
||||
}
|
||||
)
|
||||
|
||||
assert isinstance(error, AuthenticationError)
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
mock_client.get.return_value = poll_resp
|
||||
|
||||
def test_get_error_class_internal_server_error(self):
|
||||
"""Test that get_error_class returns InternalServerError for 500+ status."""
|
||||
from litellm.exceptions import InternalServerError
|
||||
|
||||
error = self.config.get_error_class(
|
||||
error_message="Internal server error",
|
||||
status_code=500,
|
||||
headers={"Content-Type": "application/json"},
|
||||
result = self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=ImageResponse(data=[]),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
assert isinstance(error, InternalServerError)
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].url == "https://example.com/otter.png"
|
||||
# Submit POST hits /images/generations, poll GET hits /tasks/{task_id}
|
||||
assert mock_client.post.call_count == 1
|
||||
assert "/images/generations" in mock_client.post.call_args.kwargs["url"]
|
||||
assert mock_client.get.call_count == 1
|
||||
assert "/tasks/task-123" in mock_client.get.call_args.kwargs["url"]
|
||||
|
||||
def test_get_error_class_default(self):
|
||||
"""Test that get_error_class returns BadRequestError for other status codes."""
|
||||
from litellm.exceptions import BadRequestError
|
||||
def test_sync_submit_returns_error_status_raises(self):
|
||||
"""A 4xx/5xx on the submit call raises immediately."""
|
||||
submit_resp = self._mock_response({"errors": {"message": "bad model"}}, status_code=400)
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
|
||||
error = self.config.get_error_class(
|
||||
error_message="Some error",
|
||||
status_code=404,
|
||||
headers={"Content-Type": "application/json"},
|
||||
with pytest.raises(ModelScopeError):
|
||||
self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=MagicMock(),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
mock_client.get.assert_not_called()
|
||||
|
||||
def test_sync_poll_failed_status_raises(self):
|
||||
"""A FAILED task_status during polling raises."""
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
poll_resp = self._mock_response({"task_status": "FAILED", "task_id": "task-123", "output_images": []})
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
mock_client.get.return_value = poll_resp
|
||||
|
||||
with pytest.raises(ModelScopeError):
|
||||
self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=MagicMock(),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
def test_sync_submit_missing_task_id_raises(self):
|
||||
"""A submit response without task_id raises (cannot poll)."""
|
||||
submit_resp = self._mock_response({"task_status": "SUCCEED"})
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
|
||||
with pytest.raises(ModelScopeError):
|
||||
self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=MagicMock(),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
def test_submit_sends_async_mode_header(self):
|
||||
"""The submit POST must include X-ModelScope-Async-Mode: true."""
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
poll_resp = self._mock_response(
|
||||
{
|
||||
"task_status": "SUCCEED",
|
||||
"task_id": "task-123",
|
||||
"output_images": ["https://example.com/x.png"],
|
||||
}
|
||||
)
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
mock_client.get.return_value = poll_resp
|
||||
|
||||
self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=ImageResponse(data=[]),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
assert isinstance(error, BadRequestError)
|
||||
submit_headers = mock_client.post.call_args.kwargs["headers"]
|
||||
assert submit_headers.get("X-ModelScope-Async-Mode") == "true"
|
||||
# Polling must carry the task-type header, not the async-mode header
|
||||
poll_headers = mock_client.get.call_args.kwargs["headers"]
|
||||
assert poll_headers.get("X-ModelScope-Task-Type") == "image_generation"
|
||||
assert "X-ModelScope-Async-Mode" not in poll_headers
|
||||
|
||||
def test_sync_poll_running_then_succeed(self):
|
||||
"""The poll loop must keep polling while RUNNING and stop on SUCCEED."""
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []})
|
||||
succeed_resp = self._mock_response(
|
||||
{
|
||||
"task_status": "SUCCEED",
|
||||
"task_id": "task-123",
|
||||
"output_images": ["https://example.com/otter.png"],
|
||||
}
|
||||
)
|
||||
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
mock_client.get.side_effect = [running_resp, succeed_resp]
|
||||
|
||||
result = self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=ImageResponse(data=[]),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].url == "https://example.com/otter.png"
|
||||
# RUNNING did not terminate the loop; SUCCEED did on the second GET.
|
||||
assert mock_client.get.call_count == 2
|
||||
|
||||
def test_sync_poll_timeout_raises(self):
|
||||
"""A task that never leaves RUNNING must time out with a 408."""
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
running_resp = self._mock_response({"task_status": "RUNNING", "task_id": "task-123", "output_images": []})
|
||||
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.get.return_value = running_resp
|
||||
|
||||
with pytest.raises(ModelScopeError) as exc_info:
|
||||
self.handler._poll_for_result_sync(
|
||||
initial_response=submit_resp,
|
||||
api_base=None,
|
||||
headers={"Authorization": "Bearer ms-test"},
|
||||
sync_client=mock_client,
|
||||
max_wait=0.05,
|
||||
interval=0.01,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 408
|
||||
assert "timed out" in str(exc_info.value).lower()
|
||||
|
||||
def test_sync_submit_200_with_errors_raises(self):
|
||||
"""A 200 submit response carrying an `errors` key must raise."""
|
||||
submit_resp = self._mock_response({"errors": {"message": "rate limited"}}, status_code=200)
|
||||
mock_client = MagicMock(spec=HTTPHandler)
|
||||
mock_client.post.return_value = submit_resp
|
||||
|
||||
with pytest.raises(ModelScopeError) as exc_info:
|
||||
self.handler.image_generation(
|
||||
model=self.model,
|
||||
prompt=self.prompt,
|
||||
model_response=ImageResponse(data=[]),
|
||||
optional_params={},
|
||||
litellm_params={"api_key": "ms-test", "api_base": None},
|
||||
logging_obj=self.logging_obj,
|
||||
timeout=30,
|
||||
client=mock_client,
|
||||
aimg_generation=False,
|
||||
)
|
||||
|
||||
assert "rate limited" in str(exc_info.value)
|
||||
mock_client.get.assert_not_called()
|
||||
|
||||
def test_sync_poll_non_json_response_raises(self):
|
||||
"""A non-JSON poll response (e.g. gateway HTML) raises ModelScopeError,
|
||||
not a raw JSONDecodeError that bypasses error normalization."""
|
||||
import json as _json
|
||||
|
||||
submit_resp = self._mock_response({"task_id": "task-123"})
|
||||
poll_resp = MagicMock()
|
||||
poll_resp.status_code = 200
|
||||
poll_resp.text = "<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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue