diff --git a/litellm/router.py b/litellm/router.py index d89a5099b01..b3a88e738c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -114,6 +114,9 @@ from litellm.router_utils.handle_error import ( from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, +) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -1247,6 +1250,25 @@ class Router: self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + # --------------------------------------------------------------------- + # Encrypted content affinity + # --------------------------------------------------------------------- + if "encrypted_content_affinity" in optional_pre_call_checks: + if self.optional_callbacks is None: + self.optional_callbacks = [] + + already_registered = any( + isinstance(cb, EncryptedContentAffinityCheck) + for cb in self.optional_callbacks + ) + if not already_registered: + ec_callback = EncryptedContentAffinityCheck( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + ) + self.optional_callbacks.append(ec_callback) + litellm.logging_callback_manager.add_litellm_callback(ec_callback) + # --------------------------------------------------------------------- # Remaining optional pre-call checks # --------------------------------------------------------------------- @@ -1256,6 +1278,7 @@ class Router: "deployment_affinity", "responses_api_deployment_check", "session_affinity", + "encrypted_content_affinity", ): continue if pre_call_check == "prompt_caching": @@ -8661,6 +8684,13 @@ class Router: if isinstance(healthy_deployments, dict): return healthy_deployments + # When encrypted content affinity pins to a specific deployment, + if ( + request_kwargs.get("_encrypted_content_affinity_pinned") + and len(healthy_deployments) == 1 + ): + return healthy_deployments[0] + start_time = time.time() if ( self.routing_strategy == "usage-based-routing-v2" diff --git a/litellm/types/router.py b/litellm/types/router.py index aa4d7bd9a97..fca731d1f91 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -813,6 +813,7 @@ OptionalPreCallChecks = List[ "session_affinity", "forward_client_headers_by_model_group", "enforce_model_rate_limits", + "encrypted_content_affinity", ] ]