diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 883329c9fa8..2d66a280663 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -773,9 +773,6 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def rewrites_streamed_output(self) -> bool: - return self.mask_response_content - def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 85eb50c78e7..722f96ef814 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1947,15 +1947,6 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) - def rewrites_streamed_output(self) -> bool: - return ( - super().rewrites_streamed_output() - or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) - or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) - or any(action == ContentFilterAction.MASK for _, _, action in self.category_keywords.values()) - or any(action == ContentFilterAction.MASK for _, _, action in self.always_block_category_keywords.values()) - ) - async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 73020fe3e6f..be55ac47bde 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3068,59 +3068,3 @@ class TestContentFilterToolCallArguments: request_data={}, input_type="response", ) - - -class TestRewritesStreamedOutput: - def test_block_only_rules_do_not_rewrite(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.BLOCK)], - blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], - ) - - assert guardrail.rewrites_streamed_output() is False - - def test_mask_blocked_word_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - blocked_words=[BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK)], - ) - - assert guardrail.rewrites_streamed_output() is True - - def test_mask_pattern_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)], - ) - - assert guardrail.rewrites_streamed_output() is True - - def test_mask_response_content_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], - mask_response_content=True, - ) - - assert guardrail.rewrites_streamed_output() is True - - @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) - def test_category_keywords_follow_the_category_action(self, action, expected): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - categories=[{"category": "bias_gender", "enabled": True, "action": action}], - ) - - assert guardrail.category_keywords and not guardrail.always_block_category_keywords - assert guardrail.rewrites_streamed_output() is expected - - @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) - def test_always_block_category_keywords_follow_the_category_action(self, action, expected): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - categories=[{"category": "age_discrimination", "enabled": True, "action": action}], - ) - - assert guardrail.always_block_category_keywords and not guardrail.category_keywords - assert guardrail.rewrites_streamed_output() is expected