From 6d49c6856ceab74a7869d254aa952e0c290d686e Mon Sep 17 00:00:00 2001 From: pjdurden Date: Wed, 19 Aug 2026 10:30:37 -0500 Subject: [PATCH 1/4] fix: BerriAI/litellm#36535 - [Bug]: sensitive_data_routing documented but not recognized as guardrail Signed-off-by: pjdurden --- litellm/exceptions.py | 5 +- .../sensitive_data_routing/__init__.py | 67 +++++++ .../sensitive_data_routing.py | 144 ++++++++++++++++ litellm/proxy/hooks/sensitive_data_routing.py | 10 +- litellm/proxy/utils.py | 1 + litellm/types/guardrails.py | 1 + .../guardrail_hooks/sensitive_data_routing.py | 42 +++++ .../test_sensitive_data_routing.py | 163 ++++++++++++++++++ .../hooks/test_sensitive_data_routing.py | 51 +++++- 9 files changed, 479 insertions(+), 5 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index ac461a2472e..cc78b36c029 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1326,7 +1326,8 @@ class SensitiveDataRouteException(Exception): The proxy catches this exception and: 1. Reroutes the current request to the specified model 2. When sticky_session_routing is True, stores the routing decision in session - cache so all subsequent requests in the same session are routed to the same model + cache so all subsequent requests in the same session are routed to the same model, + for session_ttl_seconds when the guardrail sets one and the proxy-wide default otherwise """ def __init__( @@ -1337,11 +1338,13 @@ class SensitiveDataRouteException(Exception): detection_info: dict[str, Any] | None = None, message: str | None = None, sticky_session_routing: bool = True, + session_ttl_seconds: int | None = None, ): self.route_to_model = route_to_model self.session_id = session_id self.guardrail_name = guardrail_name self.detection_info = detection_info or {} self.sticky_session_routing = sticky_session_routing + self.session_ttl_seconds = session_ttl_seconds self.message = message or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}" super().__init__(self.message) diff --git a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py new file mode 100644 index 00000000000..1a138b29fe0 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py @@ -0,0 +1,67 @@ +"""Built-in Sensitive Data Routing guardrail: reroutes sensitive prompts to an on-premise model.""" + +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingGuardrailConfigModel, +) + +from .sensitive_data_routing import SensitiveDataRoutingGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def _to_event_hook( + mode: str | list[str] | Mode | None, +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None: + if mode is None or isinstance(mode, Mode): + return mode + if isinstance(mode, str): + return GuardrailEventHooks(mode) + return [GuardrailEventHooks(single_mode) for single_mode in mode] + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +) -> SensitiveDataRoutingGuardrail: + """Initialize the Sensitive Data Routing guardrail from config.""" + import litellm + + guardrail_name: Final = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("sensitive_data_routing guardrail requires a guardrail_name") + + config: Final = SensitiveDataRoutingGuardrailConfigModel.model_validate(litellm_params.model_dump()) + if not config.on_premise_model: + raise ValueError("sensitive_data_routing guardrail requires 'on_premise_model'") + + instance: Final = SensitiveDataRoutingGuardrail( + on_premise_model=config.on_premise_model, + guardrail_name=guardrail_name, + prebuilt_patterns=config.prebuilt_patterns, + regex_patterns=config.regex_patterns, + keywords=config.keywords, + sticky_session=config.sticky_session, + session_ttl_seconds=config.session_ttl_seconds, + event_hook=_to_event_hook(litellm_params.mode), + default_on=bool(litellm_params.default_on), + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.SENSITIVE_DATA_ROUTING.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.SENSITIVE_DATA_ROUTING.value: SensitiveDataRoutingGuardrail, +} + +__all__ = [ + "SensitiveDataRoutingGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py new file mode 100644 index 00000000000..7a6c7ee09e1 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py @@ -0,0 +1,144 @@ +""" +Built-in Sensitive Data Routing guardrail. + +Detects sensitive data with prebuilt regex patterns, custom regex and keyword matching, +then reroutes the request to an on-premise model instead of blocking or redacting it. +""" + +import re +from collections.abc import Sequence +from dataclasses import dataclass +from re import Pattern +from typing import TYPE_CHECKING, Final, Literal, Optional + +from litellm.exceptions import SensitiveDataRouteException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + get_session_id_from_request_data, + log_guardrail_information, +) +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import ( + get_compiled_pattern, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + DEFAULT_SESSION_TTL_SECONDS, + SensitiveDataRoutingGuardrailConfigModel, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +SensitiveDataDetectorKind = Literal["prebuilt_pattern", "regex_pattern", "keyword"] + + +@dataclass(frozen=True, slots=True) +class SensitiveDataDetector: + kind: SensitiveDataDetectorKind + rule: str + matcher: Pattern[str] + + +def build_detectors( + prebuilt_patterns: Sequence[str] | None, + regex_patterns: Sequence[str] | None, + keywords: Sequence[str] | None, +) -> tuple[SensitiveDataDetector, ...]: + return ( + *( + SensitiveDataDetector("prebuilt_pattern", name, get_compiled_pattern(name)) + for name in prebuilt_patterns or () + ), + *( + SensitiveDataDetector("regex_pattern", pattern, re.compile(pattern, re.IGNORECASE)) + for pattern in regex_patterns or () + ), + *( + SensitiveDataDetector("keyword", keyword, re.compile(re.escape(keyword), re.IGNORECASE)) + for keyword in keywords or () + ), + ) + + +class SensitiveDataRoutingGuardrail(CustomGuardrail): + """ + Reroutes a request to an on-premise model when sensitive data is detected. + + Runs locally with no external API call, and never blocks or redacts: the prompt is + forwarded unchanged to the on-premise model. When sticky_session is enabled and the + request carries a session id, the whole session stays pinned to that model. + """ + + def __init__( + self, + on_premise_model: str, + guardrail_name: str | None = None, + prebuilt_patterns: Sequence[str] | None = None, + regex_patterns: Sequence[str] | None = None, + keywords: Sequence[str] | None = None, + sticky_session: bool = True, + session_ttl_seconds: int = DEFAULT_SESSION_TTL_SECONDS, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, + default_on: bool = False, + ) -> None: + if not on_premise_model: + raise ValueError("sensitive_data_routing guardrail requires 'on_premise_model'") + + detectors: Final = build_detectors(prebuilt_patterns, regex_patterns, keywords) + if not detectors: + raise ValueError( + "sensitive_data_routing guardrail requires at least one of " + "'prebuilt_patterns', 'regex_patterns' or 'keywords'" + ) + + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=self.get_supported_event_hooks(), + event_hook=event_hook or GuardrailEventHooks.pre_call, + default_on=default_on, + ) + self.on_premise_model = on_premise_model + self.detectors = detectors + self.sticky_session = sticky_session + self.session_ttl_seconds = session_ttl_seconds + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [GuardrailEventHooks.pre_call] + + @staticmethod + def get_config_model() -> type[GuardrailConfigModel] | None: + return SensitiveDataRoutingGuardrailConfigModel + + def detect(self, texts: Sequence[str]) -> SensitiveDataDetector | None: + return next( + (detector for detector in self.detectors for text in texts if text and detector.matcher.search(text)), + None, + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + detected: Final = self.detect(inputs.get("texts") or ()) + if detected is None: + return inputs + + session_id: Final = get_session_id_from_request_data(request_data) + raise SensitiveDataRouteException( + route_to_model=self.on_premise_model, + session_id=session_id or "", + guardrail_name=self.guardrail_name, + detection_info={"detection_type": detected.kind, "rule": detected.rule}, + sticky_session_routing=self.sticky_session and session_id is not None, + session_ttl_seconds=self.session_ttl_seconds, + ) diff --git a/litellm/proxy/hooks/sensitive_data_routing.py b/litellm/proxy/hooks/sensitive_data_routing.py index 0909f36dbe5..13434ec4647 100644 --- a/litellm/proxy/hooks/sensitive_data_routing.py +++ b/litellm/proxy/hooks/sensitive_data_routing.py @@ -123,6 +123,7 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): model: str, user_api_key_dict: UserAPIKeyAuth | None = None, guardrail_name: str | None = None, + ttl: int | None = None, ) -> None: """ Store a routing override for a session. @@ -130,15 +131,18 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): Called by guardrails when they detect sensitive data and want to route the session to a specific model. The override is scoped to the requesting principal so sessions from different tenants cannot collide. + A guardrail that configures its own session TTL passes it as ``ttl``; + otherwise the proxy-wide default applies. """ cache_key: Final = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict)) + effective_ttl: Final = ttl if ttl is not None else self.ttl verbose_proxy_logger.info( "SensitiveDataRoutingHandler: Setting session routing session_id=%s model=%s guardrail=%s ttl=%s", session_id, model, guardrail_name, - self.ttl, + effective_ttl, ) if self.internal_usage_cache.dual_cache.redis_cache is not None: @@ -146,7 +150,7 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): await self.internal_usage_cache.dual_cache.redis_cache.async_set_cache( key=cache_key, value=model, - ttl=self.ttl, + ttl=effective_ttl, ) except Exception as e: log_redis_failure( @@ -159,7 +163,7 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): await self.internal_usage_cache.async_set_cache( key=cache_key, value=model, - ttl=self.ttl, + ttl=effective_ttl, litellm_parent_otel_span=None, local_only=True, ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 302f55eadf4..977857127bf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2751,6 +2751,7 @@ class ProxyLogging: model=exc.route_to_model, user_api_key_dict=user_api_key_dict, guardrail_name=exc.guardrail_name, + ttl=exc.session_ttl_seconds, ) else: verbose_proxy_logger.warning( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index b3b57d53bdf..5301ef07462 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -153,6 +153,7 @@ class SupportedGuardrailIntegrations(Enum): LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" DECISION_MODEL = "decision_model" + SENSITIVE_DATA_ROUTING = "sensitive_data_routing" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py new file mode 100644 index 00000000000..0c71d41101c --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py @@ -0,0 +1,42 @@ +"""Types for the built-in Sensitive Data Routing guardrail.""" + +from typing import Final + +from pydantic import Field + +from .base import GuardrailConfigModel + +DEFAULT_SESSION_TTL_SECONDS: Final = 14400 + + +class SensitiveDataRoutingGuardrailConfigModel(GuardrailConfigModel): + """Configuration for the built-in Sensitive Data Routing guardrail.""" + + on_premise_model: str | None = Field( + default=None, + description="Model group (from model_list) to route the request to when sensitive data is detected.", + ) + prebuilt_patterns: list[str] | None = Field( + default=None, + description="Prebuilt pattern names to match (e.g. us_ssn, credit_card, email).", + ) + regex_patterns: list[str] | None = Field( + default=None, + description="Custom regular expressions; a match in any message reroutes the request.", + ) + keywords: list[str] | None = Field( + default=None, + description="Case-insensitive keywords; a match in any message reroutes the request.", + ) + sticky_session: bool = Field( + default=True, + description="Keep the whole session on the on-premise model after sensitive data is first detected.", + ) + session_ttl_seconds: int = Field( + default=DEFAULT_SESSION_TTL_SECONDS, + description="How long a session stays pinned to the on-premise model after detection.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Sensitive Data Routing" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py new file mode 100644 index 00000000000..78508278481 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py @@ -0,0 +1,163 @@ +"""Tests for the built-in Sensitive Data Routing guardrail.""" + +import pytest + +from litellm.exceptions import SensitiveDataRouteException +from litellm.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingGuardrail, + initialize_guardrail, +) +from litellm.proxy.guardrails.guardrail_registry import ( + IN_MEMORY_GUARDRAIL_HANDLER, + guardrail_class_registry, + guardrail_initializer_registry, +) +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams + +DOCUMENTED_LITELLM_PARAMS = { + "guardrail": "sensitive_data_routing", + "mode": "pre_call", + "default_on": True, + "on_premise_model": "on-prem-model", + "prebuilt_patterns": ["us_ssn", "credit_card", "email"], + "regex_patterns": [r"project\s+titan"], + "keywords": ["confidential", "internal only"], + "sticky_session": True, + "session_ttl_seconds": 14400, +} + + +def make_guardrail(**overrides) -> SensitiveDataRoutingGuardrail: + litellm_params = LitellmParams(**{**DOCUMENTED_LITELLM_PARAMS, **overrides}) + return initialize_guardrail(litellm_params, {"guardrail_name": "sensitive-data-routing"}) + + +class TestSensitiveDataRoutingGuardrailRegistration: + def test_guardrail_type_is_registered(self): + """Regression: `guardrail: sensitive_data_routing` used to raise "Unsupported guardrail".""" + assert "sensitive_data_routing" in guardrail_initializer_registry + assert guardrail_class_registry["sensitive_data_routing"] is SensitiveDataRoutingGuardrail + + def test_documented_config_initializes_through_the_registry(self): + """The config.yaml from the docs initializes instead of failing proxy startup.""" + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( + { + "guardrail_name": "docs-sensitive-data-routing", + "litellm_params": dict(DOCUMENTED_LITELLM_PARAMS), + } + ) + + assert guardrail is not None + callback = IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail[guardrail["guardrail_id"]] + assert isinstance(callback, SensitiveDataRoutingGuardrail) + assert callback.on_premise_model == "on-prem-model" + assert callback.event_hook == GuardrailEventHooks.pre_call + + def test_on_premise_model_is_required(self): + with pytest.raises(ValueError, match="on_premise_model"): + make_guardrail(on_premise_model=None) + + def test_at_least_one_detector_is_required(self): + with pytest.raises(ValueError, match="prebuilt_patterns"): + make_guardrail(prebuilt_patterns=None, regex_patterns=None, keywords=None) + + def test_unknown_prebuilt_pattern_fails_fast(self): + with pytest.raises(ValueError, match="Unknown pattern name"): + make_guardrail(prebuilt_patterns=["not_a_real_pattern"]) + + +class TestSensitiveDataRoutingGuardrailDetection: + @pytest.mark.asyncio + async def test_clean_request_is_untouched(self): + guardrail = make_guardrail() + inputs = {"texts": ["What is the capital of France?"]} + + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + assert result["texts"] == ["What is the capital of France?"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "text, detection_type, rule", + [ + ("My SSN is 123-45-6789, summarize my record", "prebuilt_pattern", "us_ssn"), + ("notes on project titan", "regex_pattern", r"project\s+titan"), + ("This is INTERNAL ONLY", "keyword", "internal only"), + ], + ) + async def test_detection_reroutes_to_the_on_premise_model(self, text, detection_type, rule): + guardrail = make_guardrail() + + with pytest.raises(SensitiveDataRouteException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": [text]}, + request_data={"model": "cloud-model", "metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + exc = exc_info.value + assert exc.route_to_model == "on-prem-model" + assert exc.session_id == "abc-123" + assert exc.guardrail_name == "sensitive-data-routing" + assert exc.detection_info == {"detection_type": detection_type, "rule": rule} + assert exc.sticky_session_routing is True + assert exc.session_ttl_seconds == 14400 + + @pytest.mark.asyncio + async def test_detection_never_leaks_the_matched_value(self): + guardrail = make_guardrail() + + with pytest.raises(SensitiveDataRouteException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["My SSN is 123-45-6789"]}, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + assert "123-45-6789" not in str(exc_info.value.detection_info) + + @pytest.mark.asyncio + async def test_request_without_session_id_reroutes_without_pinning(self): + """Docs: turns without a session id are still routed, but never pinned.""" + guardrail = make_guardrail() + + with pytest.raises(SensitiveDataRouteException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["My SSN is 123-45-6789"]}, + request_data={"model": "cloud-model"}, + input_type="request", + ) + + assert exc_info.value.route_to_model == "on-prem-model" + assert exc_info.value.sticky_session_routing is False + + @pytest.mark.asyncio + async def test_sticky_session_disabled_does_not_pin(self): + guardrail = make_guardrail(sticky_session=False) + + with pytest.raises(SensitiveDataRouteException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["My SSN is 123-45-6789"]}, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + assert exc_info.value.sticky_session_routing is False + + @pytest.mark.asyncio + async def test_responses_are_not_scanned(self): + """The guardrail only picks the model, so it has nothing to do on the response.""" + guardrail = make_guardrail() + inputs = {"texts": ["My SSN is 123-45-6789"]} + + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="response", + ) + + assert result["texts"] == ["My SSN is 123-45-6789"] diff --git a/tests/unit/proxy/hooks/test_sensitive_data_routing.py b/tests/unit/proxy/hooks/test_sensitive_data_routing.py index fbe97e2917a..e045392f32d 100644 --- a/tests/unit/proxy/hooks/test_sensitive_data_routing.py +++ b/tests/unit/proxy/hooks/test_sensitive_data_routing.py @@ -739,7 +739,56 @@ class TestProxyHandleSensitiveDataRouteException: ) @pytest.mark.asyncio - async def test_non_sticky_routing_does_not_persist_override(self, proxy_logging, routing_hook): + async def test_guardrail_session_ttl_overrides_the_proxy_default( + self, proxy_logging, routing_hook + ): + proxy_logging.proxy_hook_mapping["sensitive_data_routing"] = routing_hook + exc = SensitiveDataRouteException( + route_to_model="on-premise-model", + session_id="sess-ttl", + guardrail_name="sensitive-data-routing", + sticky_session_routing=True, + session_ttl_seconds=14400, + ) + + await proxy_logging._handle_sensitive_data_route_exception( + exc, + {"model": "gpt-4", "metadata": {"session_id": "sess-ttl"}}, + UserAPIKeyAuth(api_key="tenant-a"), + ) + + cache_key = routing_hook._make_cache_key("sess-ttl", "tenant-a") + assert routing_hook.internal_usage_cache._ttls[cache_key] == 14400 + assert routing_hook.ttl == DEFAULT_SENSITIVE_ROUTING_TTL + + @pytest.mark.asyncio + async def test_routing_without_guardrail_ttl_uses_the_proxy_default( + self, proxy_logging, routing_hook + ): + proxy_logging.proxy_hook_mapping["sensitive_data_routing"] = routing_hook + exc = SensitiveDataRouteException( + route_to_model="on-premise-model", + session_id="sess-default-ttl", + guardrail_name="pii", + sticky_session_routing=True, + ) + + await proxy_logging._handle_sensitive_data_route_exception( + exc, + {"model": "gpt-4", "metadata": {"session_id": "sess-default-ttl"}}, + UserAPIKeyAuth(api_key="tenant-a"), + ) + + cache_key = routing_hook._make_cache_key("sess-default-ttl", "tenant-a") + assert ( + routing_hook.internal_usage_cache._ttls[cache_key] + == DEFAULT_SENSITIVE_ROUTING_TTL + ) + + @pytest.mark.asyncio + async def test_non_sticky_routing_does_not_persist_override( + self, proxy_logging, routing_hook + ): proxy_logging.proxy_hook_mapping["sensitive_data_routing"] = routing_hook exc = SensitiveDataRouteException( route_to_model="on-premise-model", From 03b391d3b4ba9299dbc4c612739b284720f256f8 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Thu, 20 Aug 2026 21:41:11 -0500 Subject: [PATCH 2/4] fix: reject non-positive session_ttl_seconds on sensitive_data_routing A non-positive TTL either expires the in-memory pin immediately or makes Redis reject the SETEX, so sticky_session silently never pinned. Bound the field at config-parse time instead of failing quietly at request time. Signed-off-by: pjdurden --- .../guardrail_hooks/sensitive_data_routing.py | 1 + .../test_sensitive_data_routing.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py index 0c71d41101c..65aa4c077cf 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py @@ -34,6 +34,7 @@ class SensitiveDataRoutingGuardrailConfigModel(GuardrailConfigModel): ) session_ttl_seconds: int = Field( default=DEFAULT_SESSION_TTL_SECONDS, + ge=1, description="How long a session stays pinned to the on-premise model after detection.", ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py index 78508278481..b913d99b97d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py @@ -1,6 +1,7 @@ """Tests for the built-in Sensitive Data Routing guardrail.""" import pytest +from pydantic import ValidationError from litellm.exceptions import SensitiveDataRouteException from litellm.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( @@ -13,6 +14,9 @@ from litellm.proxy.guardrails.guardrail_registry import ( guardrail_initializer_registry, ) from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingGuardrailConfigModel, +) DOCUMENTED_LITELLM_PARAMS = { "guardrail": "sensitive_data_routing", @@ -65,6 +69,17 @@ class TestSensitiveDataRoutingGuardrailRegistration: with pytest.raises(ValueError, match="Unknown pattern name"): make_guardrail(prebuilt_patterns=["not_a_real_pattern"]) + @pytest.mark.parametrize("ttl", [0, -1]) + def test_non_positive_session_ttl_is_rejected(self, ttl): + """A non-positive TTL expires the pin immediately, so reject it at config time + instead of silently never pinning the session.""" + with pytest.raises(ValidationError): + SensitiveDataRoutingGuardrailConfigModel( + on_premise_model="on-prem-model", + keywords=["confidential"], + session_ttl_seconds=ttl, + ) + class TestSensitiveDataRoutingGuardrailDetection: @pytest.mark.asyncio From 8a9fef3877c89299a3c3da13af7f77638813fe60 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Thu, 20 Aug 2026 22:02:40 -0500 Subject: [PATCH 3/4] chore: satisfy the LIT001 mutable-annotation gate The config-model pattern lists become Sequence[str], which the guardrail already accepted and which generates an identical JSON schema. The remaining four are fixed by CustomGuardrail's own signatures, so they carry a mutable-ok reason instead. Comments are kept short enough that ruff format leaves them on the annotation line, which is where the gate reads them. Signed-off-by: pjdurden --- .../guardrail_hooks/sensitive_data_routing/__init__.py | 5 +++-- .../sensitive_data_routing/sensitive_data_routing.py | 6 +++--- .../guardrails/guardrail_hooks/sensitive_data_routing.py | 7 ++++--- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py index 1a138b29fe0..d38b851680a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py @@ -1,5 +1,6 @@ """Built-in Sensitive Data Routing guardrail: reroutes sensitive prompts to an on-premise model.""" +from collections.abc import Sequence from typing import TYPE_CHECKING, Final from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations @@ -14,8 +15,8 @@ if TYPE_CHECKING: def _to_event_hook( - mode: str | list[str] | Mode | None, -) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None: + mode: str | Sequence[str] | Mode | None, +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None: # mutable-ok: base type if mode is None or isinstance(mode, Mode): return mode if isinstance(mode, str): diff --git a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py index 7a6c7ee09e1..4f2aa2c1e0e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py @@ -80,7 +80,7 @@ class SensitiveDataRoutingGuardrail(CustomGuardrail): keywords: Sequence[str] | None = None, sticky_session: bool = True, session_ttl_seconds: int = DEFAULT_SESSION_TTL_SECONDS, - event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, # mutable-ok: base type default_on: bool = False, ) -> None: if not on_premise_model: @@ -105,7 +105,7 @@ class SensitiveDataRoutingGuardrail(CustomGuardrail): self.session_ttl_seconds = session_ttl_seconds @classmethod - def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: base type return [GuardrailEventHooks.pre_call] @staticmethod @@ -122,7 +122,7 @@ class SensitiveDataRoutingGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict, # mutable-ok: proxy request body, owned by the caller input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py index 65aa4c077cf..326dae595b3 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py @@ -1,5 +1,6 @@ """Types for the built-in Sensitive Data Routing guardrail.""" +from collections.abc import Sequence from typing import Final from pydantic import Field @@ -16,15 +17,15 @@ class SensitiveDataRoutingGuardrailConfigModel(GuardrailConfigModel): default=None, description="Model group (from model_list) to route the request to when sensitive data is detected.", ) - prebuilt_patterns: list[str] | None = Field( + prebuilt_patterns: Sequence[str] | None = Field( default=None, description="Prebuilt pattern names to match (e.g. us_ssn, credit_card, email).", ) - regex_patterns: list[str] | None = Field( + regex_patterns: Sequence[str] | None = Field( default=None, description="Custom regular expressions; a match in any message reroutes the request.", ) - keywords: list[str] | None = Field( + keywords: Sequence[str] | None = Field( default=None, description="Case-insensitive keywords; a match in any message reroutes the request.", ) From d4c30627287e512ca935344566ac34a1e6b92bb3 Mon Sep 17 00:00:00 2001 From: pjdurden Date: Mon, 12 Oct 2026 15:21:39 -0500 Subject: [PATCH 4/4] fix(guardrails): scan tool-call arguments in sensitive_data_routing Assistant tool_calls reach the guardrail through inputs["tool_calls"], not texts, so an SSN inside function.arguments was forwarded to the cloud model without triggering a reroute. Scan the raw argument string and its decoded JSON string values alongside the message text. Also replaces the two-for generator in detect() with a plain loop and declares SensitiveDataDetectorKind as a TypeAlias, for LIT014/LIT010. --- .../sensitive_data_routing.py | 46 +++++++++++++++---- .../test_sensitive_data_routing.py | 46 +++++++++++++++++++ 2 files changed, 84 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py index 4f2aa2c1e0e..d95560cc02e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py @@ -5,11 +5,12 @@ Detects sensitive data with prebuilt regex patterns, custom regex and keyword ma then reroutes the request to an on-premise model instead of blocking or redacting it. """ +import json import re -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from dataclasses import dataclass from re import Pattern -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, TypeAlias from litellm.exceptions import SensitiveDataRouteException from litellm.integrations.custom_guardrail import ( @@ -31,7 +32,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -SensitiveDataDetectorKind = Literal["prebuilt_pattern", "regex_pattern", "keyword"] +SensitiveDataDetectorKind: TypeAlias = Literal["prebuilt_pattern", "regex_pattern", "keyword"] @dataclass(frozen=True, slots=True) @@ -62,6 +63,33 @@ def build_detectors( ) +def _string_leaves(value: object) -> Iterator[str]: + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for key, item in value.items(): + yield from _string_leaves(key) + yield from _string_leaves(item) + elif isinstance(value, list): + for item in value: + yield from _string_leaves(item) + + +def tool_call_argument_texts(tool_calls: Sequence[object] | None) -> Iterator[str]: + """Yield tool-call argument strings, plus their decoded JSON string values so escapes cannot hide a match.""" + for tool_call in tool_calls or (): + function = tool_call.get("function") if isinstance(tool_call, dict) else getattr(tool_call, "function", None) + arguments = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) + if not isinstance(arguments, str) or not arguments: + continue + yield arguments + try: + parsed = json.loads(arguments) + except (json.JSONDecodeError, TypeError, ValueError): + continue + yield from _string_leaves(parsed) + + class SensitiveDataRoutingGuardrail(CustomGuardrail): """ Reroutes a request to an on-premise model when sensitive data is detected. @@ -113,10 +141,10 @@ class SensitiveDataRoutingGuardrail(CustomGuardrail): return SensitiveDataRoutingGuardrailConfigModel def detect(self, texts: Sequence[str]) -> SensitiveDataDetector | None: - return next( - (detector for detector in self.detectors for text in texts if text and detector.matcher.search(text)), - None, - ) + for detector in self.detectors: + if any(text and detector.matcher.search(text) for text in texts): + return detector + return None @log_guardrail_information async def apply_guardrail( @@ -129,7 +157,9 @@ class SensitiveDataRoutingGuardrail(CustomGuardrail): if input_type != "request": return inputs - detected: Final = self.detect(inputs.get("texts") or ()) + detected: Final = self.detect( + [*(inputs.get("texts") or ()), *tool_call_argument_texts(inputs.get("tool_calls"))] + ) if detected is None: return inputs diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py index b913d99b97d..bdd130a0a84 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py @@ -122,6 +122,52 @@ class TestSensitiveDataRoutingGuardrailDetection: assert exc.sticky_session_routing is True assert exc.session_ttl_seconds == 14400 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "arguments", + [ + '{"ssn": "123-45-6789"}', + '{"ssn": "123\\u002d45\\u002d6789"}', + '{"record": {"ids": ["123-45-6789"]}}', + "not json 123-45-6789", + ], + ) + async def test_tool_call_arguments_are_scanned(self, arguments): + guardrail = make_guardrail(prebuilt_patterns=["us_ssn"], regex_patterns=None, keywords=None) + inputs = { + "texts": ["look this person up"], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + } + + with pytest.raises(SensitiveDataRouteException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + assert exc_info.value.detection_info == {"detection_type": "prebuilt_pattern", "rule": "us_ssn"} + + @pytest.mark.asyncio + async def test_clean_tool_call_arguments_are_untouched(self): + guardrail = make_guardrail() + inputs = { + "texts": ["weather please"], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "weather", "arguments": '{"city": "Paris"}'}} + ], + } + + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"metadata": {"session_id": "abc-123"}}, + input_type="request", + ) + + assert result is inputs + @pytest.mark.asyncio async def test_detection_never_leaks_the_matched_value(self): guardrail = make_guardrail()