From eedf4d86481bcbbeae183be1b8f37f6f359c1f4e Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 24 Sep 2026 00:23:13 +0000 Subject: [PATCH] refactor(images): expose reference pixels through a typed ImageResponse accessor --- litellm/images/dimensions.py | 2 +- .../image_generation/cost_calculator.py | 28 +++++++------------ litellm/llms/custom_httpx/llm_http_handler.py | 3 +- litellm/types/utils.py | 7 +++++ .../image_generation/test_cost_calculator.py | 27 ++++++++++++++++-- .../custom_httpx/test_llm_http_handler.py | 7 +++-- 6 files changed, 47 insertions(+), 27 deletions(-) diff --git a/litellm/images/dimensions.py b/litellm/images/dimensions.py index 979e777ae81..a54cc7b1947 100644 --- a/litellm/images/dimensions.py +++ b/litellm/images/dimensions.py @@ -105,7 +105,7 @@ def _jpeg_sof_dimensions(stream: IO[bytes], start: int, offset: int, segments_le sof = stream.read(_JPEG_SOF_PAYLOAD_SIZE) if len(sof) < _JPEG_SOF_PAYLOAD_SIZE: return None - _precision, height, width = cast(tuple[int, int, int], struct.unpack(">BHH", sof)) + height, width = cast(tuple[int, int], struct.unpack(">HH", sof[1:5])) return ImageDimensions(width=width, height=height) segment_length = int.from_bytes(marker[2:4], "big") if segment_length < 2: diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index a0ddbd85618..c65232aa196 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping from typing import Any, Final @@ -14,23 +15,16 @@ def _rate(table: ModelInfo, key: str) -> float | None: return _get_cost_per_unit(table, key, default_value=None) +_SIZE_PATTERN: Final = re.compile(r"(\d+)(?:x|-x-)(\d+)") + + def _size_pixels(size: str | None) -> int: if size is None: return 0 - for separator in ("x", "-x-"): - if separator in size: - parts = size.split(separator) - if len(parts) != 2: - continue - try: - width = int(parts[0]) - height = int(parts[1]) - except ValueError: - continue - if width > 0 and height > 0: - return width * height - continue - return 0 + match: Final = _SIZE_PATTERN.fullmatch(size) + if match is None or int(match[1]) <= 0 or int(match[2]) <= 0: + return 0 + return int(match[1]) * int(match[2]) def _generated_pixels( @@ -73,8 +67,7 @@ def cost_calculator( if usage_cost is not None: return usage_cost - data: Final = image_response.data - num_images: Final = n if n is not None else (len(data) if isinstance(data, list) else 0) + num_images: Final = n if n is not None else len(image_response.data or ()) generated_meters: Final = ( ("output_cost_per_image", num_images), ("input_cost_per_pixel", _generated_pixels(optional_params, size, image_response) * num_images), @@ -83,8 +76,7 @@ def cost_calculator( (rate * units for key, units in generated_meters if (rate := _rate(resolved, key)) is not None), 0.0, ) - reference_pixels: Final = getattr(image_response, "_reference_pixels", None) reference: Final = (_rate(resolved, "input_cost_per_reference_pixel") or 0.0) * ( - reference_pixels if isinstance(reference_pixels, int) and not isinstance(reference_pixels, bool) else 0 + image_response.reference_pixels or 0 ) return generated + reference diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a5e36ef50d7..f7615396da3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -385,8 +385,7 @@ def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: def _with_reference_pixels(response: ImageResponse, reference_pixels: int | None) -> ImageResponse: if reference_pixels is not None: - # mutable-ok: handler stamps this billing field - response._reference_pixels = reference_pixels # pyright: ignore[reportPrivateUsage] # billing stamp + response.set_reference_pixels(reference_pixels) return response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c31b7ef886a..456c3b9de9a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2586,6 +2586,13 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} _reference_pixels: int | None = None + @property + def reference_pixels(self) -> int | None: + return self._reference_pixels + + def set_reference_pixels(self, pixels: int) -> None: + self._reference_pixels = pixels + usage: ImageUsage | None = None """ Users might use litellm with older python versions, we don't want this to break for them. diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py b/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py index 7f79e518bf3..608a04959fc 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py @@ -24,16 +24,17 @@ def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): def _edit_response(reference_pixels: int | None = REFERENCE_PIXELS, **kwargs: object) -> ImageResponse: response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n")], size="1024x1024", **kwargs) - response._reference_pixels = reference_pixels + if reference_pixels is not None: + response.set_reference_pixels(reference_pixels) return response -def _edit_cost(response: ImageResponse, **kwargs: object) -> float: +def _edit_cost(response: ImageResponse, size: str = "1024x1024", **kwargs: object) -> float: return CostCalculatorUtils.route_image_generation_cost_calculator( model="FLUX.2-flex", completion_response=response, custom_llm_provider="azure_ai", - size="1024x1024", + size=size, call_type="image_edit", **kwargs, ) @@ -136,6 +137,26 @@ def test_edit_without_measurement_bills_generated_pixels_only() -> None: assert cost == pytest.approx(catalog_rate * 1024 * 1024) +@pytest.mark.parametrize( + ("size", "expected_generated"), + [ + ("1024-x-1024", "computed"), + ("auto", "reference-only"), + ("garbage", "reference-only"), + ], + ids=["dashed-size", "auto-size-unparsed", "garbage-size-unparsed"], +) +def test_size_string_parsing(size: str, expected_generated: str) -> None: + catalog_rate: Final = litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] + expected: Final = ( + catalog_rate * 1024 * 1024 if expected_generated == "computed" else 0.0 + ) + catalog_rate * REFERENCE_PIXELS + + cost: Final = _edit_cost(_edit_response(), size=size) + + assert cost == pytest.approx(expected) + + def test_optional_params_dimensions_beat_response_size() -> None: cost: Final = _edit_cost( _edit_response(reference_pixels=None), diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 7023e4ff6e1..899b0d95465 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3927,7 +3927,8 @@ def test_image_edit_handler_stamps_measured_reference_pixels(): client=client, ) - assert response._reference_pixels == 4 * 2 + 1 * 1 + assert response.reference_pixels == 4 * 2 + 1 * 1 + assert "reference_pixels" not in response.model_dump() async def test_async_image_edit_handler_stamps_measured_reference_pixels(): @@ -3947,7 +3948,7 @@ async def test_async_image_edit_handler_stamps_measured_reference_pixels(): client=client, ) - assert response._reference_pixels == 8 + assert response.reference_pixels == 8 def test_image_edit_handler_leaves_reference_pixels_unset_when_unmeasurable(): @@ -3967,7 +3968,7 @@ def test_image_edit_handler_leaves_reference_pixels_unset_when_unmeasurable(): client=client, ) - assert response._reference_pixels is None + assert response.reference_pixels is None class _ScriptedClientWebSocket(_FakeClientWebSocket):