From dd08db6f793ea068bbfafaddd6dd4fe804f3f940 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 24 Sep 2026 00:19:45 +0000 Subject: [PATCH] feat(images): stamp measured reference pixels on image edit responses --- litellm/llms/custom_httpx/llm_http_handler.py | 31 +++++--- litellm/types/utils.py | 4 ++ .../custom_httpx/test_llm_http_handler.py | 72 +++++++++++++++++++ 3 files changed, 99 insertions(+), 8 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 052978c2680..2a40d578af4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -33,6 +33,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult +from litellm.images.dimensions import total_reference_pixels from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, @@ -382,6 +383,12 @@ 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: + response._reference_pixels = reference_pixels # mutable-ok: stamp the measured reference pixels + return response + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -6881,6 +6888,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) + reference_pixels: Final = total_reference_pixels(image if isinstance(image, list) else [image]) data, files = image_edit_provider_config.transform_image_edit_request( model=model, image=image, @@ -6928,10 +6936,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( - model=model, - raw_response=response, - logging_obj=logging_obj, + return _with_reference_pixels( + image_edit_provider_config.transform_image_edit_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ), + reference_pixels, ) async def async_image_edit_handler( @@ -6980,6 +6991,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) + reference_pixels: Final = total_reference_pixels(image if isinstance(image, list) else [image]) data, files = await image_edit_provider_config.async_transform_image_edit_request( model=model, image=image, @@ -7027,10 +7039,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( - model=model, - raw_response=response, - logging_obj=logging_obj, + return _with_reference_pixels( + image_edit_provider_config.transform_image_edit_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ), + reference_pixels, ) def image_generation_handler( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3f8471cbde1..c31b7ef886a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -318,6 +318,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_audio_token_batches: ReadOnly[float | None] input_cost_per_image_token_batches: ReadOnly[float | None] input_cost_per_second: float | None # for OpenAI Speech models + input_cost_per_pixel: ReadOnly[float | None] + input_cost_per_reference_pixel: ReadOnly[float | None] input_cost_per_token_batches: float | None input_cost_per_video_token_batches: ReadOnly[float | None] input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None] @@ -2582,6 +2584,7 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} + _reference_pixels: int | None = None usage: ImageUsage | None = None """ @@ -3702,6 +3705,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_image_1024: float | None = None output_cost_per_image_1536: float | None = None input_cost_per_pixel: float | None = None + input_cost_per_reference_pixel: float | None = None output_cost_per_pixel: float | None = None # Include all ModelInfoBase fields as optional 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 68f37c8ffcc..7023e4ff6e1 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 @@ -2,6 +2,7 @@ import asyncio import base64 import json import logging +import struct import threading import time from typing import Final @@ -3898,6 +3899,77 @@ def test_image_edit_handler_keeps_the_sync_transform(): assert response.data[0].b64_json == "sync" +def _tiny_png(width: int, height: int) -> bytes: + return ( + b"\x89PNG\r\n\x1a\n" + + struct.pack(">I", 13) + + b"IHDR" + + struct.pack(">II", width, height) + + b"\x08\x06\x00\x00\x00" + + bytes(4) + ) + + +def test_image_edit_handler_stamps_measured_reference_pixels(): + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport({})) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=[_tiny_png(4, 2), _tiny_png(1, 1)], + prompt="add a hat", + image_edit_provider_config=_ImageEditRecordingConfig(), + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert response._reference_pixels == 4 * 2 + 1 * 1 + + +async def test_async_image_edit_handler_stamps_measured_reference_pixels(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_echo_json_transport({})) + + response = await BaseLLMHTTPHandler().async_image_edit_handler( + model="edit-model", + image=_tiny_png(4, 2), + prompt="add a hat", + image_edit_provider_config=_ImageEditRecordingConfig(), + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert response._reference_pixels == 8 + + +def test_image_edit_handler_leaves_reference_pixels_unset_when_unmeasurable(): + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport({})) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=b"not-an-image", + prompt="add a hat", + image_edit_provider_config=_ImageEditRecordingConfig(), + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert response._reference_pixels is None + + class _ScriptedClientWebSocket(_FakeClientWebSocket): def __init__(self, messages: list[str], last_event_type: str) -> None: super().__init__()