diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 9ada3674d33..0ffd603deac 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -432,6 +432,25 @@ def _safe_read_response(response: httpx.Response, timeout: float | None = None) return b"" +def _build_masked_request(original_request: httpx.Request, masked_url: str) -> httpx.Request: + """Build a masked request without forcing unread streaming bodies into memory.""" + headers = httpx.Headers(original_request.headers) + try: + request_content = original_request.content + except httpx.RequestNotRead: + request_content = b"" + headers.pop("content-length", None) + headers.pop("transfer-encoding", None) + headers.pop("content-type", None) + + return httpx.Request( + method=original_request.method, + url=masked_url, + headers=headers, + content=request_content, + ) + + def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None: """Raise a MaskedHTTPStatusError for sync HTTP handlers.""" if stream: @@ -494,17 +513,7 @@ class MaskedHTTPStatusError(httpx.HTTPStatusError): if k.lower() not in ("content-encoding", "content-length") } - try: - request_content = original_error.request.content - except httpx.RequestNotRead: - request_content = b"" - - masked_request: Final = httpx.Request( - method=original_error.request.method, - url=masked_url, - headers=original_error.request.headers, - content=request_content, - ) + masked_request: Final = _build_masked_request(original_error.request, masked_url) super().__init__( message=masked_original_message, diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py index 0a3bf403bf8..29d30e78bdb 100644 --- a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py +++ b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py @@ -9,6 +9,7 @@ Covers: import os import sys +from io import BytesIO from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -160,6 +161,31 @@ class TestMaskedHTTPStatusError: # Content-Encoding must have been stripped from the rebuilt headers. assert "content-encoding" not in {k.lower() for k in masked.response.headers} + def test_handles_unread_multipart_request_content(self): + """Multipart file requests can have unread streaming request bodies.""" + request = httpx.Request( + "POST", + "https://api.example.com/v1/images/edits?key=SECRET", + data={"model": "gpt-image-1"}, + files=[ + ("image[]", ("image.png", BytesIO(b"image"), "image/png")), + ("mask", ("mask.png", BytesIO(b"mask"), "image/png")), + ], + ) + response = httpx.Response( + status_code=500, + content=b'{"error": "upstream error"}', + request=request, + ) + orig = httpx.HTTPStatusError("500", request=request, response=response) + + masked = MaskedHTTPStatusError(orig) + + assert "SECRET" not in str(masked.request.url) + assert masked.request.content == b"" + assert masked.request.headers.get("content-length") == "0" + assert "content-type" not in masked.request.headers + class TestSafeResponseHelpers: def test_safe_get_response_text_normal(self): diff --git a/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py b/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py index da3352a4943..279eb0d4258 100644 --- a/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py +++ b/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py @@ -1,9 +1,7 @@ -from io import BufferedReader, BytesIO -from typing import Dict +from io import BytesIO import pytest -from litellm import image_edit from litellm.images.utils import ImageEditRequestUtils from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.types.router import GenericLiteLLMParams @@ -266,6 +264,40 @@ def test_transform_image_edit_request_with_mask_list( assert mask_file[1][1] == mask1 # Should be the first mask, not the second +def test_transform_image_edit_request_with_bytesio_mask_list( + image_edit_config: OpenAIImageEditConfig, +): + """Test transformation with proxy-style BytesIO mask list.""" + model = "gpt-image-1" + prompt = "Make the background blue" + image = BytesIO(b"fake_image_data") + image.name = "image.png" + mask = BytesIO(b"fake_mask_data") + mask.name = "mask.png" + image_edit_optional_request_params = {"mask": [mask]} + litellm_params = GenericLiteLLMParams() + headers = {} + + data, files = image_edit_config.transform_image_edit_request( + model=model, + prompt=prompt, + image=[image], + image_edit_optional_request_params=image_edit_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + + assert data["model"] == model + assert data["prompt"] == prompt + assert "image" not in data + assert "mask" not in data + image_file = next(f for f in files if f[0] == "image[]") + mask_file = next(f for f in files if f[0] == "mask") + assert image_file[1][1] is image + assert mask_file[1][0] == "mask.png" + assert mask_file[1][1] is mask + + def test_transform_image_edit_request_with_input_fidelity( image_edit_config: OpenAIImageEditConfig, ): diff --git a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py index f5410ef0d70..20ba3db0802 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py +++ b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py @@ -145,3 +145,34 @@ def test_azure_image_generation_route_allows_ordinary_path_model(client_no_auth) assert response.status_code == 200 mock_aimage_generation.assert_called_once() + + +def test_image_edit_route_with_mask(client_no_auth): + client, _, mock_aimage_edit = client_no_auth + image_path = os.path.join( + os.path.dirname(__file__), + "../../../image_gen_tests/test_image.png", + ) + with open(image_path, "rb") as image_file, open(image_path, "rb") as mask_file: + files = { + "image": ("test_image.png", image_file, "image/png"), + "mask": ("mask.png", mask_file, "image/png"), + } + data = { + "model": "dall-e-3", + "prompt": "A cute baby sea otter", + "response_format": "b64_json", + "size": "1024x1024", + } + response = client.post("/v1/images/edits", files=files, data=data) + + mock_aimage_edit.assert_called_once() + called_kwargs = mock_aimage_edit.call_args.kwargs + assert "dall-e-3" in called_kwargs["model"] + assert called_kwargs["prompt"] == "A cute baby sea otter" + assert called_kwargs["response_format"] == "b64_json" + assert called_kwargs["size"] == "1024x1024" + assert called_kwargs["image"][0].name == "test_image.png" + assert called_kwargs["mask"][0].name == "mask.png" + assert response.status_code == 200 + assert response.json()["data"]