diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py index bbe3ded791d..d0f3c5ebe7e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py @@ -38,6 +38,8 @@ def initialize_guardrail( patterns=litellm_params.patterns, blocked_words=litellm_params.blocked_words, blocked_words_file=litellm_params.blocked_words_file, + pattern_redaction_format=litellm_params.pattern_redaction_format, + keyword_redaction_tag=litellm_params.keyword_redaction_tag, event_hook=litellm_params.mode, default_on=litellm_params.default_on or False, categories=getattr(litellm_params, "categories", None), 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 130b0da000b..4f513277f7e 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 @@ -4,13 +4,17 @@ Tests for the Content Filter Guardrail import json import os -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +import litellm import pytest from fastapi import HTTPException +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( + initialize_guardrail, +) from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( ContentFilterGuardrail, ) @@ -19,6 +23,7 @@ from litellm.types.guardrails import ( ContentFilterAction, ContentFilterPattern, GuardrailEventHooks, + LitellmParams, ) from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, @@ -73,6 +78,45 @@ class TestContentFilterGuardrail: assert "secret_project" in guardrail.blocked_words assert guardrail.blocked_words["secret_project"][0] == ContentFilterAction.BLOCK + @pytest.mark.asyncio + async def test_initializer_preserves_custom_redaction_formats(self): + litellm_params = LitellmParams( + guardrail="litellm_content_filter", + mode="pre_call", + patterns=[ + ContentFilterPattern( + pattern_type="regex", + pattern=r"1[3-9]\d{9}", + name="phone_number", + action=ContentFilterAction.MASK, + ) + ], + blocked_words=[ + BlockedWord( + keyword="secret_project", + action=ContentFilterAction.MASK, + ) + ], + pattern_redaction_format="<<{pattern_name}>>", + keyword_redaction_tag="<>", + ) + + guardrail = initialize_guardrail( + litellm_params=litellm_params, + guardrail={"guardrail_name": "test-content-filter"}, + ) + + try: + result = await guardrail.apply_guardrail( + inputs={"texts": ["Call 13912345678 about secret_project"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["Call <> about <>"] + finally: + litellm.callbacks.remove(guardrail) + def test_check_patterns_ssn(self): """ Test SSN pattern detection