mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix: extend case normalization to ALL guardrail types (not just tool_permission)
This extends the previous fix to handle capitalized fields across ALL guardrail types, including Presidio, Azure, Lakera, Bedrock, etc. Discovery: - Database investigation revealed the issue affects multiple guardrail types - Found 4 affected guardrails in staging: 3 Presidio + 1 Azure - All had default_action: 'Deny' causing the same validation failures - The initial fix only covered ToolPermissionGuardrailConfigModel Root Cause (Deeper): - LitellmParams inherits from 13+ different guardrail config models - Models use ConfigDict(extra="allow") allowing any field to be set - Users can set default_action/on_disallowed_action on ANY guardrail type - Only ToolPermissionGuardrailConfigModel was validating these fields Solution: - Added field validators to LitellmParams class (parent of all guardrails) - Validators run for ALL guardrail types: Presidio, Azure, Bedrock, Lakera, etc. - Added comprehensive tests covering multiple guardrail types Changes: - litellm/types/guardrails.py: * Added @field_validator for default_action in LitellmParams * Added @field_validator for on_disallowed_action in LitellmParams * Added normalization in LitellmParams.__init__ as backup * Imported field_validator from pydantic - tests/test_litellm/types/test_guardrails_case_normalization.py: * New test file with 7 tests covering multiple guardrail types * Tests verify Presidio, Azure, Tool Permission, Lakera, Bedrock * All tests passing Impact: - Previous fix: Only tool_permission guardrails protected - This fix: ALL guardrail types now protected (13+ types) - Handles both new writes and existing database records - Tested against actual database with Presidio/Azure guardrails Testing: - 7 new cross-guardrail tests (all passing) - 27 existing tool_permission tests (all passing) - Verified fix works for real database records
This commit is contained in:
parent
b5aa7d1838
commit
053e0a9077
2 changed files with 108 additions and 1 deletions
|
|
@ -2,7 +2,7 @@ from datetime import datetime
|
|||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
from typing_extensions import Required, TypedDict
|
||||
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam
|
||||
|
|
@ -671,12 +671,29 @@ class LitellmParams(
|
|||
description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)"
|
||||
)
|
||||
|
||||
@field_validator("default_action", mode="before", check_fields=False)
|
||||
@classmethod
|
||||
def normalize_default_action_litellm_params(cls, v):
|
||||
"""Normalize default_action to lowercase for ALL guardrail types."""
|
||||
if isinstance(v, str):
|
||||
return v.lower()
|
||||
return v
|
||||
|
||||
@field_validator("on_disallowed_action", mode="before", check_fields=False)
|
||||
@classmethod
|
||||
def normalize_on_disallowed_action_litellm_params(cls, v):
|
||||
"""Normalize on_disallowed_action to lowercase for ALL guardrail types."""
|
||||
if isinstance(v, str):
|
||||
return v.lower()
|
||||
return v
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
default_on = kwargs.pop("default_on", None)
|
||||
if default_on is not None:
|
||||
kwargs["default_on"] = default_on
|
||||
else:
|
||||
kwargs["default_on"] = False
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def __contains__(self, key):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
"""
|
||||
Test case normalization in LitellmParams for all guardrail types
|
||||
"""
|
||||
import pytest
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
|
||||
|
||||
class TestLitellmParamsCaseNormalization:
|
||||
"""Test that LitellmParams normalizes case for all guardrail types"""
|
||||
|
||||
def test_presidio_guardrail_with_capitalized_default_action(self):
|
||||
"""Test Presidio guardrail with capitalized default_action"""
|
||||
params = LitellmParams(
|
||||
guardrail="presidio",
|
||||
mode="post_call",
|
||||
default_action="Deny", # Capitalized
|
||||
)
|
||||
assert params.default_action == "deny"
|
||||
|
||||
def test_azure_guardrail_with_capitalized_default_action(self):
|
||||
"""Test Azure guardrail with capitalized default_action"""
|
||||
params = LitellmParams(
|
||||
guardrail="azure/text_moderations",
|
||||
mode="pre_call",
|
||||
default_action="Allow", # Capitalized
|
||||
)
|
||||
assert params.default_action == "allow"
|
||||
|
||||
def test_tool_permission_with_capitalized_fields(self):
|
||||
"""Test tool_permission with capitalized fields"""
|
||||
params = LitellmParams(
|
||||
guardrail="tool_permission",
|
||||
mode="post_call",
|
||||
default_action="DENY", # Uppercase
|
||||
on_disallowed_action="BLOCK", # Uppercase
|
||||
)
|
||||
assert params.default_action == "deny"
|
||||
assert params.on_disallowed_action == "block"
|
||||
|
||||
def test_lakera_with_capitalized_default_action(self):
|
||||
"""Test Lakera guardrail with capitalized default_action"""
|
||||
params = LitellmParams(
|
||||
guardrail="lakera_v2",
|
||||
mode="pre_call",
|
||||
default_action="Deny", # Capitalized
|
||||
)
|
||||
assert params.default_action == "deny"
|
||||
|
||||
def test_bedrock_with_capitalized_default_action(self):
|
||||
"""Test Bedrock guardrail with capitalized default_action"""
|
||||
params = LitellmParams(
|
||||
guardrail="bedrock",
|
||||
mode="pre_call",
|
||||
default_action="Allow", # Capitalized
|
||||
)
|
||||
assert params.default_action == "allow"
|
||||
|
||||
def test_multiple_guardrails_all_normalized(self):
|
||||
"""Test that all guardrail types benefit from normalization"""
|
||||
test_cases = [
|
||||
("presidio", "Deny"),
|
||||
("azure/text_moderations", "Allow"),
|
||||
("tool_permission", "DENY"),
|
||||
("lakera_v2", "allow"), # Already lowercase - should still work
|
||||
("bedrock", "Deny"),
|
||||
]
|
||||
|
||||
for guardrail_type, default_action_input in test_cases:
|
||||
params = LitellmParams(
|
||||
guardrail=guardrail_type,
|
||||
mode="pre_call",
|
||||
default_action=default_action_input,
|
||||
)
|
||||
# Should always be lowercase
|
||||
assert params.default_action.lower() == params.default_action
|
||||
# Should match the expected lowercase value
|
||||
assert params.default_action in ["allow", "deny"]
|
||||
|
||||
def test_on_disallowed_action_all_cases(self):
|
||||
"""Test on_disallowed_action normalization across all cases"""
|
||||
test_cases = ["block", "Block", "BLOCK", "rewrite", "Rewrite", "REWRITE"]
|
||||
|
||||
for action in test_cases:
|
||||
params = LitellmParams(
|
||||
guardrail="tool_permission",
|
||||
mode="post_call",
|
||||
on_disallowed_action=action,
|
||||
)
|
||||
assert params.on_disallowed_action in ["block", "rewrite"]
|
||||
assert params.on_disallowed_action.islower()
|
||||
Loading…
Add table
Reference in a new issue