mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix: _process_image_response
This commit is contained in:
parent
610c268e82
commit
22b1a67127
1 changed files with 19 additions and 0 deletions
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue