fix: _process_image_response

This commit is contained in:
Ishaan Jaffer 2026-01-16 15:27:02 -08:00
parent 610c268e82
commit 22b1a67127

View file

@ -9,6 +9,7 @@ from httpx import Response
import litellm
from litellm import verbose_logger
from litellm.caching.caching import InMemoryCache
from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
MAX_IMGS_IN_MEMORY = 10
@ -21,7 +22,25 @@ def _process_image_response(response: Response, url: str) -> str:
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 MAX_IMAGE_URL_DOWNLOAD_SIZE_MB > 0 and 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}"
)
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)
if MAX_IMAGE_URL_DOWNLOAD_SIZE_MB > 0 and 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}"
)
base64_image = base64.b64encode(image_bytes).decode("utf-8")
image_type = response.headers.get("Content-Type")