diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 52a5026c44f..8660ccd5b7f 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -266,20 +266,30 @@ 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] = header_value(response.headers, "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 - _header, encoded = data.split(",", 1) - img_data = base64.b64decode(encoded) + 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: + 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 return _image_dimensions_from_bytes(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 d0b837646b4..05b8b024b57 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1038,6 +1038,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: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.token_counter import get_image_dimensions + + def _failing_get(client: object, url: str, **kw: object) -> object: + raise Exception("Connection timeout") + + monkeypatch.setattr( + "litellm.litellm_core_utils.token_counter.safe_get", + _failing_get, + ) + + with pytest.raises(ValueError, match="Failed to fetch image from URL"): + 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|>")