scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(guardrails): remove TrendAI mask setting

Browse source
This commit is contained in:
Oliver Fei 2026-09-15 13:13:08 -04:00
parent 5701da0884
commit fc65a9a333
4 changed files with 12 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py
View file

@ -28,7 +28,6 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
api_base=litellm_params.api_base,
app_name=settings.app_name,
fallback_on_error=settings.fallback_on_error,
mask_pii=settings.mask_pii,
timeout=settings.timeout,
stream_batch_size=settings.stream_batch_size,
stream_overlap_size=settings.stream_overlap_size,

1
litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py
View file

@ -11,7 +11,6 @@ from pydantic import BaseModel, ConfigDict, Field
class TrendAISettings(BaseModel):
app_name: str | None = None
fallback_on_error: Literal["block", "allow"] = "block"
mask_pii: bool = True
timeout: float = 5.0
stream_batch_size: int = 2048
stream_overlap_size: int = 256

8
litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py
View file

@ -74,7 +74,6 @@ class TrendAIGuardrail(CustomGuardrail):
api_base: str | None = None,
app_name: str | None = None,
fallback_on_error: Literal["block", "allow"] = "block",
mask_pii: bool = True,
timeout: float = 5.0,
stream_batch_size: int = 2048,
stream_overlap_size: int = 256,
@ -114,7 +113,6 @@ class TrendAIGuardrail(CustomGuardrail):
self.api_url: str = _build_apply_guardrails_url(resolved_api_base)
self.app_name: str = app_name or os.environ.get("TMV1_APPLICATION_NAME", "litellm")
self.fallback_on_error: Literal["block", "allow"] = fallback_on_error
self.mask_pii: bool = mask_pii
self.timeout: float = timeout
self.stream_batch_size: int = stream_batch_size
self.stream_overlap_size: int = stream_overlap_size
@ -129,8 +127,8 @@ class TrendAIGuardrail(CustomGuardrail):
event_hook=event_hook,
default_on=default_on,
supported_event_hooks=self.get_supported_event_hooks(),
mask_request_content=mask_pii,
mask_response_content=mask_pii,
mask_request_content=True,
mask_response_content=True,
)
@classmethod
@ -369,7 +367,7 @@ class TrendAIGuardrail(CustomGuardrail):
("TMV1-Client-Name", TMV1_CLIENT_NAME),
("TMV1-Client-Version", litellm_version),
("TMV1-Plugin-Version", PLUGIN_VERSION),
*_optional_header("prefer", "redact-pii,return=representation" if self.mask_pii else None),
("prefer", "redact-pii,return=representation"),
*_optional_header("TMV1-Request-Type", request_type),
)
return MappingProxyType(dict(headers))

11
tests/test_litellm/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py
View file

@ -33,7 +33,6 @@ def _guardrail(
*,
app_name: str | None = None,
fallback_on_error: Literal["block", "allow"] = "block",
mask_pii: bool = True,
timeout: float = 5.0,
stream_batch_size: int = 2048,
stream_overlap_size: int = 256,
@ -48,7 +47,6 @@ def _guardrail(
api_base=api_base,
app_name=app_name,
fallback_on_error=fallback_on_error,
mask_pii=mask_pii,
timeout=timeout,
stream_batch_size=stream_batch_size,
stream_overlap_size=stream_overlap_size,
@ -133,6 +131,15 @@ def test_invalid_timeout_is_rejected() -> None:
_guardrail(timeout=0)
def test_pii_masking_is_not_configurable_on_the_guardrail() -> None:
with pytest.raises(TypeError, match="mask_pii"):
TrendAIGuardrail(
api_key="test-key",
api_base="https://guard.example.com/v3.0/aiSecurity",
mask_pii=False,
)
@pytest.mark.asyncio
async def test_scan_uses_injected_client_and_required_headers() -> None:
captured_request: httpx.Request | None = None