Fix: Stream the download in chunks

litellm/litellm_core_utils/prompt_templates/image_handling.py

@@ -31,15 +31,19 @@ def _process_image_response(response: Response, url: str) -> str:
                 f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}"
             )
 
-    image_bytes = response.content
+    # Stream download with size checking to prevent downloading huge files
+    max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024)
+    image_bytes = bytearray()
+    bytes_downloaded = 0
     
-    # Check actual size after download if Content-Length was not available
-    if content_length is None:
-        size_mb = len(image_bytes) / (1024 * 1024)
-        if size_mb > MAX_IMAGE_URL_DOWNLOAD_SIZE_MB:
+    for chunk in response.iter_bytes(chunk_size=8192):
+        bytes_downloaded += len(chunk)
+        if bytes_downloaded > max_bytes:
+            size_mb = bytes_downloaded / (1024 * 1024)
             raise litellm.ImageFetchError(
                 f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}"
             )
+        image_bytes.extend(chunk)
     
     base64_image = base64.b64encode(image_bytes).decode("utf-8")
 

tests/test_litellm/litellm_core_utils/test_image_handling.py

@@ -66,6 +66,41 @@ class LargeImageClient:
         )
 
 
+class StreamingLargeImageClient:
+    """
+    Client that streams a large image to test streaming download protection.
+    This simulates a huge file without actually creating it all in memory.
+    """
+
+    def __init__(self, size_mb=100, include_content_length=False):
+        self.size_mb = size_mb
+        self.include_content_length = include_content_length
+
+    def get(self, url, follow_redirects=True):
+        size_bytes = int(self.size_mb * 1024 * 1024)
+        headers = {"Content-Type": "image/jpeg"}
+        if self.include_content_length:
+            headers["Content-Length"] = str(size_bytes)
+        
+        # Create a generator that yields chunks without creating the whole file in memory
+        def generate_chunks(total_size, chunk_size=8192):
+            bytes_sent = 0
+            while bytes_sent < total_size:
+                chunk = b"x" * min(chunk_size, total_size - bytes_sent)
+                bytes_sent += len(chunk)
+                yield chunk
+        
+        # Create response with streaming content
+        response = Response(
+            status_code=200,
+            headers=headers,
+            request=Request("GET", url),
+        )
+        # Mock the iter_bytes method to return our generator
+        response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size)
+        return response
+
+
 def test_image_exceeds_size_limit_with_content_length(monkeypatch):
     """
     Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present.
@@ -83,6 +118,7 @@ def test_image_exceeds_size_limit_with_content_length(monkeypatch):
 def test_image_exceeds_size_limit_without_content_length(monkeypatch):
     """
     Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header.
+    This uses the old non-streaming mock for backward compatibility.
     """
     monkeypatch.setattr(
         litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False)
@@ -94,6 +130,29 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch):
     assert "exceeds maximum allowed size" in str(excinfo.value)
 
 
+def test_streaming_download_protects_against_huge_files(monkeypatch):
+    """
+    Test that streaming download aborts early when file exceeds size limit,
+    preventing memory exhaustion from huge files (e.g., petabyte-sized files).
+    
+    This test verifies that the streaming implementation doesn't download the entire
+    file into memory before checking size. Instead, it should abort as soon as the
+    limit is exceeded during streaming.
+    """
+    # Simulate a 1GB file - far larger than the 50MB default limit
+    client = StreamingLargeImageClient(size_mb=1024, include_content_length=False)
+    monkeypatch.setattr(litellm, "module_level_client", client)
+
+    with pytest.raises(litellm.ImageFetchError) as excinfo:
+        convert_url_to_base64("https://example.com/huge-image.jpg")
+
+    # Verify the error message shows it was caught during streaming
+    assert "exceeds maximum allowed size" in str(excinfo.value)
+    
+    # The error should be raised after downloading just slightly more than the limit
+    # not after downloading the full 1GB
+
+
 class SmallImageClient:
     """
     Client that returns a small valid image.
@@ -124,6 +183,26 @@ def test_image_within_size_limit(monkeypatch):
     assert result.startswith("data:image/jpeg;base64,")
 
 
+def test_streaming_download_handles_petabyte_file(monkeypatch):
+    """
+    Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized)
+    without attempting to download the entire file or causing memory exhaustion.
+    
+    This simulates what happens if a malicious actor or misconfiguration provides
+    a URL to an extremely large file.
+    """
+    # Simulate a 1 petabyte file (1,000,000 GB)
+    # Without streaming protection, this would cause OOM or hang indefinitely
+    client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False)
+    monkeypatch.setattr(litellm, "module_level_client", client)
+
+    with pytest.raises(litellm.ImageFetchError) as excinfo:
+        convert_url_to_base64("https://example.com/petabyte-file.jpg")
+
+    # Should fail fast without downloading anywhere near 1 petabyte
+    assert "exceeds maximum allowed size" in str(excinfo.value)
+
+
 def test_image_size_limit_disabled(monkeypatch):
     """
     Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads.