From 793371270b7c36cc2f1650e658e7f7c3b3690a62 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 24 Sep 2026 04:02:34 +0000 Subject: [PATCH] fix(images): carry deployment pricing kwargs to logging and harden dimension parsers - forward CustomPricingLiteLLMParams kwargs through the image_generation and aimage_edit litellm_params dicts so deployment-scoped rates reach the spend ledger (council H2: litellm_params-slot rates were dropped) - require PNG IHDR, WebP VP8 start code, VP8L signature, and known BMP DIB sizes so signature-matching malformed files fall back instead of fabricating dimensions (council H3) - make read_image_dimensions and total_reference_pixels never raise so odd streams and short tuples degrade to the generated-pixel fallback --- litellm/images/dimensions.py | 38 +++++++++++--- litellm/images/main.py | 3 ++ tests/test_litellm/images/test_dimensions.py | 41 ++++++++++++++++ .../images/test_image_edit_utils.py | 49 +++++++++++++++++++ .../test_image_generation_extra_headers.py | 35 ++++++++++++- 5 files changed, 157 insertions(+), 9 deletions(-) diff --git a/litellm/images/dimensions.py b/litellm/images/dimensions.py index f9cc2a27c93..849b11535f9 100644 --- a/litellm/images/dimensions.py +++ b/litellm/images/dimensions.py @@ -18,6 +18,12 @@ _JPEG_FIRST_SEGMENT_OFFSET: Final = 2 _JPEG_SOF_PAYLOAD_SIZE: Final = 5 _JPEG_FILL_CHUNK: Final = 64 * 1024 +_PNG_IHDR: Final = b"IHDR" +_WEBP_VP8_START_CODE: Final = b"\x9d\x01\x2a" +_WEBP_VP8L_SIGNATURE: Final = 0x2F +_BMP_CORE_DIB_SIZE: Final = 12 +_BMP_KNOWN_DIB_SIZES: Final = frozenset({12, 40, 52, 56, 64, 108, 124}) + @dataclass(frozen=True, slots=True) class ImageDimensions: @@ -30,7 +36,11 @@ class ImageDimensions: def total_reference_pixels(images: Sequence[FileTypes]) -> int | None: - measured: Final = tuple(read_image_dimensions(image) for image in images) + """All-or-nothing sum of measured reference pixels; never raises.""" + try: + measured: Final = tuple(read_image_dimensions(image) for image in images) + except Exception: # noqa: BLE001 # a billing helper must fall back, never fail the request + return None dimensions: Final = tuple(size for size in measured if size is not None) if len(dimensions) != len(measured): verbose_logger.debug( @@ -44,10 +54,13 @@ def total_reference_pixels(images: Sequence[FileTypes]) -> int | None: def read_image_dimensions(image: FileTypes) -> ImageDimensions | None: try: - content = image[1] if isinstance(image, tuple) else image - if isinstance(content, (str, os.PathLike)): + content: Final = cast( + "IO[bytes] | bytes | str | os.PathLike[str] | None", + image[1] if isinstance(image, tuple) else image, + ) + if content is None or isinstance(content, (str, os.PathLike)): return None - stream = BytesIO(content) if isinstance(content, bytes) else content + stream = BytesIO(content) if isinstance(content, (bytes, bytearray, memoryview)) else content if not stream.seekable(): return None position = stream.tell() @@ -58,13 +71,15 @@ def read_image_dimensions(image: FileTypes) -> ImageDimensions | None: if dimensions is None or dimensions.width <= 0 or dimensions.height <= 0: return None return dimensions - except (OSError, ValueError, AttributeError, struct.error): + except Exception: # noqa: BLE001 # an odd stream or malformed file must fall back, never raise return None def _header_dimensions(stream: IO[bytes], position: int) -> ImageDimensions | None: stream.seek(position) - head: Final = stream.read(_HEADER_READ_SIZE) + head: Final = cast("bytes | None", stream.read(_HEADER_READ_SIZE)) + if head is None: + return None match get_image_type(head): # pyright: ignore[reportMatchNotExhaustive] # unmatched types fall through to the None below case "png": return _png_dimensions(head) @@ -80,7 +95,7 @@ def _header_dimensions(stream: IO[bytes], position: int) -> ImageDimensions | No def _png_dimensions(head: bytes) -> ImageDimensions | None: - if len(head) < 24: + if len(head) < 24 or head[12:16] != _PNG_IHDR: return None width, height = cast(tuple[int, int], struct.unpack(">II", head[16:24])) return ImageDimensions(width=width, height=height) @@ -97,7 +112,10 @@ def _gif_dimensions(head: bytes) -> ImageDimensions | None: def _bmp_dimensions(head: bytes) -> ImageDimensions | None: if len(head) < 26: return None - if int.from_bytes(head[14:18], "little") == 12: + dib_size: Final = int.from_bytes(head[14:18], "little") + if dib_size not in _BMP_KNOWN_DIB_SIZES: + return None + if dib_size == _BMP_CORE_DIB_SIZE: core_width: Final = int.from_bytes(head[18:20], "little") core_height: Final = int.from_bytes(head[20:22], "little") return ImageDimensions(width=core_width, height=core_height) @@ -116,9 +134,13 @@ def _webp_dimensions(head: bytes) -> ImageDimensions | None: height=int.from_bytes(head[27:30], "little") + 1, ) case b"VP8 ": + if head[23:26] != _WEBP_VP8_START_CODE: + return None width, height = cast(tuple[int, int], struct.unpack("> 14) & 0x3FFF) + 1) return None diff --git a/litellm/images/main.py b/litellm/images/main.py index 1f722eb752a..df930c82e81 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -51,6 +51,7 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( LITELLM_IMAGE_VARIATION_PROVIDERS, + CustomPricingLiteLLMParams, LlmProviders, all_litellm_params, ) @@ -292,6 +293,7 @@ def image_generation( "model_info": model_info, "preset_cache_key": None, "stream_response": {}, + **{k: v for k, v in kwargs.items() if k in CustomPricingLiteLLMParams.model_fields}, }, custom_llm_provider=custom_llm_provider, ) @@ -885,6 +887,7 @@ def image_edit( **image_edit_request_params, "litellm_call_id": litellm_call_id, "model_info": model_info, + **{k: v for k, v in kwargs.items() if k in CustomPricingLiteLLMParams.model_fields}, }, custom_llm_provider=custom_llm_provider, ) diff --git a/tests/test_litellm/images/test_dimensions.py b/tests/test_litellm/images/test_dimensions.py index 08b3b970687..2ab6d4930e2 100644 --- a/tests/test_litellm/images/test_dimensions.py +++ b/tests/test_litellm/images/test_dimensions.py @@ -129,7 +129,24 @@ def _dims(width: int, height: int) -> ImageDimensions: pytest.param(_webp_chunk(b"VP8X", VP8X_BODY), _dims(700, 350), id="webp-vp8x-extended"), pytest.param(_webp_chunk(b"VP8Z", bytes(10)), None, id="webp-unknown-fourcc"), pytest.param(_webp_chunk(b"VP8X", VP8X_BODY)[:29], None, id="webp-shorter-than-30-bytes"), + pytest.param( + _webp_chunk(b"VP8 ", bytes(3) + bytes(3) + struct.pack("II", 4000, 3000) + bytes(8)) + + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", 1024, 1024, 8, 0, 0, 0, 0)), + None, + id="png-cgbi-first-chunk", + ), pytest.param(_gif(b"GIF89a", 1024, 768), _dims(1024, 768), id="gif89a"), pytest.param(_gif(b"GIF87a", 320, 200), _dims(320, 200), id="gif87a"), pytest.param(b"GIF89a" + bytes(2), None, id="gif-shorter-than-10-bytes"), @@ -138,6 +155,11 @@ def _dims(width: int, height: int) -> ImageDimensions: pytest.param(_bmp(-1024, 768), None, id="bmp-negative-width"), pytest.param(_bmp(0, 768), None, id="bmp-zero-width"), pytest.param(b"BM" + bytes(18), None, id="bmp-shorter-than-26-bytes"), + pytest.param( + b"BM" + bytes(12) + struct.pack(" bytes | None: + return None + + +def test_read_image_dimensions_returns_none_when_read_returns_none(): + assert read_image_dimensions(cast(Any, _NonBlockingStream(_png(64, 64)))) is None + + +def test_read_image_dimensions_returns_none_for_short_and_none_tuples(): + assert read_image_dimensions(("ref.png",)) is None + assert read_image_dimensions(("ref.png", None)) is None + + +def test_total_reference_pixels_returns_none_instead_of_raising(): + assert total_reference_pixels([("ref.png",)]) is None + assert total_reference_pixels(cast(Any, 123)) is None + + def test_total_reference_pixels_returns_none_for_closed_and_read_only_streams(): closed: Final = io.BytesIO(_png(64, 64)) closed.close() diff --git a/tests/test_litellm/images/test_image_edit_utils.py b/tests/test_litellm/images/test_image_edit_utils.py index 2146c1fab01..e9d50bcb7fe 100644 --- a/tests/test_litellm/images/test_image_edit_utils.py +++ b/tests/test_litellm/images/test_image_edit_utils.py @@ -268,6 +268,55 @@ class TestImageEditCustomPricing: litellm_params = {"litellm_call_id": "test-call-id"} assert use_custom_pricing_for_model(litellm_params) is False + def test_image_edit_forwards_custom_pricing_kwargs_to_logging(self): + """Pricing a deployment declares under litellm_params (e.g. input_cost_per_pixel) + reaches self.litellm_params so use_custom_pricing_for_model fires on image edits.""" + from litellm.images.main import image_edit + + captured_litellm_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + original_update = mock_logging_obj.update_from_kwargs + + def capturing_update(**update_kwargs): + captured_litellm_params.update(update_kwargs.get("litellm_params", {})) + return original_update(**update_kwargs) + + mock_logging_obj.update_from_kwargs = capturing_update + + with ( + patch( + "litellm.images.main.get_llm_provider", + return_value=("test-model", "openai", None, None), + ), + patch( + "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", + return_value=MagicMock(), + ), + patch( + "litellm.images.main._get_ImageEditRequestUtils", + return_value=MagicMock( + get_requested_image_edit_optional_param=MagicMock(return_value={}), + get_optional_params_image_edit=MagicMock(return_value={}), + ), + ), + patch("litellm.images.main.base_llm_http_handler") as mock_handler, + ): + mock_handler.image_edit_handler.return_value = MagicMock() + + image_edit( + image=b"fake-image-data", + prompt="test prompt", + model="openai/test-model", + litellm_logging_obj=mock_logging_obj, + input_cost_per_pixel=4.76837158203125e-08, + ) + + assert captured_litellm_params["input_cost_per_pixel"] == pytest.approx(4.76837158203125e-08) + assert use_custom_pricing_for_model(captured_litellm_params) is True + class TestImageEditHandlerCredentialsForwarding: """ diff --git a/tests/test_litellm/images/test_image_generation_extra_headers.py b/tests/test_litellm/images/test_image_generation_extra_headers.py index a65bdeb892b..2f893bffcf3 100644 --- a/tests/test_litellm/images/test_image_generation_extra_headers.py +++ b/tests/test_litellm/images/test_image_generation_extra_headers.py @@ -10,9 +10,9 @@ from unittest.mock import MagicMock, patch import pytest - import litellm from litellm.images.main import image_generation +from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model class TestImageGenerationExtraHeaders: @@ -73,3 +73,36 @@ class TestImageGenerationExtraHeaders: ) assert "extra_headers" not in optional_params + + @patch("litellm.images.main.openai_chat_completions") + def test_image_generation_forwards_custom_pricing_kwargs_to_logging( + self, mock_openai_chat_completions + ): + """Pricing a deployment declares under litellm_params (e.g. input_cost_per_pixel) + reaches self.litellm_params so use_custom_pricing_for_model fires on generations.""" + mock_openai_chat_completions.image_generation.return_value = litellm.utils.ImageResponse( + created=1234567890, + data=[{"url": "https://example.com/image.png"}], + ) + captured_litellm_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + original_update = mock_logging_obj.update_from_kwargs + + def capturing_update(**update_kwargs): + captured_litellm_params.update(update_kwargs.get("litellm_params", {})) + return original_update(**update_kwargs) + + mock_logging_obj.update_from_kwargs = capturing_update + + image_generation( + model="openai/dall-e-3", + prompt="A red circle", + litellm_logging_obj=mock_logging_obj, + input_cost_per_pixel=4.76837158203125e-08, + ) + + assert captured_litellm_params["input_cost_per_pixel"] == pytest.approx(4.76837158203125e-08) + assert use_custom_pricing_for_model(captured_litellm_params) is True