diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index fadf6e4f804..f56ff3f1b1b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -626,9 +626,6 @@ 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 1c680f095d1..4788efe9216 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1753,35 +1753,23 @@ class ProxyLogging: return caps @staticmethod - def _stream_requires_guardrail_translation(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 call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES - - @staticmethod - def _guardrail_supports_streaming_route( - guardrail: CustomGuardrail, + def _stream_guardrail_translation_call_type( user_api_key_dict: UserAPIKeyAuth, - ) -> bool: + ) -> CallTypes | None: 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 + return None 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) + return None + call_type: Final = call_types[0] + if call_type not in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: + return None + return call_type @staticmethod def has_post_call_response_headers_callbacks() -> bool: @@ -2744,7 +2732,7 @@ class ProxyLogging: request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) current_response = response - stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) + stream_translation_call_type: Final = ProxyLogging._stream_guardrail_translation_call_type(user_api_key_dict) for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): @@ -2757,14 +2745,11 @@ class ProxyLogging: "apply_guardrail" if ( kind == "override" - and stream_needs_translation + and stream_translation_call_type is not None 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, - ) + and stream_translation_call_type not in resolved_callback.supported_streaming_call_types ) else kind ) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 12817902765..51fc20e7a33 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -51,17 +51,6 @@ def test_has_streaming_callbacks_detects_guardrails(monkeypatch): assert ProxyLogging.has_streaming_callbacks() is True -def test_custom_guardrail_streaming_call_type_capability_is_opt_in(): - from litellm.types.utils import CallTypes - - class _ProtocolAwareGuardrail(CustomGuardrail): - supported_streaming_call_types = frozenset({CallTypes.anthropic_messages}) - - assert CustomGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is False - assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is True - assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.aresponses) is False - - @pytest.mark.asyncio async def test_post_call_response_headers_hook_returns_early_without_callbacks( monkeypatch, @@ -244,27 +233,28 @@ def _streaming_logging_obj(): ) -def test_stream_requires_guardrail_translation_route_detection(): +def test_stream_guardrail_translation_call_type_route_detection(): from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes assert ( - ProxyLogging._stream_requires_guardrail_translation( + ProxyLogging._stream_guardrail_translation_call_type( UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages") ) - is True + is CallTypes.anthropic_messages ) assert ( - ProxyLogging._stream_requires_guardrail_translation( + ProxyLogging._stream_guardrail_translation_call_type( UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions") ) - is False + is None ) - assert ProxyLogging._stream_requires_guardrail_translation(UserAPIKeyAuth(api_key="sk-1234")) is False + assert ProxyLogging._stream_guardrail_translation_call_type(UserAPIKeyAuth(api_key="sk-1234")) is None assert ( - ProxyLogging._stream_requires_guardrail_translation( + ProxyLogging._stream_guardrail_translation_call_type( UserAPIKeyAuth(api_key="sk-1234", request_route="/route/without/call/types") ) - is False + is None )