feat check check_response_size_is_safe

This commit is contained in:
Ishaan Jaff 2024-07-27 16:53:39 -07:00
parent 5f07afa268
commit 805d04f7f3

View file

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