From e38b01c9737f690515c89873c1f56fb1807efd54 Mon Sep 17 00:00:00 2001 From: icn5381 <255778606+icn5381@users.noreply.github.com> Date: Mon, 17 Aug 2026 02:31:55 +0800 Subject: [PATCH 1/3] fix(core): stop get_image_dimensions crashing on bare base64 and unfetchable URLs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_image_dimensions documents 'URL or base64 encoded string' input, but the base64 branch unconditionally split on ',' and unpacked two values, so bare base64 (no data: prefix) raised ValueError, and a URL whose download failed fell into the same branch and crashed the same way — instead of reaching the documented 'sensible default dimensions' fallback that only the unknown-image-format path could reach. The base64 branch now accepts both data-URL and bare base64 forms and treats a decode error as 'use defaults'; an unfetchable URL or undecodable payload leaves img_data unset and skips image-format detection, falling through to DEFAULT_IMAGE_WIDTH/HEIGHT. --- litellm/litellm_core_utils/token_counter.py | 17 +++++-- .../litellm_core_utils/test_token_counter.py | 47 +++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 17f3dea72ec..332f2803f3c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -213,12 +213,19 @@ def get_image_dimensions( img_data = body except Exception: pass - if img_data is None: - # Not a URL or fetch failed — assume base64 - _header, encoded = data.split(",", 1) - img_data = base64.b64decode(encoded) + if img_data is None and not data.startswith(("http://", "https://")): + # Not a URL — assume base64. Accept both data-URL form + # ('data:;base64,') and bare base64; on a decode + # error, leave img_data unset so the default dimensions are used. + encoded = data.partition(",")[2] if data.startswith("data:") else data + try: + img_data = base64.b64decode(encoded) + except ValueError: + verbose_logger.debug("Failed to decode base64 image data; using default dimensions") - img_type: Final = get_image_type(img_data) + # A URL that could not be fetched (or base64 that could not be decoded) + # leaves img_data unset — fall through to the default dimensions below. + img_type: Final = get_image_type(img_data) if img_data is not None else None if img_type == "png": w, h = struct.unpack(">LL", img_data[16:24]) 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 71e686563a5..c585f7a7792 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -526,6 +526,53 @@ def test_img_url_token_counter(img_url, monkeypatch): assert height is not None +def test_get_image_dimensions_bare_base64(): + """ + Bare base64 (no 'data:' URL prefix) is documented input for + get_image_dimensions and must not crash on the ','-split of a data URL. + """ + import base64 + import struct + + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + # Minimal PNG header carrying 100x200 dimensions. + png = base64.b64encode( + b"\x89PNG\r\n\x1a\n" + b"\x00" * 8 + struct.pack(">LL", 100, 200) + b"\x00" * 16 + ) + assert get_image_dimensions(data=png.decode()) == (100, 200) + + +def test_get_image_dimensions_unfetchable_url_returns_defaults(monkeypatch): + """A URL that cannot be fetched falls back to the default dimensions instead of raising.""" + + def _raise(client, url, **kwargs): + raise RuntimeError("connection failed") + + monkeypatch.setattr( + "litellm.litellm_core_utils.token_counter.safe_get", + _raise, + ) + from litellm.constants import DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + assert get_image_dimensions(data="https://invalid.invalid/x.png") == ( + DEFAULT_IMAGE_WIDTH, + DEFAULT_IMAGE_HEIGHT, + ) + + +def test_get_image_dimensions_undecodable_data_returns_defaults(): + """Data that is neither a fetchable URL nor decodable base64 uses the defaults.""" + from litellm.constants import DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + assert get_image_dimensions(data="not-an-image!") == ( + DEFAULT_IMAGE_WIDTH, + DEFAULT_IMAGE_HEIGHT, + ) + + def test_token_encode_disallowed_special(): encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") token_counter(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>") From 716acca88862c30b9c91e547224a832278f0dd96 Mon Sep 17 00:00:00 2001 From: icn5381 <255778606+icn5381@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:22:06 +0800 Subject: [PATCH 2/3] Refactor base64 decoding into a helper to stay within the strict-rule budget The fallback branches added to get_image_dimensions pushed the repo-wide C901 total one over its ceiling in ruff-strict-budget.json. Extracting the data-URL/bare-base64 decoding into a small helper keeps the same behavior (verified by the image tests) and lands the function one complexity point below its upstream baseline. --- litellm/litellm_core_utils/token_counter.py | 25 ++++++++++++++------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 332f2803f3c..bb34986d9b9 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -186,6 +186,20 @@ def get_image_type(image_data: bytes) -> str | None: return None +def _decode_base64_image(data: str) -> "bytes | None": + """Decode base64 image data, accepting both data-URL and bare base64 forms. + + Returns None (and logs at debug) for input that does not decode, so callers + can fall back to default dimensions instead of failing the call. + """ + encoded = data.partition(",")[2] if data.startswith("data:") else data + try: + return base64.b64decode(encoded) + except ValueError: + verbose_logger.debug("Failed to decode base64 image data; using default dimensions") + return None + + def get_image_dimensions( data: str, ) -> tuple[int, int]: @@ -214,14 +228,9 @@ def get_image_dimensions( except Exception: pass if img_data is None and not data.startswith(("http://", "https://")): - # Not a URL — assume base64. Accept both data-URL form - # ('data:;base64,') and bare base64; on a decode - # error, leave img_data unset so the default dimensions are used. - encoded = data.partition(",")[2] if data.startswith("data:") else data - try: - img_data = base64.b64decode(encoded) - except ValueError: - verbose_logger.debug("Failed to decode base64 image data; using default dimensions") + # Not a URL or fetch failed — assume base64, keeping None on decode + # errors so the default dimensions are used below. + img_data = _decode_base64_image(data) # A URL that could not be fetched (or base64 that could not be decoded) # leaves img_data unset — fall through to the default dimensions below. From 5ca8c0154fcc634a7cd3c3e4d7655fd47dcf2c17 Mon Sep 17 00:00:00 2001 From: icn5381 <255778606+icn5381@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:49:43 +0800 Subject: [PATCH 3/3] Narrow img_data before format detection to clear the basedpyright budget The None-tolerant flow left img_data Optional at every subscript in the format-detection branches, adding ten reportOptionalSubscript errors against a zero-error budget. An early return on None (plus an explicit annotation) narrows the type before detection; complexity returns to the upstream baseline, so the C901 budget stays level too. --- litellm/litellm_core_utils/token_counter.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index bb34986d9b9..f11a85fe64e 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -212,7 +212,7 @@ def get_image_dimensions( Returns: Tuple[int, int]: The width and height of the image. """ - img_data = None + img_data: bytes | None = None if data.startswith(("http://", "https://")): try: client: Final = _get_httpx_client() @@ -233,8 +233,11 @@ def get_image_dimensions( img_data = _decode_base64_image(data) # A URL that could not be fetched (or base64 that could not be decoded) - # leaves img_data unset — fall through to the default dimensions below. - img_type: Final = get_image_type(img_data) if img_data is not None else None + # leaves img_data unset — return the default dimensions. + if img_data is None: + return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT + + img_type: Final = get_image_type(img_data) if img_type == "png": w, h = struct.unpack(">LL", img_data[16:24])