From f6d2bd1c7a49fa84cae9f6f2ca37270275bd089d Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 21 Jan 2026 01:44:48 +0000 Subject: [PATCH] Refactor: Extract constants and add E2E tests for streaming validation Changes: - Extract DOWNLOAD_CHUNK_SIZE_BYTES constant (8KB) to replace magic numbers - Add comprehensive E2E unit tests for streaming image download behavior: * test_streaming_download_regular_image: Validates normal images download correctly * test_async_streaming_download_regular_image: Async variant * test_streaming_stops_for_large_image_without_content_length: Critical test proving we stop downloading at ~50MB instead of downloading entire file * test_async_streaming_stops_for_large_image_without_content_length: Async variant Tests validate: - No regression in normal image downloads (1-2MB images work) - Streaming stops early for oversized files (downloads ~50MB of 100MB file) - Both sync and async code paths work correctly This ensures the streaming validation fix works as intended and prevents downloading huge files entirely. --- .../prompt_templates/image_handling.py | 7 +- .../litellm_core_utils/test_image_handling.py | 153 +++++++++++++++++- 2 files changed, 155 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 030bd97452d..82b3cbdee7a 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -12,6 +12,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB MAX_IMGS_IN_MEMORY = 10 +DOWNLOAD_CHUNK_SIZE_BYTES = 8192 # 8KB chunks for streaming downloads in_memory_cache = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) @@ -38,9 +39,8 @@ async def _async_stream_and_validate_image(response: Response, url: str) -> byte # 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): + async for chunk in response.aiter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE_BYTES): image_bytes += chunk if len(image_bytes) > max_bytes: size_mb = len(image_bytes) / (1024 * 1024) @@ -73,9 +73,8 @@ def _stream_and_validate_image(response: Response, url: str) -> bytes: # 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): + for chunk in response.iter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE_BYTES): image_bytes += chunk if len(image_bytes) > max_bytes: size_mb = len(image_bytes) / (1024 * 1024) diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index b15d75a4145..dab15638d8f 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -7,6 +7,7 @@ import litellm from litellm import constants from litellm.litellm_core_utils.prompt_templates.image_handling import ( convert_url_to_base64, + async_convert_url_to_base64, ) @@ -135,6 +136,156 @@ def test_image_size_limit_disabled(monkeypatch): with pytest.raises(litellm.ImageFetchError) as excinfo: convert_url_to_base64("https://example.com/image.jpg") - + assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +class StreamingImageClient: + """ + Client that simulates streaming image downloads. + Tracks how many bytes were actually downloaded via streaming. + """ + + def __init__(self, total_size_bytes, include_content_length=False): + self.total_size_bytes = total_size_bytes + self.include_content_length = include_content_length + self.bytes_downloaded = 0 + + def get(self, url, follow_redirects=True): + client_ref = self + + class StreamingResponse: + def __init__(self): + self.status_code = 200 + self.headers = {"Content-Type": "image/jpeg"} + if client_ref.include_content_length: + self.headers["Content-Length"] = str(client_ref.total_size_bytes) + + def iter_bytes(self, chunk_size=8192): + """Simulate streaming bytes in chunks""" + bytes_sent = 0 + while bytes_sent < client_ref.total_size_bytes: + chunk_bytes = min(chunk_size, client_ref.total_size_bytes - bytes_sent) + client_ref.bytes_downloaded += chunk_bytes + bytes_sent += chunk_bytes + yield b"x" * chunk_bytes + + return StreamingResponse() + + +class AsyncStreamingImageClient: + """ + Async client that simulates streaming image downloads. + Tracks how many bytes were actually downloaded via streaming. + """ + + def __init__(self, total_size_bytes, include_content_length=False): + self.total_size_bytes = total_size_bytes + self.include_content_length = include_content_length + self.bytes_downloaded = 0 + + async def get(self, url, follow_redirects=True): + client_ref = self + + class AsyncStreamingResponse: + def __init__(self): + self.status_code = 200 + self.headers = {"Content-Type": "image/jpeg"} + if client_ref.include_content_length: + self.headers["Content-Length"] = str(client_ref.total_size_bytes) + + async def aiter_bytes(self, chunk_size=8192): + """Simulate async streaming bytes in chunks""" + bytes_sent = 0 + while bytes_sent < client_ref.total_size_bytes: + chunk_bytes = min(chunk_size, client_ref.total_size_bytes - bytes_sent) + client_ref.bytes_downloaded += chunk_bytes + bytes_sent += chunk_bytes + yield b"x" * chunk_bytes + + return AsyncStreamingResponse() + + +def test_streaming_download_regular_image(monkeypatch): + """ + E2E test: Verify that streaming download works correctly for regular-sized images. + This ensures no regression in normal image download functionality. + """ + # Use a 1MB image (well within the 50MB limit) + client = StreamingImageClient(total_size_bytes=1 * 1024 * 1024) + monkeypatch.setattr(litellm, "module_level_client", client) + + result = convert_url_to_base64("https://example.com/regular-image.jpg") + + # Verify the image was downloaded successfully + assert result.startswith("data:image/jpeg;base64,") + # Verify all bytes were downloaded + assert client.bytes_downloaded == 1 * 1024 * 1024 + + +@pytest.mark.asyncio +async def test_async_streaming_download_regular_image(monkeypatch): + """ + E2E test: Verify that async streaming download works correctly for regular-sized images. + """ + # Use a 2MB image (well within the 50MB limit) + client = AsyncStreamingImageClient(total_size_bytes=2 * 1024 * 1024) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + result = await async_convert_url_to_base64("https://example.com/regular-image.jpg") + + # Verify the image was downloaded successfully + assert result.startswith("data:image/jpeg;base64,") + # Verify all bytes were downloaded + assert client.bytes_downloaded == 2 * 1024 * 1024 + + +def test_streaming_stops_for_large_image_without_content_length(monkeypatch): + """ + E2E test: Verify that streaming download stops early for images exceeding size limit + when Content-Length header is not provided. + + This is the critical test for the streaming validation fix - it ensures we don't + download the entire file before validating size. + """ + # Simulate a 100MB image without Content-Length header + client = StreamingImageClient( + total_size_bytes=100 * 1024 * 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 it was rejected + assert "exceeds maximum allowed size" in str(excinfo.value) + + # Critical assertion: verify we did NOT download the entire 100MB + # We should have stopped around 50MB (the default limit) + mb_downloaded = client.bytes_downloaded / (1024 * 1024) + assert mb_downloaded < 55 # Allow for chunk size overhead + assert mb_downloaded > 45 # Should download around the limit + + +@pytest.mark.asyncio +async def test_async_streaming_stops_for_large_image_without_content_length(monkeypatch): + """ + E2E test: Verify that async streaming download stops early for large images. + """ + # Simulate a 100MB image without Content-Length header + client = AsyncStreamingImageClient( + total_size_bytes=100 * 1024 * 1024, include_content_length=False + ) + monkeypatch.setattr(litellm, "module_level_aclient", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + await async_convert_url_to_base64("https://example.com/huge-image.jpg") + + # Verify it was rejected + assert "exceeds maximum allowed size" in str(excinfo.value) + + # Critical assertion: verify we did NOT download the entire 100MB + mb_downloaded = client.bytes_downloaded / (1024 * 1024) + assert mb_downloaded < 55 # Allow for chunk size overhead + assert mb_downloaded > 45 # Should download around the limit