Merge pull request #4928 from BerriAI/litellm_check_response_size

[Feat Enterprise] - check max response size
This commit is contained in:
Ishaan Jaff 2024-07-27 17:03:56 -07:00 committed by GitHub
commit 10e70f842d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 51 additions and 3 deletions

View file

@ -124,6 +124,7 @@ from litellm.proxy.auth.auth_checks import (
get_user_object,
log_to_opentelemetry,
)
from litellm.proxy.auth.auth_utils import check_response_size_is_safe
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.litellm_license import LicenseCheck
from litellm.proxy.auth.model_checks import (
@ -3000,6 +3001,7 @@ async def chat_completion(
**additional_headers,
)
)
await check_response_size_is_safe(response=response)
return response
except RejectedRequestError as e:
@ -3241,7 +3243,7 @@ async def completion(
response_cost=response_cost,
)
)
await check_response_size_is_safe(response=response)
return response
except RejectedRequestError as e:
_data = e.request_data
@ -3491,6 +3493,7 @@ async def embeddings(
call_id=litellm_call_id,
)
)
await check_response_size_is_safe(response=response)
return response
except Exception as e: