mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
d80f8c28ca
* refactor(types): prove runtime types at harness, search, rag and client boundaries Replace Any with adapter-validated types in the litellm.agent() harness, the search provider transformations, RAG ingestion and query, the vector store pre-call hook and registry, the galileo and opik logging integrations and the proxy client CLI. Each boundary gets unit tests for well-formed and malformed payloads. * chore(typing): prove types at more provider boundaries and restore search transformations Second pass over a2a, embedding, rerank, image, audio and small provider modules. Search transformations go back to their previous form because validating their response bodies would change the proxy status for malformed upstream bodies from 400 to 500. * refactor(types): prove types at logging, files, rerank, image, audio and management boundaries Replace Any with validated or annotation-only types in 37 more files: logging integrations, token counters, provider files/rerank/image generation/audio transcription transformations, pass-through logging handlers and management endpoints. No proxy HTTP status or error type changes. * refactor(types): prove types at repository, spend, files and router boundaries Replace Any with repository table accessors, validated mappings and annotation-only types in 33 more files: Prisma repositories, the enterprise batch and responses cost checkers, budget reservation, files endpoints, management endpoints, the policy registry, the adaptive and complexity routers and the secret managers. No proxy HTTP status or error type changes. * test(types): run the aiohttp transformation test in-process and cover repository row conversion The aiohttp chat transformation test no longer starts a server. It feeds the transformation a response whose json() returns the body under test. The proxy unit shards now exercise stored model rows whose params are JSON strings and the object permission create and update paths.
553 lines
19 KiB
Python
553 lines
19 KiB
Python
"""
|
|
Unit tests for DashScope image generation support (qwen-image-2.0, qwen-image-2.0-pro,
|
|
qwen-image-3.0, qwen-image-3.0-pro).
|
|
|
|
Run in docker: pytest tests/unit/test_dashscope_image_generation.py -v
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
import litellm
|
|
from litellm.llms.dashscope.image_generation.transformation import (
|
|
DashScopeImageGenerationConfig,
|
|
DEFAULT_API_BASE,
|
|
)
|
|
from litellm.types.utils import ImageResponse
|
|
from litellm.utils import get_llm_provider
|
|
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Provider detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_string",
|
|
[
|
|
"dashscope/qwen-image-2.0",
|
|
"dashscope/qwen-image-2.0-pro",
|
|
"dashscope/qwen-image-3.0",
|
|
"dashscope/qwen-image-3.0-pro",
|
|
],
|
|
)
|
|
def test_get_llm_provider_returns_dashscope(model_string: str):
|
|
model, provider, _, _ = get_llm_provider(model_string)
|
|
assert provider == "dashscope", f"Expected 'dashscope', got '{provider}'"
|
|
assert "qwen-image" in model
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Model info: mode == "image_generation"
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. Request transformation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDashScopeImageGenerationConfig:
|
|
def setup_method(self):
|
|
self.cfg = DashScopeImageGenerationConfig()
|
|
|
|
def test_get_complete_url_default(self):
|
|
url = self.cfg.get_complete_url(None, None, "qwen-image-2.0", {}, {})
|
|
assert url == DEFAULT_API_BASE
|
|
|
|
def test_get_complete_url_custom(self):
|
|
custom = "https://custom.endpoint/generate"
|
|
url = self.cfg.get_complete_url(custom, None, "qwen-image-2.0", {}, {})
|
|
assert url == custom
|
|
|
|
@pytest.mark.parametrize(
|
|
"chat_api_base",
|
|
[
|
|
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|
"https://dashscope-intl.aliyuncs.com/compatible-mode/v1/",
|
|
],
|
|
)
|
|
def test_get_complete_url_ignores_chat_compatible_mode_base(
|
|
self, chat_api_base: str
|
|
):
|
|
url = self.cfg.get_complete_url(chat_api_base, None, "qwen-image-3.0", {}, {})
|
|
assert url == DEFAULT_API_BASE
|
|
|
|
def test_validate_environment_sets_auth_header(self):
|
|
headers = self.cfg.validate_environment(
|
|
headers={},
|
|
model="qwen-image-2.0",
|
|
messages=[],
|
|
optional_params={},
|
|
litellm_params={},
|
|
api_key="sk-test-key",
|
|
)
|
|
assert headers["Authorization"] == "Bearer sk-test-key"
|
|
assert headers["Content-Type"] == "application/json"
|
|
|
|
def test_validate_environment_raises_without_key(self):
|
|
with patch(
|
|
"litellm.llms.dashscope.image_generation.transformation.get_secret_str",
|
|
return_value=None,
|
|
):
|
|
with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"):
|
|
self.cfg.validate_environment(
|
|
headers={},
|
|
model="qwen-image-2.0",
|
|
messages=[],
|
|
optional_params={},
|
|
litellm_params={},
|
|
api_key=None,
|
|
)
|
|
|
|
def test_transform_request_structure(self):
|
|
req = self.cfg.transform_image_generation_request(
|
|
model="qwen-image-2.0",
|
|
prompt="a puppy on green grass",
|
|
optional_params={"size": "1024*1024"},
|
|
litellm_params={},
|
|
headers={},
|
|
)
|
|
assert req["model"] == "qwen-image-2.0"
|
|
messages = req["input"]["messages"]
|
|
assert len(messages) == 1
|
|
assert messages[0]["role"] == "user"
|
|
assert messages[0]["content"][0]["text"] == "a puppy on green grass"
|
|
assert req["parameters"]["size"] == "1024*1024"
|
|
|
|
@pytest.mark.parametrize("model", ["qwen-image-3.0", "qwen-image-3.0-pro"])
|
|
def test_transform_request_qwen_image_3(self, model: str):
|
|
req = self.cfg.transform_image_generation_request(
|
|
model=model,
|
|
prompt="a poster with small multilingual text",
|
|
optional_params=self.cfg.map_openai_params(
|
|
non_default_params={"size": "2048x2048", "n": 6},
|
|
optional_params={},
|
|
model=model,
|
|
drop_params=False,
|
|
),
|
|
litellm_params={},
|
|
headers={},
|
|
)
|
|
assert req["model"] == model
|
|
assert req["input"]["messages"][0]["content"][0]["text"] == (
|
|
"a poster with small multilingual text"
|
|
)
|
|
assert req["parameters"]["size"] == "2048*2048"
|
|
assert req["parameters"]["n"] == 6
|
|
|
|
def test_transform_request_empty_params(self):
|
|
req = self.cfg.transform_image_generation_request(
|
|
model="qwen-image-2.0-pro",
|
|
prompt="sunset over the ocean",
|
|
optional_params={},
|
|
litellm_params={},
|
|
headers={},
|
|
)
|
|
assert req["parameters"] == {}
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 4. Response transformation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_mock_response(self, image_url: str) -> httpx.Response:
|
|
body = {
|
|
"status_code": 200,
|
|
"request_id": "test-request-id",
|
|
"output": {
|
|
"choices": [
|
|
{
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"image": image_url}],
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"usage": {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"image_count": 1,
|
|
},
|
|
}
|
|
mock_resp = MagicMock(spec=httpx.Response)
|
|
mock_resp.status_code = 200
|
|
mock_resp.headers = {}
|
|
mock_resp.json.return_value = body
|
|
return mock_resp
|
|
|
|
def test_transform_response_extracts_url(self):
|
|
image_url = "https://example.oss.aliyuncs.com/generated/test.png"
|
|
mock_resp = self._make_mock_response(image_url)
|
|
model_response = ImageResponse()
|
|
result = self.cfg.transform_image_generation_response(
|
|
model="qwen-image-2.0",
|
|
raw_response=mock_resp,
|
|
model_response=model_response,
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
assert result.data is not None
|
|
assert len(result.data) == 1
|
|
assert result.data[0].url == image_url
|
|
|
|
def test_transform_response_multiple_images(self):
|
|
body = {
|
|
"output": {
|
|
"choices": [
|
|
{
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"image": "https://example.com/img1.png"}],
|
|
},
|
|
},
|
|
{
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"image": "https://example.com/img2.png"}],
|
|
},
|
|
},
|
|
]
|
|
},
|
|
"usage": {},
|
|
}
|
|
mock_resp = MagicMock(spec=httpx.Response)
|
|
mock_resp.status_code = 200
|
|
mock_resp.headers = {}
|
|
mock_resp.json.return_value = body
|
|
|
|
model_response = ImageResponse()
|
|
result = self.cfg.transform_image_generation_response(
|
|
model="qwen-image-2.0",
|
|
raw_response=mock_resp,
|
|
model_response=model_response,
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
assert len(result.data) == 2
|
|
assert result.data[0].url == "https://example.com/img1.png"
|
|
assert result.data[1].url == "https://example.com/img2.png"
|
|
|
|
def test_transform_response_multiple_images_in_one_choice(self):
|
|
body = {
|
|
"output": {
|
|
"choices": [
|
|
{
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{"image": "https://example.com/img1.png", "type": "image"},
|
|
{"image": "https://example.com/img2.png", "type": "image"},
|
|
],
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"usage": {
|
|
"output_width": 1024,
|
|
"output_height": 1024,
|
|
"output_image_count": 2,
|
|
},
|
|
}
|
|
mock_resp = MagicMock(spec=httpx.Response)
|
|
mock_resp.status_code = 200
|
|
mock_resp.headers = {}
|
|
mock_resp.json.return_value = body
|
|
|
|
result = self.cfg.transform_image_generation_response(
|
|
model="qwen-image-3.0",
|
|
raw_response=mock_resp,
|
|
model_response=ImageResponse(),
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
assert [image.url for image in result.data] == [
|
|
"https://example.com/img1.png",
|
|
"https://example.com/img2.png",
|
|
]
|
|
|
|
def test_transform_response_raises_on_non_200_status(self):
|
|
mock_resp = MagicMock(spec=httpx.Response)
|
|
mock_resp.status_code = 400
|
|
mock_resp.headers = {}
|
|
mock_resp.text = '{"code":"InvalidParameter","message":"Size not supported"}'
|
|
mock_resp.json.return_value = {
|
|
"code": "InvalidParameter",
|
|
"message": "Size not supported",
|
|
}
|
|
|
|
with pytest.raises(BaseLLMException):
|
|
self.cfg.transform_image_generation_response(
|
|
model="qwen-image-2.0",
|
|
raw_response=mock_resp,
|
|
model_response=ImageResponse(),
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
|
|
def test_transform_response_raises_on_api_error_body(self):
|
|
mock_resp = MagicMock(spec=httpx.Response)
|
|
mock_resp.status_code = 200
|
|
mock_resp.headers = {}
|
|
mock_resp.json.return_value = {
|
|
"code": "InvalidParameter",
|
|
"message": "Size not supported",
|
|
}
|
|
|
|
with pytest.raises(BaseLLMException):
|
|
self.cfg.transform_image_generation_response(
|
|
model="qwen-image-2.0",
|
|
raw_response=mock_resp,
|
|
model_response=ImageResponse(),
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 5. OpenAI → DashScope parameter mapping
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_map_openai_params_size_conversion(self):
|
|
mapped = self.cfg.map_openai_params(
|
|
non_default_params={"size": "1024x1024"},
|
|
optional_params={},
|
|
model="qwen-image-2.0",
|
|
drop_params=False,
|
|
)
|
|
assert mapped["size"] == "1024*1024"
|
|
|
|
def test_map_openai_params_n_passthrough(self):
|
|
mapped = self.cfg.map_openai_params(
|
|
non_default_params={"n": 2},
|
|
optional_params={},
|
|
model="qwen-image-2.0",
|
|
drop_params=False,
|
|
)
|
|
assert mapped == {"n": 2}
|
|
|
|
def test_map_openai_params_unknown_size_uses_asterisk(self):
|
|
mapped = self.cfg.map_openai_params(
|
|
non_default_params={"size": "768x768"},
|
|
optional_params={},
|
|
model="qwen-image-2.0",
|
|
drop_params=False,
|
|
)
|
|
assert mapped["size"] == "768*768"
|
|
|
|
@pytest.mark.parametrize(
|
|
"openai_size, expected",
|
|
[
|
|
("256x256", "256*256"),
|
|
("512x512", "512*512"),
|
|
("1024x1024", "1024*1024"),
|
|
("1792x1024", "1792*1024"),
|
|
("1024x1792", "1024*1792"),
|
|
("2048x2048", "2048*2048"),
|
|
],
|
|
)
|
|
def test_map_openai_params_size_table(self, openai_size: str, expected: str):
|
|
mapped = self.cfg.map_openai_params(
|
|
non_default_params={"size": openai_size},
|
|
optional_params={},
|
|
model="qwen-image-2.0",
|
|
drop_params=False,
|
|
)
|
|
assert mapped["size"] == expected
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 6. End-to-end flow via litellm.image_generation (HTTP mocked)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model",
|
|
[
|
|
"dashscope/qwen-image-2.0",
|
|
"dashscope/qwen-image-3.0",
|
|
"dashscope/qwen-image-3.0-pro",
|
|
],
|
|
)
|
|
def test_litellm_image_generation_dashscope_end_to_end(model: str):
|
|
mock_response_body = {
|
|
"output": {
|
|
"choices": [
|
|
{
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"image": "https://dashscope-result.oss.aliyuncs.com/test.png"
|
|
}
|
|
],
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"usage": {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"image_count": 1,
|
|
},
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
|
|
) as mock_post:
|
|
mock_http_response = MagicMock()
|
|
mock_http_response.json.return_value = mock_response_body
|
|
mock_http_response.status_code = 200
|
|
mock_http_response.headers = {}
|
|
mock_post.return_value = mock_http_response
|
|
|
|
response = litellm.image_generation(
|
|
model=model,
|
|
prompt="a puppy playing on green grass",
|
|
api_key="sk-test-key",
|
|
size="1024x1024",
|
|
)
|
|
|
|
assert response is not None
|
|
assert response.data is not None
|
|
assert len(response.data) == 1
|
|
assert (
|
|
response.data[0].url == "https://dashscope-result.oss.aliyuncs.com/test.png"
|
|
)
|
|
|
|
# Verify the HTTP call was made to the DashScope endpoint
|
|
call_args = mock_post.call_args
|
|
called_url = (
|
|
call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
|
|
)
|
|
assert called_url == DEFAULT_API_BASE
|
|
|
|
# Verify request body contains DashScope format
|
|
call_kwargs = call_args[1] if call_args[1] else {}
|
|
if "json" in call_kwargs:
|
|
body = call_kwargs["json"]
|
|
assert "input" in body
|
|
assert "messages" in body["input"]
|
|
assert body["parameters"]["size"] == "1024*1024"
|
|
|
|
|
|
def _transform_response(payload: object) -> ImageResponse:
|
|
return DashScopeImageGenerationConfig().transform_image_generation_response(
|
|
model="qwen-image-2.0",
|
|
raw_response=httpx.Response(200, json=payload),
|
|
model_response=ImageResponse(),
|
|
logging_obj=MagicMock(),
|
|
request_data={},
|
|
optional_params={},
|
|
litellm_params={},
|
|
encoding=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
{},
|
|
{"output": {}},
|
|
{"output": {"choices": []}},
|
|
{"output": {"choices": ""}},
|
|
{"output": {"choices": [{}]}},
|
|
{"output": {"choices": [{"message": {}}]}},
|
|
{"output": {"choices": [{"message": {"content": {}}}]}},
|
|
{"output": {"choices": [{"message": {"content": [{}]}}]}},
|
|
{"output": {"choices": [{"message": {"content": [{"image": ""}]}}]}},
|
|
{"output": {"choices": [{"message": {"content": [{"text": "hi"}]}}]}},
|
|
{"code": "Partial", "output": {"choices": []}},
|
|
],
|
|
)
|
|
def test_transform_response_without_image_content_has_no_images(
|
|
payload: dict[str, object],
|
|
):
|
|
assert _transform_response(payload).data == []
|
|
|
|
|
|
def test_transform_response_skips_content_items_without_an_image():
|
|
response = _transform_response(
|
|
{
|
|
"output": {
|
|
"choices": [
|
|
{"message": {"content": [{"text": "caption"}, {"image": "https://a.example/1.png"}]}},
|
|
{"message": {"content": [{"image": None}, {"image": "https://a.example/2.png"}]}},
|
|
]
|
|
}
|
|
}
|
|
)
|
|
|
|
assert [image.url for image in response.data] == [
|
|
"https://a.example/1.png",
|
|
"https://a.example/2.png",
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("payload", "message"),
|
|
[
|
|
({"code": "InvalidParameter", "message": "Size not supported"}, "Size not supported"),
|
|
({"code": "Throttled"}, "{'code': 'Throttled'}"),
|
|
({"code": 429, "message": {"detail": "slow down"}}, "{'detail': 'slow down'}"),
|
|
],
|
|
)
|
|
def test_transform_response_reports_api_error_bodies(
|
|
payload: dict[str, object], message: str
|
|
):
|
|
with pytest.raises(BaseLLMException) as exc_info:
|
|
_transform_response(payload)
|
|
|
|
assert exc_info.value.message == message
|
|
assert exc_info.value.status_code == 200
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
["code"],
|
|
"code",
|
|
7,
|
|
{"output": None},
|
|
{"output": ["not", "an", "object"]},
|
|
{"output": {"choices": 7}},
|
|
{"output": {"choices": ["not an object"]}},
|
|
{"output": {"choices": [{"message": "not an object"}]}},
|
|
{"output": {"choices": [{"message": {"content": None}}]}},
|
|
{"output": {"choices": [{"message": {"content": ["not an object"]}}]}},
|
|
],
|
|
)
|
|
def test_transform_response_rejects_malformed_payloads_without_echoing_them(
|
|
payload: object,
|
|
):
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
_transform_response(payload)
|
|
|
|
assert "input_value" not in str(exc_info.value)
|