diff --git a/docs/my-website/docs/proxy/guardrails/quick_start.md b/docs/my-website/docs/proxy/guardrails/quick_start.md index 0016f24ec1..5abe499e30 100644 --- a/docs/my-website/docs/proxy/guardrails/quick_start.md +++ b/docs/my-website/docs/proxy/guardrails/quick_start.md @@ -497,7 +497,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ Run guardrails based on the user-agent header. This is useful for running pre-call checks on OpenWebUI but only masking in logs for Claude CLI. -`default` can be a single mode string or a list of modes. +Both `default` and tag values can be a single mode string or a list of modes. @@ -545,6 +545,29 @@ guardrails: default_on: true ``` + + + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "guardrails_ai-guard" + litellm_params: + guardrail: guardrails_ai + guard_name: "pii_detect" + mode: + tags: + "User-Agent: claude-cli": ["pre_call", "post_call"] # Run both pre and post call for claude-cli + default: "logging_only" # Default to logging only when no tags match + api_base: os.environ/GUARDRAILS_AI_API_BASE + default_on: true +``` + @@ -669,7 +692,7 @@ guardrails: Mode Specification -`default` accepts either a single string or a list of strings. +Both `default` and tag values accept either a single string or a list of strings. ```python from litellm.types.guardrails import Mode @@ -685,6 +708,12 @@ mode = Mode( tags={"User-Agent: claude-cli": "logging_only"}, default=["pre_call", "post_call"] ) + +# Multiple modes on a tag value +mode = Mode( + tags={"User-Agent: claude-cli": ["pre_call", "post_call"]}, + default="logging_only" +) ``` ### `guardrails` Request Parameter diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index 8ed3bfcac4..f07752d5c1 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -50,8 +50,10 @@ class EnterpriseCustomGuardrailHelper: break if matched_mode is not None: - # Tag matched: only run if event_type matches the tag's mode value + # Tag matched: only run if event_type matches the tag's mode value(s) if event_type is not None: + if isinstance(matched_mode, list): + return event_type.value in matched_mode return event_type.value == matched_mode return True diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 269797b987..aed77ab2b3 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -231,8 +231,14 @@ class CustomGuardrail(CustomLogger): event_hook, supported_event_hooks ) elif isinstance(event_hook, Mode): + tag_values_flat: list = [] + for v in event_hook.tags.values(): + if isinstance(v, list): + tag_values_flat.extend(v) + else: + tag_values_flat.append(v) _validate_event_hook_list_is_in_supported_event_hooks( - list(event_hook.tags.values()), supported_event_hooks + tag_values_flat, supported_event_hooks ) if event_hook.default: default_list = ( @@ -466,8 +472,12 @@ class CustomGuardrail(CustomLogger): if isinstance(self.event_hook, list): return event_type.value in self.event_hook if isinstance(self.event_hook, Mode): - if event_type.value in self.event_hook.tags.values(): - return True + for tag_value in self.event_hook.tags.values(): + if isinstance(tag_value, list): + if event_type.value in tag_value: + return True + elif event_type.value == tag_value: + return True if self.event_hook.default: default_list = ( self.event_hook.default diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index dc95ed3314..d5abf5c8fb 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -717,7 +717,9 @@ class BaseLitellmParams( class Mode(BaseModel): - tags: Dict[str, str] = Field(description="Tags for the guardrail mode") + tags: Dict[str, Union[str, List[str]]] = Field( + description="Tags for the guardrail mode" + ) default: Optional[Union[str, List[str]]] = Field( default=None, description="Default mode when no tags match" ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index fb646f7534..67b6c3ea0a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2625,8 +2625,8 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): class GuardrailMode(TypedDict, total=False): - tags: Optional[Dict[str, str]] - default: Optional[str] + tags: Optional[Dict[str, Union[str, List[str]]]] + default: Optional[Union[str, List[str]]] GuardrailStatus = Literal[ diff --git a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py b/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py index f4e06f9f31..8a29e5c1ce 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py @@ -134,6 +134,78 @@ def test_custom_guardrail_with_mode_no_default(monkeypatch): ) +def test_custom_guardrail_with_mode_tag_value_list(monkeypatch): + """Test Mode with tag value as a list of modes (e.g. tags: {"tag": ["pre_call", "post_call"]})""" + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.logging_only, + ], + event_hook=Mode( + tags={"test_tag": ["pre_call", "post_call"]}, + default="logging_only", + ), + default_on=True, + ) + + # Tag matches → pre_call should fire (in tag's list) + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.pre_call, + ) + is True + ) + + # Tag matches → post_call should fire (in tag's list) + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.post_call, + ) + is True + ) + + # Tag matches → logging_only should NOT fire (not in tag's list) + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.logging_only, + ) + is False + ) + + # No tag match → default fires for logging_only + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.logging_only, + ) + is True + ) + + # No tag match → pre_call should NOT fire (not in default) + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.pre_call, + ) + is False + ) + + def test_custom_guardrail_with_mode(monkeypatch): monkeypatch.setattr( "litellm.proxy.proxy_server.premium_user", True