scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

fix(proxy/utils.py): run guardrails before running other logging hooks on "async_post_call_success_hook"

Browse source 
Closes LIT-1152
This commit is contained in:
Krrish Dholakia 2025-09-30 12:50:21 -07:00
parent 39d1d2a379
commit 6006320251
1 changed files with 31 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
litellm/proxy/utils.py
View file

@ -1395,9 +1395,12 @@ class ProxyLogging:
3. /image/generation
4. /files
"""
from litellm.types.guardrails import GuardrailEventHooks
for callback in litellm.callbacks:
try:
guardrail_callbacks: List[CustomGuardrail] = []
other_callbacks: List[CustomLogger] = []
try:
for callback in litellm.callbacks:
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
@ -1407,36 +1410,37 @@ class ProxyLogging:
_callback = callback # type: ignore
if _callback is not None:
if isinstance(_callback, CustomGuardrail):
guardrail_callbacks.append(_callback)
else:
other_callbacks.append(_callback)
############## Handle Guardrails ########################################
#############################################################################
if isinstance(callback, CustomGuardrail):
# Main - V2 Guardrails implementation
from litellm.types.guardrails import GuardrailEventHooks
if (
callback.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.post_call
)
is not True
):
continue
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if (
callback.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.post_call
)
is not True
):
continue
await callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
)
await callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
)
############ Handle CustomLogger ###############################
#################################################################
elif isinstance(_callback, CustomLogger):
await _callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
)
except Exception as e:
raise e
############ Handle CustomLogger ###############################
#################################################################
for callback in other_callbacks:
await callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, data=data, response=response
)
except Exception as e:
raise e
return response
async def async_post_call_streaming_hook(