From 4ad18797abbd0e7bf076a526b209100d64c0eb0d Mon Sep 17 00:00:00 2001 From: Shreshth Kharbanda Date: Wed, 23 Sep 2026 13:33:38 -0700 Subject: [PATCH] refactor(token_counter): split image_dimensions_from_bytes out of get_image_dimensions get_image_dimensions fetched or decoded its input, parsed the header, and substituted a 300x300 default when the header was unreadable, all in one function. Billing code needs the middle step on its own: given bytes, return the dimensions or None, never a default. The parsers are unchanged and the string entry point still applies the default, so its callers see no difference. Unreadable headers used to raise TypeError or struct.error out of the parser; the bytes entry point turns those into None --- litellm/litellm_core_utils/token_counter.py | 32 +++++--- .../litellm_core_utils/test_token_counter.py | 74 +++++++++++++++++++ 2 files changed, 94 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 5d7956059e4..bf0673ec816 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -248,15 +248,29 @@ def get_image_dimensions( _header, encoded = data.split(",", 1) img_data = base64.b64decode(encoded) + dimensions: Final = image_dimensions_from_bytes(img_data) + if dimensions is None: + return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT + return dimensions + + +def image_dimensions_from_bytes(img_data: bytes) -> tuple[int, int] | None: + try: + return _header_dimensions(img_data) + except (struct.error, TypeError): + return None + + +def _header_dimensions(img_data: bytes) -> tuple[int, int] | None: img_type: Final = get_image_type(img_data) if img_type == "png": w, h = _unpack_ints(">LL", img_data[16:24]) return w, h - elif img_type == "gif": + if img_type == "gif": w, h = _unpack_ints("HH", fhandle.read(4)) return w, h - elif img_type == "webp": - # For WebP, the dimensions are stored at different offsets depending on the format - # Check for VP8X (extended format) + if img_type == "webp": if img_data[12:16] == b"VP8X": w = _unpack_ints("> 14) & 0x3FFF) + 1 return w, h - - # return sensible default image dimensions if unable to get dimensions - return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT + return None def calculate_img_tokens( diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index eccf44a1bda..53a09d8feed 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -3,6 +3,7 @@ import asyncio import base64 import importlib +import struct import threading import time import traceback @@ -28,6 +29,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_tiktoken_count_function, calculate_img_tokens, high_detail_image_token_upper_bound, + image_dimensions_from_bytes, offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new @@ -1599,3 +1601,75 @@ def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() + + +def _png_bytes(width: int, height: int) -> bytes: + return ( + b"\x89PNG\r\n\x1a\n" + + (13).to_bytes(4, "big") + + b"IHDR" + + struct.pack(">II", width, height) + + b"\x08\x02\x00\x00\x00" + ) + + +def _gif_bytes(width: int, height: int) -> bytes: + return b"GIF89a" + struct.pack(" bytes: + app: Final = b"".join( + b"\xff\xe0" + struct.pack(">H", 16) + b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00" + for _ in range(app_segments) + ) + sof: Final = ( + b"\xff" + sof_marker + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" + ) + return b"\xff\xd8" + app + sof + + +def _webp_bytes(chunk: bytes, payload: bytes) -> bytes: + body: Final = chunk + struct.pack(" bytes: + return _webp_bytes(b"VP8 ", b"\x00\x00\x00\x9d\x01\x2a" + struct.pack(" bytes: + return _webp_bytes(b"VP8L", b"\x2f" + struct.pack(" bytes: + return _webp_bytes(b"VP8X", b"\x00" * 4 + (width - 1).to_bytes(3, "little") + (height - 1).to_bytes(3, "little")) + + +@pytest.mark.parametrize( + ("image", "expected"), + [ + pytest.param(_png_bytes(1024, 768), (1024, 768), id="png"), + pytest.param(_gif_bytes(100, 50), (100, 50), id="gif"), + pytest.param(_jpeg_bytes(800, 600, b"\xc0", 1), (800, 600), id="jpeg-baseline"), + pytest.param(_jpeg_bytes(640, 480, b"\xc2", 3), (640, 480), id="jpeg-progressive-after-app-segments"), + pytest.param(_webp_vp8_bytes(640, 480), (640, 480), id="webp-vp8"), + pytest.param(_webp_vp8l_bytes(320, 240), (320, 240), id="webp-vp8l"), + pytest.param(_webp_vp8x_bytes(1920, 1080), (1920, 1080), id="webp-vp8x"), + ], +) +def test_image_dimensions_from_bytes_reads_each_header_format(image: bytes, expected: tuple[int, int]) -> None: + assert image_dimensions_from_bytes(image) == expected + + +@pytest.mark.parametrize( + "image", + [ + pytest.param(b"", id="empty"), + pytest.param(b"BM" + b"\x00" * 30, id="unknown-format"), + pytest.param(b"\x89PNG\r\n\x1a\n\x00\x00", id="png-truncated-before-ihdr"), + pytest.param(b"\xff\xd8\xff\xe0\x00\x10JFIF", id="jpeg-truncated-inside-app0"), + pytest.param(b"\xff\xd8\xff\xe0\x00\x04\x00\x00", id="jpeg-ends-before-sof"), + ], +) +def test_image_dimensions_from_bytes_returns_none_for_unreadable_headers(image: bytes) -> None: + assert image_dimensions_from_bytes(image) is None