diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2e91e082bd4..fadf6e4f804 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -103,6 +103,7 @@ class CustomGuardrail(CustomLogger): use_native_during_call_hook: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + supported_streaming_call_types: ClassVar[frozenset[CallTypes]] = frozenset() def __init__( self, @@ -625,6 +626,9 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + def supports_streaming_call_type(self, call_type: CallTypes) -> bool: + return call_type in self.supported_streaming_call_types + def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface(): return self diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..1c680f095d1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1766,6 +1766,23 @@ class ProxyLogging: return False return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES + @staticmethod + def _guardrail_supports_streaming_route( + guardrail: CustomGuardrail, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + + route: Final = user_api_key_dict.request_route + if not route: + return False + call_types: Final = get_call_types_for_route(route) + if not call_types: + return False + return any(guardrail.supports_streaming_call_type(call_type) for call_type in call_types) + @staticmethod def has_post_call_response_headers_callbacks() -> bool: return ProxyLogging._callback_capabilities().has_post_call_response_headers @@ -2744,6 +2761,10 @@ class ProxyLogging: and isinstance(resolved_callback, CustomGuardrail) and resolved_callback.uses_apply_guardrail_interface() and not resolved_callback.mask_response_content + and not ProxyLogging._guardrail_supports_streaming_route( + resolved_callback, + user_api_key_dict, + ) ) else kind )