diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 83c676518d..f9be71c35d 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1,4 +1,5 @@ import re +import sys from fastapi import Request @@ -88,7 +89,11 @@ async def check_if_request_size_is_safe(request: Request) -> bool: request (Request): The incoming request. Returns: - bool: True if the request size is within the limit, False otherwise. + bool: True if the request size is within the limit + + Raises: + ProxyException: If the request size is too large + """ from litellm.proxy.proxy_server import general_settings, premium_user @@ -138,6 +143,46 @@ async def check_if_request_size_is_safe(request: Request) -> bool: return True +async def check_response_size_is_safe(response: Any) -> bool: + """ + Enterprise Only: + - Checks if the response size is within the limit + + Args: + response (Any): The response to check. + + Returns: + bool: True if the response size is within the limit + + Raises: + ProxyException: If the response size is too large + + """ + + from litellm.proxy.proxy_server import general_settings, premium_user + + max_response_size_mb = general_settings.get("max_response_size_mb", None) + if max_response_size_mb is not None: + # Check if premium user + if premium_user is not True: + verbose_proxy_logger.warning( + f"using max_response_size_mb - not checking - this is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}" + ) + return True + + response_size_mb = bytes_to_mb(bytes_value=sys.getsizeof(response)) + verbose_proxy_logger.debug(f"response size in MB={response_size_mb}") + if response_size_mb > max_response_size_mb: + raise ProxyException( + message=f"Response size is too large. Response size is {response_size_mb} MB. Max size is {max_response_size_mb} MB", + type=ProxyErrorTypes.bad_request_error.value, + code=400, + param="content-length", + ) + + return True + + def bytes_to_mb(bytes_value: int): """ Helper to convert bytes to MB