Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions litellm/proxy/image_endpoints/endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,10 @@

IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams))

IMAGE_ARRAY_FIELD: Final = "image[]"
MASK_ARRAY_FIELD: Final = "mask[]"
BRACKETED_FILE_FIELDS: Final = frozenset({IMAGE_ARRAY_FIELD, MASK_ARRAY_FIELD})


async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO:
"""
Expand Down Expand Up @@ -241,9 +245,9 @@ async def image_edit_api(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
image: list[UploadFile] | None = File(None),
image_array: list[UploadFile] | None = File(None, alias="image[]"),
image_array: list[UploadFile] | None = File(None, alias=IMAGE_ARRAY_FIELD),
mask: list[UploadFile] | None = File(None),
mask_array: list[UploadFile] | None = File(None, alias="mask[]"),
mask_array: list[UploadFile] | None = File(None, alias=MASK_ARRAY_FIELD),
model: str | None = None,
):
"""
Expand Down Expand Up @@ -291,12 +295,14 @@ async def image_edit_api(
#########################################################
# Read request body and convert UploadFiles to BytesIO
#########################################################
data: Final = dict(
coerce_numeric_form_fields(
data: Final = {
key: value
for key, value in coerce_numeric_form_fields(
parsed_body=await _read_request_body(request=request),
numeric_fields=IMAGE_EDIT_NUMERIC_FORM_FIELDS,
)
)
).items()
if key not in BRACKETED_FILE_FIELDS
}
image_files: Final = await batch_to_bytesio(image)
mask_files: Final = await batch_to_bytesio(mask)
if image_files:
Expand Down
70 changes: 62 additions & 8 deletions tests/test_litellm/proxy/image_endpoints/test_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,18 +92,14 @@ async def receive():
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", fake_proxy_logger
)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_proxy_logger)
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
monkeypatch.setattr(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers",
classmethod(lambda *args, **kwargs: {}),
)
monkeypatch.setattr(
"litellm.proxy.image_endpoints.endpoints.route_request", fake_route_request
)
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", fake_route_request)

result = await endpoints.image_generation(
request=request,
Expand Down Expand Up @@ -138,6 +134,60 @@ async def base_process_llm_request(self, **_: Any) -> Dict[str, Any]:
return TestClient(app)


def test_image_edit_image_array_alias_is_not_forwarded(monkeypatch: pytest.MonkeyPatch):
"""The documented `image[]` alias must reach the provider only as `image`."""
captured: Dict[str, Any] = {}

response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
files={"image[]": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png")},
data={"model": "gpt-image-1", "prompt": "add a hat"},
)

assert response.status_code == 200
assert "image[]" not in captured
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]
assert [buffer.name for buffer in captured["image"]] == ["tree.png"]


def test_image_edit_mask_array_alias_is_not_forwarded(monkeypatch: pytest.MonkeyPatch):
"""`mask[]` has the same shape as `image[]` and must be dropped the same way."""
captured: Dict[str, Any] = {}

response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
files={
"image": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png"),
"mask[]": ("mask.png", b"\x89PNG\r\n\x1a\nmask", "image/png"),
},
data={"model": "gpt-image-1", "prompt": "add a hat"},
)

assert response.status_code == 200
assert "mask[]" not in captured
assert [buffer.getvalue() for buffer in captured["mask"]] == [b"\x89PNG\r\n\x1a\nmask"]
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]


def test_image_edit_canonical_file_fields_still_reach_the_provider(monkeypatch: pytest.MonkeyPatch):
"""Dropping the bracketed aliases must not touch the canonical fields."""
captured: Dict[str, Any] = {}

response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
files={
"image": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png"),
"mask": ("mask.png", b"\x89PNG\r\n\x1a\nmask", "image/png"),
},
data={"model": "gpt-image-1", "prompt": "add a hat"},
)

assert response.status_code == 200
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]
assert [buffer.getvalue() for buffer in captured["mask"]] == [b"\x89PNG\r\n\x1a\nmask"]
assert captured["prompt"] == "add a hat"


def test_image_edit_multipart_n_reaches_the_provider_as_an_int(monkeypatch):
"""A multipart `n` must not arrive as the string Starlette parsed it into."""
captured: Dict[str, Any] = {}
Expand Down Expand Up @@ -177,7 +227,9 @@ async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(mon
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
return kwargs["data"]

async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]:
async def fake_pre_call_hook(
*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str
) -> dict[str, object]:
return data

async def fake_post_call_failure_hook(**_: object) -> None:
Expand Down Expand Up @@ -208,6 +260,8 @@ async def receive() -> dict[str, object]:
request = Request({"type": "http", "method": "POST", "path": "/v1/images/generations", "headers": []}, receive)

with pytest.raises(ProxyException) as raised:
await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth())
await endpoints.image_generation(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)

assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404")
Loading