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