From 46929b6d1f69a071b4e2b1123d26f15ab0d714ec Mon Sep 17 00:00:00 2001 From: Chau Vu / CPF-FAMILY Date: Wed, 30 Sep 2026 10:25:32 +0700 Subject: [PATCH 1/2] fix(token_counter): safely parse raw base64 and handle failed image URL fetches in get_image_dimensions - Check whether comma exists before splitting header in base64 data URIs - Safely handle raw base64 image strings without data: URI header - Raise informative ValueError on failed URL fetches and decoding failures - Add regression tests in tests/unit/litellm_core_utils/test_token_counter.py - Fixes #43805 --- litellm/litellm_core_utils/token_counter.py | 12 ++++++++-- .../litellm_core_utils/test_token_counter.py | 24 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index cdd2d0654be..513aaf6a691 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -245,8 +245,16 @@ def get_image_dimensions( 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 "," in data: + _header, encoded = data.split(",", 1) + else: + encoded = data + try: + img_data = base64.b64decode(encoded) + except Exception as e: + raise ValueError( + f"Unable to parse or decode image data: {e}" + ) from e img_type: Final = get_image_type(img_data) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index c71b1496bdd..609be552be2 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -719,6 +719,30 @@ def test_img_url_token_counter(img_url, monkeypatch): assert height is not None +def test_get_image_dimensions_raw_base64(): + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + raw_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + w, h = get_image_dimensions(data=raw_b64) + assert w == 1 + assert h == 1 + + +def test_get_image_dimensions_failed_url_raises_valueerror(monkeypatch): + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + def _failing_get(client, url, **kw): + raise Exception("Connection timeout") + + monkeypatch.setattr( + "litellm.litellm_core_utils.token_counter.safe_get", + _failing_get, + ) + + with pytest.raises(ValueError, match="Unable to parse or decode image data"): + get_image_dimensions(data="https://unreachable-endpoint.com/test.png") + + 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 4150fe74ab5528f330cb22ff76d459d61a8be326 Mon Sep 17 00:00:00 2001 From: Chau Vu / CPF-FAMILY Date: Wed, 30 Sep 2026 10:31:34 +0700 Subject: [PATCH 2/2] fix(token_counter): address review comments by distinguishing URL fetch failures and adding type annotations --- litellm/litellm_core_utils/token_counter.py | 18 ++++++++++-------- .../litellm_core_utils/test_token_counter.py | 6 +++--- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 513aaf6a691..77d09543526 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -233,18 +233,20 @@ def get_image_dimensions( try: client: Final = _get_httpx_client() response: Final[httpx.Response] = safe_get(client, data) + response.raise_for_status() max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) content_length: Final[str | None] = response.headers.get("Content-Length") if content_length is not None and int(content_length) > max_bytes: - pass # skip download; img_data stays None + raise ValueError(f"Image at {data} exceeds maximum allowed size of {MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB") + body: Final = response.read() + if len(body) <= max_bytes: + img_data = body else: - body: Final = response.read() - if len(body) <= max_bytes: - img_data = body - except Exception: - pass - if img_data is None: - # Not a URL or fetch failed — assume base64 + raise ValueError(f"Image at {data} exceeds maximum allowed size of {MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB") + except Exception as e: + raise ValueError(f"Failed to fetch image from URL {data}: {e}") from e + else: + # Base64 string (with or without data URI header) if "," in data: _header, encoded = data.split(",", 1) else: diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 609be552be2..3e2b3f870a1 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -728,10 +728,10 @@ def test_get_image_dimensions_raw_base64(): assert h == 1 -def test_get_image_dimensions_failed_url_raises_valueerror(monkeypatch): +def test_get_image_dimensions_failed_url_raises_valueerror(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.litellm_core_utils.token_counter import get_image_dimensions - def _failing_get(client, url, **kw): + def _failing_get(client: object, url: str, **kw: object) -> object: raise Exception("Connection timeout") monkeypatch.setattr( @@ -739,7 +739,7 @@ def test_get_image_dimensions_failed_url_raises_valueerror(monkeypatch): _failing_get, ) - with pytest.raises(ValueError, match="Unable to parse or decode image data"): + with pytest.raises(ValueError, match="Failed to fetch image from URL"): get_image_dimensions(data="https://unreachable-endpoint.com/test.png")