Add encrypted_content_affinity in router

This commit is contained in:
Sameer Kankute 2026-02-24 16:21:20 +05:30
parent b44e29fd36
commit 8bdd5c3081
2 changed files with 31 additions and 0 deletions

View file

@ -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"

View file

@ -813,6 +813,7 @@ OptionalPreCallChecks = List[
"session_affinity",
"forward_client_headers_by_model_group",
"enforce_model_rate_limits",
"encrypted_content_affinity",
]
]