Merge 4150fe74ab into 527d882508

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)
 

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