From f994a164c0bf7cbbf827e2d82b2fab90739456a0 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 21 Jan 2026 00:53:33 +0000 Subject: [PATCH] Fix: Stream image downloads with incremental size validation Previously, the image download handler would download the entire file into memory before checking the size (when Content-Length header was not provided). This could cause memory issues with very large files. Changes: - Implement streaming download using iter_bytes() in 8KB chunks - Check cumulative size after each chunk during download - Stop immediately when size exceeds MAX_IMAGE_URL_DOWNLOAD_SIZE_MB - Add both sync and async versions of streaming validation - Maintain efficient Content-Length header check for early rejection This prevents downloading potentially huge files (e.g., petabyte-sized) by stopping as soon as the limit is exceeded (~50MB + 8KB max). Addresses feedback on PR #19257 --- .../prompt_templates/image_handling.py | 94 +++++++++++++++++-- 1 file changed, 86 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5d0bedb776d..030bd97452d 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -16,7 +16,11 @@ MAX_IMGS_IN_MEMORY = 10 in_memory_cache = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) -def _process_image_response(response: Response, url: str) -> str: +async def _async_stream_and_validate_image(response: Response, url: str) -> bytes: + """ + Async stream download an image and validate size during download. + Stops immediately if size exceeds limit. + """ if response.status_code != 200: raise litellm.ImageFetchError( f"Error: Unable to fetch image from URL. Status code: {response.status_code}, url={url}" @@ -31,16 +35,90 @@ 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 - - # Check actual size after download if Content-Length was not available - if content_length is None: - size_mb = len(image_bytes) / (1024 * 1024) + # Stream the response and check size incrementally + max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + image_bytes = b"" + chunk_size = 8192 # 8KB chunks + + async for chunk in response.aiter_bytes(chunk_size=chunk_size): + image_bytes += chunk + if len(image_bytes) > max_bytes: + size_mb = len(image_bytes) / (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}" + ) + + return image_bytes + + +def _stream_and_validate_image(response: Response, url: str) -> bytes: + """ + Stream download an image and validate size during download. + Stops immediately if size exceeds limit. + """ + if response.status_code != 200: + raise litellm.ImageFetchError( + f"Error: Unable to fetch image from URL. Status code: {response.status_code}, url={url}" + ) + + # Check size before downloading if Content-Length header is present + content_length = response.headers.get("Content-Length") + if content_length is not None: + size_mb = int(content_length) / (1024 * 1024) if size_mb > MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: raise litellm.ImageFetchError( f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) - + + # Stream the response and check size incrementally + max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + image_bytes = b"" + chunk_size = 8192 # 8KB chunks + + for chunk in response.iter_bytes(chunk_size=chunk_size): + image_bytes += chunk + if len(image_bytes) > max_bytes: + size_mb = len(image_bytes) / (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}" + ) + + return image_bytes + + +def _process_image_response(response: Response, url: str) -> str: + """Process image response and convert to base64 data URL.""" + image_bytes = _stream_and_validate_image(response, url) + + base64_image = base64.b64encode(image_bytes).decode("utf-8") + + image_type = response.headers.get("Content-Type") + if image_type is None: + img_type = url.split(".")[-1].lower() + _img_type = { + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "png": "image/png", + "gif": "image/gif", + "webp": "image/webp", + }.get(img_type) + if _img_type is None: + raise Exception( + f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" + ) + img_type = _img_type + else: + img_type = image_type + + result = f"data:{img_type};base64,{base64_image}" + in_memory_cache.set_cache(url, result) + return result + + +async def _async_process_image_response(response: Response, url: str) -> str: + """Async process image response and convert to base64 data URL.""" + image_bytes = await _async_stream_and_validate_image(response, url) + base64_image = base64.b64encode(image_bytes).decode("utf-8") image_type = response.headers.get("Content-Type") @@ -81,7 +159,7 @@ async def async_convert_url_to_base64(url: str) -> str: for _ in range(3): try: response = await client.get(url, follow_redirects=True) - return _process_image_response(response, url) + return await _async_process_image_response(response, url) except litellm.ImageFetchError: raise except Exception: