Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/hooks/sensitive_data_routing.py: 31%
71 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Sensitive Data Routing Hook for LiteLLM Proxy.
4When a guardrail detects sensitive data and is configured with on_sensitive_data='route',
5this hook manages:
61. Storing the routing decision (session_id -> model) in cache
72. Checking incoming requests for existing routing overrides
83. Applying sticky routing so all subsequent requests in a session go to the same model
10Works across multiple proxy instances via DualCache (in-memory + Redis).
11"""
13import logging
14import os
15from typing import TYPE_CHECKING, Any, Final
17from litellm._logging import verbose_proxy_logger
18from litellm.caching.caching import DualCache
19from litellm.caching.redis_cache import log_redis_failure
20from litellm.integrations.custom_guardrail import get_session_id_from_request_data
21from litellm.integrations.custom_logger import CustomLogger
22from litellm.proxy._types import UserAPIKeyAuth
24if TYPE_CHECKING: 24 ↛ 25line 24 didn't jump to line 25 because the condition on line 24 was never true
25 from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
27 InternalUsageCache = _InternalUsageCache
28else:
29 InternalUsageCache = Any
32SENSITIVE_ROUTING_CACHE_PREFIX: Final = "sensitive_route"
33DEFAULT_SENSITIVE_ROUTING_TTL: Final = 3600
36class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
37 """
38 Pre-call hook that checks for existing sensitive data routing overrides
39 and applies them to incoming requests.
41 This hook runs early in the pre-call chain and modifies the request's
42 model field if a routing override exists for the session.
43 """
45 def __init__(self, internal_usage_cache: InternalUsageCache):
46 self.internal_usage_cache = internal_usage_cache
47 self.ttl = int(
48 os.getenv(
49 "LITELLM_SENSITIVE_ROUTING_TTL",
50 str(DEFAULT_SENSITIVE_ROUTING_TTL),
51 )
52 )
54 def _make_cache_key(self, session_id: str, tenant: str) -> str:
55 return f"{{{SENSITIVE_ROUTING_CACHE_PREFIX}:{tenant}:{session_id}}}:model"
57 @staticmethod
58 def _resolve_tenant(user_api_key_dict: UserAPIKeyAuth | None) -> str:
59 """
60 Identify the authenticated principal the routing override belongs to.
62 API-key auth is scoped by the hashed key. JWT (and other keyless) auth
63 has no api_key, so fall back to a stable identity claim. Without this,
64 every keyless caller would share the ``default`` namespace and could read
65 or overwrite another principal's session routing.
66 """
67 if user_api_key_dict is None:
68 return "default"
69 if user_api_key_dict.api_key:
70 return user_api_key_dict.api_key
71 principal: Final = [
72 f"{label}:{value}"
73 for label, value in (
74 ("user", user_api_key_dict.user_id),
75 ("team", user_api_key_dict.team_id),
76 ("org", user_api_key_dict.org_id),
77 )
78 if value
79 ]
80 return "|".join(principal) if principal else "default"
82 async def _get_routed_model(self, session_id: str, user_api_key_dict: UserAPIKeyAuth | None) -> str | None:
83 """Get the model this session should be routed to, if any."""
84 cache_key: Final = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict))
86 if self.internal_usage_cache.dual_cache.redis_cache is not None:
87 try:
88 result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key)
89 if result is not None:
90 routed_model: Final = str(result)
91 remaining_ttl = await self.internal_usage_cache.dual_cache.redis_cache.async_get_ttl(key=cache_key)
92 await self.internal_usage_cache.async_set_cache(
93 key=cache_key,
94 value=routed_model,
95 ttl=remaining_ttl if remaining_ttl is not None else self.ttl,
96 litellm_parent_otel_span=None,
97 local_only=True,
98 )
99 return routed_model
100 except Exception as e:
101 log_redis_failure(
102 verbose_proxy_logger,
103 logging.WARNING,
104 "SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory",
105 e,
106 )
108 result = await self.internal_usage_cache.async_get_cache(
109 key=cache_key,
110 litellm_parent_otel_span=None,
111 local_only=True,
112 )
113 if result is not None:
114 return str(result)
115 return None
117 async def set_session_routing(
118 self,
119 session_id: str,
120 model: str,
121 user_api_key_dict: UserAPIKeyAuth | None = None,
122 guardrail_name: str | None = None,
123 ) -> None:
124 """
125 Store a routing override for a session.
127 Called by guardrails when they detect sensitive data and want to
128 route the session to a specific model. The override is scoped to the
129 requesting principal so sessions from different tenants cannot collide.
130 """
131 cache_key: Final = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict))
133 verbose_proxy_logger.info(
134 "SensitiveDataRoutingHandler: Setting session routing session_id=%s model=%s guardrail=%s ttl=%s",
135 session_id,
136 model,
137 guardrail_name,
138 self.ttl,
139 )
141 if self.internal_usage_cache.dual_cache.redis_cache is not None:
142 try:
143 await self.internal_usage_cache.dual_cache.redis_cache.async_set_cache(
144 key=cache_key,
145 value=model,
146 ttl=self.ttl,
147 )
148 except Exception as e:
149 log_redis_failure(
150 verbose_proxy_logger,
151 logging.WARNING,
152 "SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory",
153 e,
154 )
156 await self.internal_usage_cache.async_set_cache(
157 key=cache_key,
158 value=model,
159 ttl=self.ttl,
160 litellm_parent_otel_span=None,
161 local_only=True,
162 )
164 async def async_pre_call_hook(
165 self,
166 user_api_key_dict: UserAPIKeyAuth,
167 cache: DualCache,
168 data: dict,
169 call_type: str,
170 ) -> Exception | str | dict | None:
171 """
172 Before each LLM call, check if this session has a routing override.
173 If so, modify the request's model field.
174 """
175 session_id: Final = get_session_id_from_request_data(data)
176 if session_id is None: 176 ↛ 179line 176 didn't jump to line 179 because the condition on line 176 was always true
177 return None
179 routed_model: Final = await self._get_routed_model(session_id, user_api_key_dict)
180 if routed_model is None:
181 return None
183 original_model: Final = data.get("model")
184 if original_model == routed_model:
185 return None
187 verbose_proxy_logger.info(
188 "SensitiveDataRoutingHandler: Applying session routing override "
189 "session_id=%s original_model=%s routed_model=%s",
190 session_id,
191 original_model,
192 routed_model,
193 )
195 data["model"] = routed_model
197 metadata: Final = data.get("metadata") or {}
198 metadata["sensitive_data_routing_applied"] = True
199 metadata["sensitive_data_routing_original_model"] = original_model
200 data["metadata"] = metadata
202 return data