Merge 5ca8c0154f into c2c2a623c0

litellm/litellm_core_utils/token_counter.py

@@ -213,6 +213,20 @@ def get_image_type(image_data: bytes) -> str | None:
     return None
 
 
+def _decode_base64_image(data: str) -> "bytes | None":
+    """Decode base64 image data, accepting both data-URL and bare base64 forms.
+
+    Returns None (and logs at debug) for input that does not decode, so callers
+    can fall back to default dimensions instead of failing the call.
+    """
+    encoded = data.partition(",")[2] if data.startswith("data:") else data
+    try:
+        return base64.b64decode(encoded)
+    except ValueError:
+        verbose_logger.debug("Failed to decode base64 image data; using default dimensions")
+        return None
+
+
 def get_image_dimensions(
     data: str,
 ) -> tuple[int, int]:
@@ -225,7 +239,7 @@ def get_image_dimensions(
     Returns:
         Tuple[int, int]: The width and height of the image.
     """
-    img_data = None
+    img_data: bytes | None = None
     if data.startswith(("http://", "https://")):
         try:
             client: Final = _get_httpx_client()
@@ -240,10 +254,15 @@ def get_image_dimensions(
                     img_data = body
         except Exception:
             pass
+    if img_data is None and not data.startswith(("http://", "https://")):
+        # Not a URL or fetch failed — assume base64, keeping None on decode
+        # errors so the default dimensions are used below.
+        img_data = _decode_base64_image(data)
+
+    # A URL that could not be fetched (or base64 that could not be decoded)
+    # leaves img_data unset — return the default dimensions.
     if img_data is None:
-        # Not a URL or fetch failed — assume base64
-        _header, encoded = data.split(",", 1)
-        img_data = base64.b64decode(encoded)
+        return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT
 
     img_type: Final = get_image_type(img_data)
 

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -741,6 +741,53 @@ def test_img_url_token_counter(img_url, monkeypatch):
     assert height is not None
 
 
+def test_get_image_dimensions_bare_base64():
+    """
+    Bare base64 (no 'data:' URL prefix) is documented input for
+    get_image_dimensions and must not crash on the ','-split of a data URL.
+    """
+    import base64
+    import struct
+
+    from litellm.litellm_core_utils.token_counter import get_image_dimensions
+
+    # Minimal PNG header carrying 100x200 dimensions.
+    png = base64.b64encode(
+        b"\x89PNG\r\n\x1a\n" + b"\x00" * 8 + struct.pack(">LL", 100, 200) + b"\x00" * 16
+    )
+    assert get_image_dimensions(data=png.decode()) == (100, 200)
+
+
+def test_get_image_dimensions_unfetchable_url_returns_defaults(monkeypatch):
+    """A URL that cannot be fetched falls back to the default dimensions instead of raising."""
+
+    def _raise(client, url, **kwargs):
+        raise RuntimeError("connection failed")
+
+    monkeypatch.setattr(
+        "litellm.litellm_core_utils.token_counter.safe_get",
+        _raise,
+    )
+    from litellm.constants import DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH
+    from litellm.litellm_core_utils.token_counter import get_image_dimensions
+
+    assert get_image_dimensions(data="https://invalid.invalid/x.png") == (
+        DEFAULT_IMAGE_WIDTH,
+        DEFAULT_IMAGE_HEIGHT,
+    )
+
+
+def test_get_image_dimensions_undecodable_data_returns_defaults():
+    """Data that is neither a fetchable URL nor decodable base64 uses the defaults."""
+    from litellm.constants import DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH
+    from litellm.litellm_core_utils.token_counter import get_image_dimensions
+
+    assert get_image_dimensions(data="not-an-image!") == (
+        DEFAULT_IMAGE_WIDTH,
+        DEFAULT_IMAGE_HEIGHT,
+    )
+
+
 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|>")