This commit is contained in:
xrwang8 2026-09-16 09:03:14 +00:00 committed by GitHub
commit 07dde25c26
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 48 additions and 1 deletions

View file

@ -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),

View file

@ -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,46 @@ 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="<<KEYWORD>>",
)
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 <<PHONE_NUMBER>> about <<KEYWORD>>"]
finally:
if guardrail in litellm.callbacks:
litellm.callbacks.remove(guardrail)
def test_check_patterns_ssn(self):
"""
Test SSN pattern detection