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] 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")