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

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

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