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..c33c37a2f28 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/__init__.py @@ -0,0 +1,75 @@ +"""Sensitive Data Routing guardrail: reroutes requests with sensitive data to an on-premise model.""" + +from typing import TYPE_CHECKING, Any, List + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .sensitive_data_routing import SensitiveDataRoutingGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def _get_param( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + key: str, + default: Any = None, +) -> Any: + value = getattr(litellm_params, key, None) + if value is not None: + return value + raw = guardrail.get("litellm_params") + if isinstance(raw, dict) and key in raw: + return raw[key] + return default + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +) -> SensitiveDataRoutingGuardrail: + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("sensitive_data_routing guardrail requires a guardrail_name") + + on_premise_model = _get_param(litellm_params, guardrail, "on_premise_model") + if not on_premise_model: + raise ValueError( + "sensitive_data_routing guardrail requires 'on_premise_model' (the model_list " + "name to route sensitive requests to)" + ) + + instance = SensitiveDataRoutingGuardrail( + guardrail_name=guardrail_name, + on_premise_model=on_premise_model, + prebuilt_patterns=_get_param(litellm_params, guardrail, "prebuilt_patterns"), + regex_patterns=_get_param(litellm_params, guardrail, "regex_patterns"), + keywords=_get_param(litellm_params, guardrail, "keywords"), + sticky_session=bool( + _get_param(litellm_params, guardrail, "sticky_session", True) + ), + session_ttl_seconds=int( + _get_param(litellm_params, guardrail, "session_ttl_seconds", 14400) + ), + event_hook=_get_param(litellm_params, guardrail, "mode"), + default_on=bool(_get_param(litellm_params, guardrail, "default_on", False)), + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.SENSITIVE_DATA_ROUTING.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.SENSITIVE_DATA_ROUTING.value: SensitiveDataRoutingGuardrail, +} + +__all__: List[str] = [ + "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..e801fdfe011 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/sensitive_data_routing/sensitive_data_routing.py @@ -0,0 +1,200 @@ +""" +Sensitive Data Routing guardrail. + +Detects sensitive data in a request and, instead of blocking or redacting it, +reroutes the request to an on-premise model. When sticky sessions are enabled, +every later turn in the same session is also routed on-premise so a conversation +that once touched sensitive data never leaves the on-premise model. +""" + +import re +from typing import ( + TYPE_CHECKING, + Any, + Iterator, + List, + Optional, + Pattern, + Type, + Union, +) + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes + +if TYPE_CHECKING: + from litellm.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + +CACHE_KEY_PREFIX = "sensitive_data_routing" + + +class SensitiveDataRoutingGuardrail(CustomGuardrail): + def __init__( + self, + on_premise_model: str, + guardrail_name: Optional[str] = None, + prebuilt_patterns: Optional[List[str]] = None, + regex_patterns: Optional[List[str]] = None, + keywords: Optional[List[str]] = None, + sticky_session: bool = True, + session_ttl_seconds: int = 14400, + event_hook: Optional[Union[str, GuardrailEventHooks]] = None, + default_on: bool = False, + **kwargs: Any, + ) -> None: + super().__init__( + guardrail_name=guardrail_name or "sensitive_data_routing", + supported_event_hooks=[GuardrailEventHooks.pre_call], + event_hook=event_hook or GuardrailEventHooks.pre_call, + default_on=default_on, + **kwargs, + ) + self.on_premise_model = on_premise_model + self.sticky_session = sticky_session + self.session_ttl_seconds = session_ttl_seconds + self._patterns = self._compile_patterns(prebuilt_patterns, regex_patterns) + self._keywords = [k.lower() for k in (keywords or [])] + if not self._patterns and not self._keywords: + raise ValueError( + "sensitive_data_routing requires at least one of prebuilt_patterns, " + "regex_patterns, or keywords" + ) + + @staticmethod + def _compile_patterns( + prebuilt_patterns: Optional[List[str]], + regex_patterns: Optional[List[str]], + ) -> List[Pattern]: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import ( + get_compiled_pattern, + ) + + compiled: List[Pattern] = [ + get_compiled_pattern(name) for name in (prebuilt_patterns or []) + ] + compiled.extend( + re.compile(pattern, re.IGNORECASE) for pattern in (regex_patterns or []) + ) + return compiled + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingConfigModel, + ) + + return SensitiveDataRoutingConfigModel + + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: "DualCache", + data: dict, + call_type: str, + ) -> Optional[dict]: + session_id = self._get_session_id(data) + + session_pinned = ( + self.sticky_session + and session_id is not None + and await self._is_session_pinned(cache, session_id) + ) + detected = (not session_pinned) and self._contains_sensitive_data( + data, call_type + ) + + if not session_pinned and not detected: + return None + + if detected and self.sticky_session and session_id is not None: + await self._pin_session(cache, session_id) + + original_model = data.get("model") + data["model"] = self.on_premise_model + self._log_route( + data=data, + original_model=original_model, + detected=detected, + session_id=session_id, + ) + return data + + def _contains_sensitive_data(self, data: dict, call_type: str) -> bool: + messages = self.get_guardrails_messages_for_call_type( + call_type=CallTypes(call_type), data=data + ) + for text in self._iter_message_texts(messages): + if any(pattern.search(text) for pattern in self._patterns): + return True + lowered = text.lower() + if any(keyword in lowered for keyword in self._keywords): + return True + return False + + @staticmethod + def _iter_message_texts(messages: Optional[List[Any]]) -> Iterator[str]: + for message in messages or []: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, str): + yield content + elif isinstance(content, list): + for part in content: + if isinstance(part, dict) and isinstance(part.get("text"), str): + yield part["text"] + + @staticmethod + def _get_session_id(data: dict) -> Optional[str]: + session_id = data.get("litellm_session_id") + if session_id: + return str(session_id) + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict) and meta.get("session_id"): + return str(meta["session_id"]) + return None + + def _session_cache_key(self, session_id: str) -> str: + return f"{CACHE_KEY_PREFIX}:{self.guardrail_name}:{session_id}" + + async def _is_session_pinned(self, cache: "DualCache", session_id: str) -> bool: + return bool( + await cache.async_get_cache(key=self._session_cache_key(session_id)) + ) + + async def _pin_session(self, cache: "DualCache", session_id: str) -> None: + await cache.async_set_cache( + key=self._session_cache_key(session_id), + value=True, + ttl=self.session_ttl_seconds, + ) + + def _log_route( + self, + data: dict, + original_model: Optional[str], + detected: bool, + session_id: Optional[str], + ) -> None: + verbose_proxy_logger.info( + "sensitive_data_routing: rerouting model=%s -> %s (detected=%s, session=%s)", + original_model, + self.on_premise_model, + detected, + session_id, + ) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "action": "route", + "on_premise_model": self.on_premise_model, + "trigger": "detection" if detected else "sticky_session", + }, + request_data=data, + guardrail_status="guardrail_intervened", + event_type=GuardrailEventHooks.pre_call, + ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 25c0bcabb4a..3b8232823ab 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -26,6 +26,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor from litellm.types.proxy.guardrails.guardrail_hooks.promptguard import ( PromptGuardConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( XecGuardConfigModel, ) @@ -107,6 +110,7 @@ class SupportedGuardrailIntegrations(Enum): QOSTODIAN_NEXUS = "qostodian_nexus" RUBRIK = "rubrik" VIGIL_GUARD = "vigil_guard" + SENSITIVE_DATA_ROUTING = "sensitive_data_routing" class Role(Enum): @@ -805,6 +809,7 @@ class LitellmParams( HiddenlayerGuardrailConfigModel, QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, + SensitiveDataRoutingConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: Union[str, List[str], Mode] = Field( 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..cf9cc636b3a --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/sensitive_data_routing.py @@ -0,0 +1,40 @@ +"""Types for the Sensitive Data Routing guardrail.""" + +from typing import List, Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class SensitiveDataRoutingConfigModel(GuardrailConfigModel): + """Configuration for the Sensitive Data Routing guardrail.""" + + on_premise_model: Optional[str] = Field( + default=None, + description="Name of the model group to route to when sensitive data is detected. Must be present in your model_list.", + ) + prebuilt_patterns: Optional[List[str]] = Field( + default=None, + description="Names of built-in detection patterns to match against, e.g. 'us_ssn', 'credit_card', 'email'. See the litellm content filter prebuilt patterns for the full list.", + ) + regex_patterns: Optional[List[str]] = Field( + default=None, + description="Custom regular expressions; a match in any request message reroutes the request on-premise.", + ) + keywords: Optional[List[str]] = Field( + default=None, + description="Case-insensitive keywords; presence in any request message reroutes the request on-premise.", + ) + sticky_session: bool = Field( + default=True, + description="When True, once sensitive data is detected in a session every later turn in that session is also routed on-premise, even turns that contain no sensitive data. Requires the client to send a stable session id.", + ) + session_ttl_seconds: int = Field( + default=14400, + description="How long, in seconds, a session stays pinned on-premise after sensitive data is detected.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Sensitive Data Routing" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py new file mode 100644 index 00000000000..cb31716bac1 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_sensitive_data_routing.py @@ -0,0 +1,207 @@ +"""Tests for the Sensitive Data Routing guardrail.""" + +from typing import Any, Dict, List, Optional, Tuple + +import pytest + +from litellm.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.sensitive_data_routing import ( + SensitiveDataRoutingGuardrail, + initialize_guardrail, +) +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams + +ON_PREM = "on-prem-model" +USER_KEY = UserAPIKeyAuth() + + +class RecordingCache(DualCache): + """DualCache that records writes so tests can assert pinning behavior.""" + + def __init__(self) -> None: + super().__init__() + self.sets: List[Tuple[str, Any, Optional[int]]] = [] + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + self.sets.append((key, value, kwargs.get("ttl"))) + await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + +def _request(text, model="gpt-4o", **extra) -> Dict[str, Any]: + return {"model": model, "messages": [{"role": "user", "content": text}], **extra} + + +def _make_guardrail(**overrides) -> SensitiveDataRoutingGuardrail: + params: Dict[str, Any] = dict( + guardrail_name="sdr", + on_premise_model=ON_PREM, + prebuilt_patterns=["us_ssn"], + keywords=["confidential"], + ) + params.update(overrides) + return SensitiveDataRoutingGuardrail(**params) + + +async def _hook(guardrail, data, cache=None): + return await guardrail.async_pre_call_hook( + user_api_key_dict=USER_KEY, + cache=cache or DualCache(), + data=data, + call_type="acompletion", + ) + + +@pytest.mark.asyncio +async def test_clean_prompt_is_not_rerouted(): + g = _make_guardrail() + data = _request("what's the weather today?") + result = await _hook(g, data) + assert result is None + assert data["model"] == "gpt-4o" + + +@pytest.mark.asyncio +async def test_prebuilt_pattern_match_reroutes_to_on_prem(): + g = _make_guardrail() + data = _request("my social security number is 123-45-6789") + result = await _hook(g, data) + assert result is not None + assert result["model"] == ON_PREM + assert data["model"] == ON_PREM + + +@pytest.mark.asyncio +async def test_custom_regex_match_reroutes(): + g = _make_guardrail(prebuilt_patterns=None, regex_patterns=[r"project\s+titan"]) + data = _request("notes on Project Titan rollout") + result = await _hook(g, data) + assert result is not None and result["model"] == ON_PREM + + +@pytest.mark.asyncio +async def test_keyword_match_is_case_insensitive(): + g = _make_guardrail(prebuilt_patterns=None, keywords=["confidential"]) + data = _request("this memo is CONFIDENTIAL") + result = await _hook(g, data) + assert result is not None and result["model"] == ON_PREM + + +@pytest.mark.asyncio +async def test_detects_text_in_content_parts_list(): + g = _make_guardrail() + data = { + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "ssn 123-45-6789"}], + } + ], + } + result = await _hook(g, data) + assert result is not None and result["model"] == ON_PREM + + +@pytest.mark.asyncio +async def test_reroute_records_guardrail_logging_information(): + g = _make_guardrail() + data = _request("ssn 123-45-6789") + await _hook(g, data) + entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "sdr" + + +@pytest.mark.asyncio +async def test_sticky_session_pins_following_clean_turns(): + g = _make_guardrail(sticky_session=True, session_ttl_seconds=999) + cache = RecordingCache() + session = {"litellm_session_id": "sess-1"} + + first = await _hook(g, _request("ssn 123-45-6789", **session), cache) + assert first is not None and first["model"] == ON_PREM + assert cache.sets and cache.sets[0][0].endswith("sess-1") + assert cache.sets[0][2] == 999 # ttl is honored + + follow_up = _request("just a normal follow-up question", **session) + result = await _hook(g, follow_up, cache) + assert result is not None + assert follow_up["model"] == ON_PREM + + +@pytest.mark.asyncio +async def test_non_sticky_does_not_pin_session(): + g = _make_guardrail(sticky_session=False) + cache = RecordingCache() + session = {"litellm_session_id": "sess-2"} + + await _hook(g, _request("ssn 123-45-6789", **session), cache) + assert cache.sets == [] # nothing pinned + + follow_up = _request("a normal follow-up", **session) + result = await _hook(g, follow_up, cache) + assert result is None + assert follow_up["model"] == "gpt-4o" + + +@pytest.mark.asyncio +async def test_sessions_are_isolated_from_each_other(): + g = _make_guardrail(sticky_session=True) + cache = RecordingCache() + + await _hook(g, _request("ssn 123-45-6789", litellm_session_id="flagged"), cache) + + other = _request("nothing sensitive here", litellm_session_id="other") + result = await _hook(g, other, cache) + assert result is None + assert other["model"] == "gpt-4o" + + +@pytest.mark.asyncio +async def test_session_id_read_from_metadata(): + g = _make_guardrail(sticky_session=True) + cache = RecordingCache() + meta = {"metadata": {"session_id": "meta-sess"}} + + await _hook(g, _request("ssn 123-45-6789", **meta), cache) + follow_up = _request("benign follow-up", **meta) + result = await _hook(g, follow_up, cache) + assert result is not None and follow_up["model"] == ON_PREM + + +def test_requires_at_least_one_detector(): + with pytest.raises(ValueError): + SensitiveDataRoutingGuardrail(guardrail_name="sdr", on_premise_model=ON_PREM) + + +def test_only_supports_pre_call_event_hook(): + with pytest.raises(ValueError): + _make_guardrail(event_hook=GuardrailEventHooks.post_call) + + +def test_initializer_requires_on_premise_model(): + params = LitellmParams(guardrail="sensitive_data_routing", mode="pre_call") + with pytest.raises(ValueError): + initialize_guardrail( + litellm_params=params, + guardrail={"guardrail_name": "sdr", "litellm_params": {}}, + ) + + +def test_initializer_builds_guardrail_from_config(): + params = LitellmParams( + guardrail="sensitive_data_routing", + mode="pre_call", + on_premise_model=ON_PREM, + prebuilt_patterns=["us_ssn"], + keywords=["confidential"], + session_ttl_seconds=120, + ) + instance = initialize_guardrail( + litellm_params=params, + guardrail={"guardrail_name": "sdr", "litellm_params": {}}, + ) + assert isinstance(instance, SensitiveDataRoutingGuardrail) + assert instance.on_premise_model == ON_PREM + assert instance.session_ttl_seconds == 120