mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 23:05:52 +00:00
Mateo Wang
812a2217ca
* feat(guardrails): add sensitive data routing to on-premise models When a guardrail detects sensitive data, route to an on-premise model instead of blocking or redacting. All subsequent requests in that session continue routing to the same model (sticky routing). New config options for guardrails: - on_sensitive_data: 'block' (default) or 'route' - sensitive_data_route_to_model: target model for rerouting - sticky_session_routing: persist routing for session (default: true) New exception SensitiveDataRouteException triggers rerouting when raised by guardrails. The proxy catches it, stores the routing decision in cache, and modifies the request's model field. New hook _PROXY_SensitiveDataRoutingHandler checks incoming requests against cached routing decisions and applies sticky routing. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix: black formatting for custom_guardrail.py https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * test: improve test coverage for sensitive data routing feature Add additional tests for: - Cache key format and TTL constants - Session ID extraction from multiple locations - Custom guardrail initialization with routing config - Exception string representation and custom messages - Redis cache paths including fallback behavior - Edge cases in pre-call hook https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix: use correct GuardrailRaisedException parameters Replace invalid 'source' parameter with 'guardrail_name' to match the exception's actual signature. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * test: move sensitive data routing tests to hooks directory Move test file to align with source code structure. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix(guardrails): honor sticky_session_routing flag and scope session routing per API key Propagate sticky_session_routing through SensitiveDataRouteException so a guardrail configured with sticky_session_routing=False reroutes only the triggering request without persisting a session override. Scope the routing cache key to the requesting API key so sessions from different tenants cannot collide, and warn when sticky routing is requested but the hook is not registered. * refactor(guardrails): dedupe session-id extraction and drop redundant import Extract the shared session-id lookup into get_session_id_from_request_data so the sensitive-data routing hook and CustomGuardrail no longer keep two identical copies of the logic. Remove the redundant local import of GuardrailRaisedException in handle_sensitive_data_detection, and document that detection_info is surfaced in request metadata and logs so it must not carry raw sensitive values. * fix(guardrails): guard None user_api_key_dict in sensitive data route handler * fix(responses): send application/json Content-Type on responses DELETE OpenAI's responses DELETE endpoint now rejects requests that arrive without a Content-Type header, defaulting them to application/octet-stream and returning 'Unsupported content type: application/octet-stream'. The delete handler sent no body and therefore no Content-Type, so the request failed. Declare application/json on the delete request, matching the OpenAI SDK. * fix(guardrails): backfill in-memory cache after redis hit in sensitive data routing When _get_routed_model resolves a routing override from Redis it now also populates the local in-memory cache. Without the write-back, a non-writing instance that only ever reads from Redis would lose the sticky routing decision the moment Redis became unavailable, silently reverting sensitive sessions to the default model. * fix(guardrails): scope sticky sensitive-data routing to JWT principal Keyless auth (JWT and similar) has no api_key, so every such caller shared the "default" cache namespace. One authenticated user could reuse another user's session_id, trip the guardrail, and silently force the other user's subsequent requests onto the cached on-prem model for the TTL. Resolve the routing tenant from the api_key when present, otherwise from a stable principal built from the user/team/org identity, before reading or writing the session route. * fix(guardrails): require route target model when on_sensitive_data='route' * fix(guardrails): mark user_api_key_dict Optional in sensitive-data route handler * fix(guardrails): use remaining redis ttl for local backfill and str env default * fix(guardrails): graceful block when routing configured but no session_id handle_sensitive_data_detection promised to raise only SensitiveDataRouteException or GuardrailRaisedException, but when routing was configured and the request had no session_id it let a ValueError from raise_sensitive_data_route_exception propagate, surfacing as an HTTP 500 instead of a block. Fall back to a graceful block in that case so the documented contract holds. * fix(guardrails): run remaining guardrails after sensitive-data reroute Defer the SensitiveDataRouteException until every guardrail in the pre-call loop has run, so downstream security guardrails are no longer skipped when an earlier guardrail triggers routing. The first reroute wins and a later guardrail that blocks still propagates. Also normalize on_sensitive_data to lowercase like sibling on_* config fields so case-insensitive values are accepted. * fix(guardrails): classify sensitive-data reroute as guardrail intervention * fix(guardrails): record sensitive-data reroute as prometheus intervention not error * fix(guardrails): record service span for routing guardrail and move case-normalizer to base params Drop the early continue so a guardrail that signals sensitive-data routing still emits its PROXY_PRE_CALL service span like every other callback. Move the lowercase normalizer onto BaseLitellmParams so on_sensitive_data is normalized consistently when BaseLitellmParams is constructed directly, matching the cross-field route->model validator that already lives on the base.
1037 lines
37 KiB
Python
1037 lines
37 KiB
Python
"""
|
|
Tests for Sensitive Data Routing feature.
|
|
|
|
This feature allows guardrails to route requests to a different model
|
|
(typically on-premise) when sensitive data is detected, instead of blocking.
|
|
All subsequent requests in the same session are routed to the same model.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Any, Dict, Optional
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.exceptions import SensitiveDataRouteException
|
|
from litellm.integrations.custom_guardrail import (
|
|
CustomGuardrail,
|
|
get_session_id_from_request_data,
|
|
)
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.hooks.sensitive_data_routing import (
|
|
_PROXY_SensitiveDataRoutingHandler,
|
|
SENSITIVE_ROUTING_CACHE_PREFIX,
|
|
DEFAULT_SENSITIVE_ROUTING_TTL,
|
|
)
|
|
|
|
|
|
class MockInternalUsageCache:
|
|
def __init__(self):
|
|
self._cache: Dict[str, Any] = {}
|
|
self._ttls: Dict[str, int] = {}
|
|
self.dual_cache = MagicMock()
|
|
self.dual_cache.redis_cache = None
|
|
|
|
async def async_get_cache(self, key: str, **kwargs) -> Optional[Any]:
|
|
return self._cache.get(key)
|
|
|
|
async def async_set_cache(self, key: str, value: Any, ttl: int = 3600, **kwargs):
|
|
self._cache[key] = value
|
|
self._ttls[key] = ttl
|
|
|
|
|
|
class TestSensitiveDataRoutingHandler:
|
|
@pytest.fixture
|
|
def handler(self):
|
|
cache = MockInternalUsageCache()
|
|
return _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
|
|
@pytest.fixture
|
|
def user_api_key_dict(self):
|
|
return UserAPIKeyAuth(api_key="test-key")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_session_routing(self, handler):
|
|
key = UserAPIKeyAuth(api_key="hashed-key")
|
|
await handler.set_session_routing(
|
|
session_id="test-session-123",
|
|
model="on-premise-model",
|
|
user_api_key_dict=key,
|
|
guardrail_name="test-guardrail",
|
|
)
|
|
|
|
routed_model = await handler._get_routed_model("test-session-123", key)
|
|
assert routed_model == "on-premise-model"
|
|
|
|
def test_get_session_id_from_metadata(self):
|
|
data = {"metadata": {"session_id": "session-from-metadata"}}
|
|
session_id = get_session_id_from_request_data(data)
|
|
assert session_id == "session-from-metadata"
|
|
|
|
def test_get_session_id_from_litellm_metadata(self):
|
|
data = {"litellm_metadata": {"session_id": "session-from-litellm-metadata"}}
|
|
session_id = get_session_id_from_request_data(data)
|
|
assert session_id == "session-from-litellm-metadata"
|
|
|
|
def test_get_session_id_from_litellm_session_id(self):
|
|
data = {"litellm_session_id": "session-direct"}
|
|
session_id = get_session_id_from_request_data(data)
|
|
assert session_id == "session-direct"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_no_session(self, handler, user_api_key_dict):
|
|
data = {"model": "gpt-4"}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
assert result is None
|
|
assert data["model"] == "gpt-4"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_with_routing_override(
|
|
self, handler, user_api_key_dict
|
|
):
|
|
await handler.set_session_routing(
|
|
session_id="routed-session",
|
|
model="on-premise-model",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "routed-session"},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert result is not None
|
|
assert result["model"] == "on-premise-model"
|
|
assert result["metadata"]["sensitive_data_routing_applied"] is True
|
|
assert result["metadata"]["sensitive_data_routing_original_model"] == "gpt-4"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_no_override_needed(self, handler, user_api_key_dict):
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "no-override-session"},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
assert result is None
|
|
assert data["model"] == "gpt-4"
|
|
|
|
|
|
class TestSensitiveDataRouteException:
|
|
def test_exception_creation(self):
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-premise-model",
|
|
session_id="test-session",
|
|
guardrail_name="test-guardrail",
|
|
detection_info={"detected_entities": ["SSN", "CREDIT_CARD"]},
|
|
)
|
|
|
|
assert exc.route_to_model == "on-premise-model"
|
|
assert exc.session_id == "test-session"
|
|
assert exc.guardrail_name == "test-guardrail"
|
|
assert "SSN" in exc.detection_info["detected_entities"]
|
|
|
|
|
|
class TestCustomGuardrailSensitiveDataRouting:
|
|
def test_should_route_on_sensitive_data_false_by_default(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
assert guardrail.should_route_on_sensitive_data() is False
|
|
|
|
def test_should_route_on_sensitive_data_true(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
assert guardrail.should_route_on_sensitive_data() is True
|
|
|
|
def test_should_route_on_sensitive_data_missing_model(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
)
|
|
assert guardrail.should_route_on_sensitive_data() is False
|
|
|
|
def test_raise_sensitive_data_route_exception(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
|
|
request_data = {"model": "gpt-4", "metadata": {"session_id": "test-session"}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException) as exc_info:
|
|
guardrail.raise_sensitive_data_route_exception(
|
|
route_to_model="on-premise-model",
|
|
request_data=request_data,
|
|
detection_info={"type": "PII"},
|
|
)
|
|
|
|
assert exc_info.value.route_to_model == "on-premise-model"
|
|
assert exc_info.value.session_id == "test-session"
|
|
|
|
def test_raise_exception_carries_sticky_flag_false(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
sticky_session_routing=False,
|
|
)
|
|
|
|
request_data = {"metadata": {"session_id": "test-session"}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException) as exc_info:
|
|
guardrail.raise_sensitive_data_route_exception(
|
|
route_to_model="on-premise-model",
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert exc_info.value.sticky_session_routing is False
|
|
|
|
def test_raise_exception_carries_sticky_flag_default_true(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
|
|
request_data = {"metadata": {"session_id": "test-session"}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException) as exc_info:
|
|
guardrail.raise_sensitive_data_route_exception(
|
|
route_to_model="on-premise-model",
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert exc_info.value.sticky_session_routing is True
|
|
|
|
def test_raise_sensitive_data_route_exception_missing_session(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
|
|
request_data = {"model": "gpt-4"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
guardrail.raise_sensitive_data_route_exception(
|
|
route_to_model="on-premise-model",
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert "session_id" in str(exc_info.value)
|
|
|
|
def test_handle_sensitive_data_detection_route(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
|
|
request_data = {"model": "gpt-4", "metadata": {"session_id": "test-session"}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException) as exc_info:
|
|
guardrail.handle_sensitive_data_detection(
|
|
request_data=request_data,
|
|
detection_info={"type": "PII"},
|
|
)
|
|
|
|
assert exc_info.value.route_to_model == "on-premise-model"
|
|
|
|
def test_handle_sensitive_data_detection_block(self):
|
|
from litellm.exceptions import GuardrailRaisedException
|
|
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
|
|
request_data = {"model": "gpt-4", "metadata": {"session_id": "test-session"}}
|
|
|
|
with pytest.raises(GuardrailRaisedException):
|
|
guardrail.handle_sensitive_data_detection(
|
|
request_data=request_data,
|
|
)
|
|
|
|
def test_handle_sensitive_data_detection_route_no_session_falls_back_to_block(self):
|
|
from litellm.exceptions import GuardrailRaisedException
|
|
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
|
|
request_data = {"model": "gpt-4"}
|
|
|
|
with pytest.raises(GuardrailRaisedException) as exc_info:
|
|
guardrail.handle_sensitive_data_detection(
|
|
request_data=request_data,
|
|
detection_info={"type": "PII"},
|
|
)
|
|
|
|
assert "session_id" in str(exc_info.value)
|
|
|
|
|
|
class TestStickySessionRouting:
|
|
@pytest.fixture
|
|
def handler(self):
|
|
cache = MockInternalUsageCache()
|
|
return _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
|
|
@pytest.fixture
|
|
def user_api_key_dict(self):
|
|
return UserAPIKeyAuth(api_key="test-key")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_routing_persists(self, handler, user_api_key_dict):
|
|
session_id = "sticky-session"
|
|
await handler.set_session_routing(
|
|
session_id=session_id,
|
|
model="on-premise-model",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
for i in range(5):
|
|
data = {
|
|
"model": f"gpt-{i}",
|
|
"metadata": {"session_id": session_id},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert result is not None
|
|
assert result["model"] == "on-premise-model"
|
|
assert (
|
|
result["metadata"]["sensitive_data_routing_original_model"]
|
|
== f"gpt-{i}"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_sessions_independent(self, handler, user_api_key_dict):
|
|
await handler.set_session_routing(
|
|
session_id="session-a",
|
|
model="on-premise-model-a",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
await handler.set_session_routing(
|
|
session_id="session-b",
|
|
model="on-premise-model-b",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
data_a = {"model": "gpt-4", "metadata": {"session_id": "session-a"}}
|
|
data_b = {"model": "gpt-4", "metadata": {"session_id": "session-b"}}
|
|
data_c = {"model": "gpt-4", "metadata": {"session_id": "session-c"}}
|
|
|
|
result_a = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data_a,
|
|
call_type="completion",
|
|
)
|
|
result_b = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data_b,
|
|
call_type="completion",
|
|
)
|
|
result_c = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data_c,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert result_a["model"] == "on-premise-model-a"
|
|
assert result_b["model"] == "on-premise-model-b"
|
|
assert result_c is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_is_isolated_per_api_key(self, handler):
|
|
shared_session = "shared-session-id"
|
|
await handler.set_session_routing(
|
|
session_id=shared_session,
|
|
model="on-premise-model",
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
)
|
|
|
|
data_for_tenant_b = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": shared_session},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-b"),
|
|
cache=DualCache(),
|
|
data=data_for_tenant_b,
|
|
call_type="completion",
|
|
)
|
|
assert result is None
|
|
assert data_for_tenant_b["model"] == "gpt-4"
|
|
|
|
data_for_tenant_a = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": shared_session},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
cache=DualCache(),
|
|
data=data_for_tenant_a,
|
|
call_type="completion",
|
|
)
|
|
assert result is not None
|
|
assert result["model"] == "on-premise-model"
|
|
|
|
|
|
class TestCacheKeyAndTTL:
|
|
def test_cache_prefix_constant(self):
|
|
assert SENSITIVE_ROUTING_CACHE_PREFIX == "sensitive_route"
|
|
|
|
def test_default_ttl_constant(self):
|
|
assert DEFAULT_SENSITIVE_ROUTING_TTL == 3600
|
|
|
|
def test_make_cache_key_format(self):
|
|
cache = MockInternalUsageCache()
|
|
handler = _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
key = handler._make_cache_key("test-session-123", "hashed-key")
|
|
assert key == "{sensitive_route:hashed-key:test-session-123}:model"
|
|
|
|
def test_make_cache_key_is_tenant_scoped(self):
|
|
cache = MockInternalUsageCache()
|
|
handler = _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
key_a = handler._make_cache_key("shared-session", "key-a")
|
|
key_b = handler._make_cache_key("shared-session", "key-b")
|
|
assert key_a != key_b
|
|
|
|
def test_resolve_tenant_prefers_api_key(self):
|
|
tenant = _PROXY_SensitiveDataRoutingHandler._resolve_tenant(
|
|
UserAPIKeyAuth(api_key="hashed-key", user_id="alice")
|
|
)
|
|
assert tenant == "hashed-key"
|
|
|
|
def test_resolve_tenant_falls_back_to_jwt_principal(self):
|
|
tenant = _PROXY_SensitiveDataRoutingHandler._resolve_tenant(
|
|
UserAPIKeyAuth(api_key=None, user_id="alice", team_id="t1", org_id="o1")
|
|
)
|
|
assert tenant == "user:alice|team:t1|org:o1"
|
|
|
|
def test_resolve_tenant_distinguishes_keyless_principals(self):
|
|
tenant_a = _PROXY_SensitiveDataRoutingHandler._resolve_tenant(
|
|
UserAPIKeyAuth(api_key=None, user_id="alice")
|
|
)
|
|
tenant_b = _PROXY_SensitiveDataRoutingHandler._resolve_tenant(
|
|
UserAPIKeyAuth(api_key=None, user_id="bob")
|
|
)
|
|
assert tenant_a != tenant_b
|
|
|
|
def test_resolve_tenant_defaults_when_anonymous(self):
|
|
assert _PROXY_SensitiveDataRoutingHandler._resolve_tenant(None) == "default"
|
|
assert (
|
|
_PROXY_SensitiveDataRoutingHandler._resolve_tenant(
|
|
UserAPIKeyAuth(api_key=None)
|
|
)
|
|
== "default"
|
|
)
|
|
|
|
|
|
class TestCustomGuardrailSessionIdExtraction:
|
|
def test_get_session_id_from_litellm_session_id(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {"litellm_session_id": "session-direct-123"}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id == "session-direct-123"
|
|
|
|
def test_get_session_id_from_metadata(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {"metadata": {"session_id": "session-metadata-456"}}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id == "session-metadata-456"
|
|
|
|
def test_get_session_id_from_litellm_metadata(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {"litellm_metadata": {"session_id": "session-litellm-meta-789"}}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id == "session-litellm-meta-789"
|
|
|
|
def test_get_session_id_returns_none_when_missing(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {"model": "gpt-4"}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id is None
|
|
|
|
def test_get_session_id_priority_litellm_session_id_first(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {
|
|
"litellm_session_id": "priority-session",
|
|
"metadata": {"session_id": "should-not-use"},
|
|
"litellm_metadata": {"session_id": "also-not-this"},
|
|
}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id == "priority-session"
|
|
|
|
def test_get_session_id_converts_to_string(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
request_data = {"litellm_session_id": 12345}
|
|
session_id = guardrail._get_session_id_from_request_data(request_data)
|
|
assert session_id == "12345"
|
|
assert isinstance(session_id, str)
|
|
|
|
|
|
class TestCustomGuardrailInit:
|
|
def test_init_with_routing_config(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test-guardrail",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
sticky_session_routing=True,
|
|
)
|
|
assert guardrail.on_sensitive_data == "route"
|
|
assert guardrail.sensitive_data_route_to_model == "on-premise-model"
|
|
assert guardrail.sticky_session_routing is True
|
|
|
|
def test_init_default_values(self):
|
|
guardrail = CustomGuardrail(guardrail_name="test")
|
|
assert guardrail.on_sensitive_data is None
|
|
assert guardrail.sensitive_data_route_to_model is None
|
|
assert guardrail.sticky_session_routing is True
|
|
|
|
|
|
class TestSensitiveDataRouteExceptionStr:
|
|
def test_exception_str_representation(self):
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-premise-model",
|
|
session_id="test-session",
|
|
guardrail_name="pii-detector",
|
|
)
|
|
assert (
|
|
str(exc)
|
|
== "Sensitive data detected by pii-detector. Routing to model: on-premise-model"
|
|
)
|
|
|
|
def test_exception_custom_message(self):
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-premise-model",
|
|
session_id="test-session",
|
|
guardrail_name="pii-detector",
|
|
message="Custom error message",
|
|
)
|
|
assert str(exc) == "Custom error message"
|
|
assert exc.message == "Custom error message"
|
|
|
|
|
|
class TestRedisCache:
|
|
@pytest.fixture
|
|
def handler_with_redis(self):
|
|
cache = MockInternalUsageCache()
|
|
mock_redis = AsyncMock()
|
|
cache.dual_cache.redis_cache = mock_redis
|
|
return _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_routed_model_from_redis(self, handler_with_redis):
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
return_value="redis-model"
|
|
)
|
|
result = await handler_with_redis._get_routed_model(
|
|
"session-123", UserAPIKeyAuth(api_key="hashed-key")
|
|
)
|
|
assert result == "redis-model"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_routed_model_backfills_in_memory_after_redis_hit(
|
|
self, handler_with_redis
|
|
):
|
|
cache_key = "{sensitive_route:hashed-key:session-123}:model"
|
|
key = UserAPIKeyAuth(api_key="hashed-key")
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
return_value="on-premise-model"
|
|
)
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_ttl = (
|
|
AsyncMock(return_value=120)
|
|
)
|
|
|
|
first = await handler_with_redis._get_routed_model("session-123", key)
|
|
assert first == "on-premise-model"
|
|
assert handler_with_redis.internal_usage_cache._cache[cache_key] == (
|
|
"on-premise-model"
|
|
)
|
|
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
side_effect=Exception("Redis went down")
|
|
)
|
|
second = await handler_with_redis._get_routed_model("session-123", key)
|
|
assert second == "on-premise-model"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_uses_remaining_redis_ttl(self, handler_with_redis):
|
|
cache_key = "{sensitive_route:hashed-key:session-123}:model"
|
|
key = UserAPIKeyAuth(api_key="hashed-key")
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
return_value="on-premise-model"
|
|
)
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_ttl = (
|
|
AsyncMock(return_value=42)
|
|
)
|
|
|
|
await handler_with_redis._get_routed_model("session-123", key)
|
|
|
|
assert handler_with_redis.internal_usage_cache._ttls[cache_key] == 42
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_falls_back_to_full_ttl_when_redis_ttl_missing(
|
|
self, handler_with_redis
|
|
):
|
|
cache_key = "{sensitive_route:hashed-key:session-123}:model"
|
|
key = UserAPIKeyAuth(api_key="hashed-key")
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
return_value="on-premise-model"
|
|
)
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_ttl = (
|
|
AsyncMock(return_value=None)
|
|
)
|
|
|
|
await handler_with_redis._get_routed_model("session-123", key)
|
|
|
|
assert (
|
|
handler_with_redis.internal_usage_cache._ttls[cache_key]
|
|
== handler_with_redis.ttl
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_routed_model_redis_fallback_on_error(self, handler_with_redis):
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_get_cache = AsyncMock(
|
|
side_effect=Exception("Redis connection error")
|
|
)
|
|
handler_with_redis.internal_usage_cache._cache[
|
|
"{sensitive_route:hashed-key:session-123}:model"
|
|
] = "fallback-model"
|
|
result = await handler_with_redis._get_routed_model(
|
|
"session-123", UserAPIKeyAuth(api_key="hashed-key")
|
|
)
|
|
assert result == "fallback-model"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_session_routing_with_redis(self, handler_with_redis):
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_set_cache = (
|
|
AsyncMock()
|
|
)
|
|
await handler_with_redis.set_session_routing(
|
|
session_id="session-456",
|
|
model="on-premise-model",
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
|
guardrail_name="test-guardrail",
|
|
)
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_set_cache.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_session_routing_redis_fallback_on_error(
|
|
self, handler_with_redis
|
|
):
|
|
handler_with_redis.internal_usage_cache.dual_cache.redis_cache.async_set_cache = AsyncMock(
|
|
side_effect=Exception("Redis connection error")
|
|
)
|
|
await handler_with_redis.set_session_routing(
|
|
session_id="session-789",
|
|
model="on-premise-model",
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
|
)
|
|
cache_key = "{sensitive_route:hashed-key:session-789}:model"
|
|
assert (
|
|
handler_with_redis.internal_usage_cache._cache[cache_key]
|
|
== "on-premise-model"
|
|
)
|
|
|
|
|
|
class TestPreCallHookEdgeCases:
|
|
@pytest.fixture
|
|
def handler(self):
|
|
cache = MockInternalUsageCache()
|
|
return _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
|
|
@pytest.fixture
|
|
def user_api_key_dict(self):
|
|
return UserAPIKeyAuth(api_key="test-key")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_call_hook_same_model_no_change(self, handler, user_api_key_dict):
|
|
await handler.set_session_routing(
|
|
session_id="same-model-session",
|
|
model="gpt-4",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "same-model-session"},
|
|
}
|
|
result = await handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=DualCache(),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
assert result is None
|
|
|
|
|
|
class TestHandleSensitiveDataDetectionWithRouting:
|
|
def test_handle_sensitive_data_detection_full_flow(self):
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="pii-guardrail",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-premise-model",
|
|
)
|
|
|
|
request_data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "flow-test-session"},
|
|
"messages": [{"role": "user", "content": "My SSN is 123-45-6789"}],
|
|
}
|
|
|
|
with pytest.raises(SensitiveDataRouteException) as exc_info:
|
|
guardrail.handle_sensitive_data_detection(
|
|
request_data=request_data,
|
|
detection_info={"detected_entities": ["SSN"]},
|
|
)
|
|
|
|
exc = exc_info.value
|
|
assert exc.route_to_model == "on-premise-model"
|
|
assert exc.session_id == "flow-test-session"
|
|
assert exc.guardrail_name == "pii-guardrail"
|
|
assert exc.detection_info == {"detected_entities": ["SSN"]}
|
|
|
|
|
|
class TestProxyHandleSensitiveDataRouteException:
|
|
@pytest.fixture
|
|
def proxy_logging(self):
|
|
from litellm.proxy.utils import ProxyLogging
|
|
|
|
return ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
@pytest.fixture
|
|
def routing_hook(self):
|
|
cache = MockInternalUsageCache()
|
|
return _PROXY_SensitiveDataRoutingHandler(internal_usage_cache=cache)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_routing_persists_override(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-sticky",
|
|
guardrail_name="pii",
|
|
sticky_session_routing=True,
|
|
)
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-sticky"}}
|
|
|
|
result = await proxy_logging._handle_sensitive_data_route_exception(
|
|
exc, data, UserAPIKeyAuth(api_key="tenant-a")
|
|
)
|
|
|
|
assert result["model"] == "on-premise-model"
|
|
assert (
|
|
await routing_hook._get_routed_model(
|
|
"sess-sticky", UserAPIKeyAuth(api_key="tenant-a")
|
|
)
|
|
== "on-premise-model"
|
|
)
|
|
|
|
@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",
|
|
session_id="sess-non-sticky",
|
|
guardrail_name="pii",
|
|
sticky_session_routing=False,
|
|
)
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-non-sticky"}}
|
|
|
|
result = await proxy_logging._handle_sensitive_data_route_exception(
|
|
exc, data, UserAPIKeyAuth(api_key="tenant-a")
|
|
)
|
|
|
|
assert result["model"] == "on-premise-model"
|
|
assert (
|
|
await routing_hook._get_routed_model(
|
|
"sess-non-sticky", UserAPIKeyAuth(api_key="tenant-a")
|
|
)
|
|
is None
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_routing_handles_none_user_api_key_dict(
|
|
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-no-key",
|
|
guardrail_name="pii",
|
|
sticky_session_routing=True,
|
|
)
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-no-key"}}
|
|
|
|
result = await proxy_logging._handle_sensitive_data_route_exception(
|
|
exc, data, None
|
|
)
|
|
|
|
assert result["model"] == "on-premise-model"
|
|
assert (
|
|
await routing_hook._get_routed_model("sess-no-key", None)
|
|
== "on-premise-model"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_routing_scopes_jwt_users_by_principal(
|
|
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="shared-jwt-session",
|
|
guardrail_name="pii",
|
|
sticky_session_routing=True,
|
|
)
|
|
attacker = UserAPIKeyAuth(api_key=None, user_id="attacker", team_id="team-x")
|
|
await proxy_logging._handle_sensitive_data_route_exception(
|
|
exc,
|
|
{"model": "gpt-4", "metadata": {"session_id": "shared-jwt-session"}},
|
|
attacker,
|
|
)
|
|
|
|
victim = UserAPIKeyAuth(api_key=None, user_id="victim", team_id="team-y")
|
|
victim_data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "shared-jwt-session"},
|
|
}
|
|
result = await routing_hook.async_pre_call_hook(
|
|
user_api_key_dict=victim,
|
|
cache=DualCache(),
|
|
data=victim_data,
|
|
call_type="completion",
|
|
)
|
|
assert result is None
|
|
assert victim_data["model"] == "gpt-4"
|
|
|
|
attacker_data = {
|
|
"model": "gpt-4",
|
|
"metadata": {"session_id": "shared-jwt-session"},
|
|
}
|
|
result = await routing_hook.async_pre_call_hook(
|
|
user_api_key_dict=attacker,
|
|
cache=DualCache(),
|
|
data=attacker_data,
|
|
call_type="completion",
|
|
)
|
|
assert result is not None
|
|
assert result["model"] == "on-premise-model"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_routing_warns_when_hook_not_registered(self, proxy_logging):
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-premise-model",
|
|
session_id="sess-no-hook",
|
|
sticky_session_routing=True,
|
|
)
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-no-hook"}}
|
|
|
|
with patch("litellm.proxy.utils.verbose_proxy_logger.warning") as mock_warning:
|
|
result = await proxy_logging._handle_sensitive_data_route_exception(
|
|
exc, data, UserAPIKeyAuth(api_key="tenant-a")
|
|
)
|
|
|
|
assert result["model"] == "on-premise-model"
|
|
mock_warning.assert_called_once()
|
|
|
|
|
|
class _RoutingGuardrail(CustomGuardrail):
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.handle_sensitive_data_detection(request_data=data)
|
|
|
|
|
|
class _RecordingGuardrail(CustomGuardrail):
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.ran = False
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
self.ran = True
|
|
return None
|
|
|
|
|
|
class _BlockingGuardrail(CustomGuardrail):
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.ran = False
|
|
|
|
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
|
from litellm.exceptions import GuardrailRaisedException
|
|
|
|
self.ran = True
|
|
raise GuardrailRaisedException(
|
|
message="blocked", guardrail_name=self.guardrail_name
|
|
)
|
|
|
|
|
|
class TestPreCallHookDeferredRouting:
|
|
"""Guardrails after the one that triggers routing must still run."""
|
|
|
|
@pytest.fixture
|
|
def proxy_logging(self):
|
|
from litellm.proxy.utils import ProxyLogging
|
|
|
|
return ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def restore_callbacks(self):
|
|
import litellm
|
|
|
|
original = litellm.callbacks
|
|
litellm.callbacks = []
|
|
yield
|
|
litellm.callbacks = original
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_later_guardrail_runs_and_routing_applied(self, proxy_logging):
|
|
import litellm
|
|
|
|
router = _RoutingGuardrail(
|
|
guardrail_name="router",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-prem-model",
|
|
sticky_session_routing=False,
|
|
)
|
|
recorder = _RecordingGuardrail(
|
|
guardrail_name="recorder",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
)
|
|
litellm.callbacks = [router, recorder]
|
|
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-defer"}}
|
|
result = await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert recorder.ran is True
|
|
assert result["model"] == "on-prem-model"
|
|
assert result["metadata"]["sensitive_data_routing_applied"] is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_later_blocking_guardrail_overrides_routing(self, proxy_logging):
|
|
import litellm
|
|
from litellm.exceptions import GuardrailRaisedException
|
|
|
|
router = _RoutingGuardrail(
|
|
guardrail_name="router",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-prem-model",
|
|
sticky_session_routing=False,
|
|
)
|
|
blocker = _BlockingGuardrail(
|
|
guardrail_name="blocker",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
)
|
|
litellm.callbacks = [router, blocker]
|
|
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-block"}}
|
|
with pytest.raises(GuardrailRaisedException):
|
|
await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert blocker.ran is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_guardrail_records_service_span(self, proxy_logging):
|
|
import litellm
|
|
from litellm.types.services import ServiceTypes
|
|
|
|
class _SlowRoutingGuardrail(CustomGuardrail):
|
|
async def async_pre_call_hook(
|
|
self, user_api_key_dict, cache, data, call_type
|
|
):
|
|
await asyncio.sleep(0.02)
|
|
self.handle_sensitive_data_detection(request_data=data)
|
|
|
|
router = _SlowRoutingGuardrail(
|
|
guardrail_name="router",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-prem-model",
|
|
sticky_session_routing=False,
|
|
)
|
|
litellm.callbacks = [router]
|
|
|
|
recorded = AsyncMock()
|
|
proxy_logging.service_logging_obj.async_service_success_hook = recorded
|
|
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-span"}}
|
|
result = await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert result["model"] == "on-prem-model"
|
|
recorded.assert_called_once()
|
|
assert recorded.call_args.kwargs["call_type"] == "_SlowRoutingGuardrail"
|
|
assert recorded.call_args.kwargs["service"] == ServiceTypes.PROXY_PRE_CALL
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_recorded_as_intervention_not_prometheus_error(
|
|
self, proxy_logging
|
|
):
|
|
import litellm
|
|
from litellm.integrations.prometheus import PrometheusLogger
|
|
|
|
router = _RoutingGuardrail(
|
|
guardrail_name="router",
|
|
default_on=True,
|
|
event_hook="pre_call",
|
|
on_sensitive_data="route",
|
|
sensitive_data_route_to_model="on-prem-model",
|
|
sticky_session_routing=False,
|
|
)
|
|
prom = MagicMock(spec=PrometheusLogger)
|
|
litellm.callbacks = [router, prom]
|
|
|
|
data = {"model": "gpt-4", "metadata": {"session_id": "sess-prom"}}
|
|
result = await proxy_logging.pre_call_hook(
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="tenant-a"),
|
|
data=data,
|
|
call_type="completion",
|
|
)
|
|
|
|
assert result["model"] == "on-prem-model"
|
|
prom._record_guardrail_metrics.assert_called_once()
|
|
metrics_kwargs = prom._record_guardrail_metrics.call_args.kwargs
|
|
assert metrics_kwargs["status"] == "intervened"
|
|
assert metrics_kwargs["error_type"] is None
|