scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 71

Support per-model encrypted content affinity config (#29760)

Browse source 
Co-authored-by: shin-berri <shin-laptop@berri.ai>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
This commit is contained in:
kejunleng 2026-06-09 20:09:42 +08:00 • committed by Sameer Kankute
parent a8c835ffc9
commit 31240b5d90
No known key found for this signature in database
6 changed files with 599 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
litellm/router.py
View file

@ -1658,6 +1658,67 @@ class Router:
f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys."
)
def _add_encrypted_content_affinity_check(
self, enable_global_affinity: bool
) -> None:
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
def _move_before_deployment_affinity(
callback_list: List[Any],
callback_to_move: EncryptedContentAffinityCheck,
) -> None:
if callback_to_move not in callback_list:
return
callback_list.remove(callback_to_move)
insert_index = next(
(
idx
for idx, callback in enumerate(callback_list)
if isinstance(callback, DeploymentAffinityCheck)
),
len(callback_list),
)
callback_list.insert(insert_index, callback_to_move)
if (
enable_global_affinity
or EncryptedContentAffinityCheck.has_model_group_affinity_enabled(
self.model_group_affinity_config
)
):
if self.optional_callbacks is None:
self.optional_callbacks = []
existing_ec_callback: Optional[EncryptedContentAffinityCheck] = None
for cb in self.optional_callbacks:
if isinstance(cb, EncryptedContentAffinityCheck):
existing_ec_callback = cb
break
if existing_ec_callback is not None:
existing_ec_callback.router = self
existing_ec_callback.enable_global_affinity = (
existing_ec_callback.enable_global_affinity
or enable_global_affinity
)
existing_ec_callback.model_group_affinity_config = (
self.model_group_affinity_config or {}
)
ec_callback = existing_ec_callback
else:
ec_callback = EncryptedContentAffinityCheck(
router=self,
enable_global_affinity=enable_global_affinity,
model_group_affinity_config=self.model_group_affinity_config,
)
self.optional_callbacks.append(ec_callback)
litellm.logging_callback_manager.add_litellm_callback(ec_callback)
_move_before_deployment_affinity(self.optional_callbacks, ec_callback)
_move_before_deployment_affinity(litellm.callbacks, ec_callback)
def add_optional_pre_call_checks(
self, optional_pre_call_checks: Optional[OptionalPreCallChecks]
):
@ -1721,22 +1782,11 @@ class Router:
# ---------------------------------------------------------------------
# Encrypted content affinity
# ---------------------------------------------------------------------
if "encrypted_content_affinity" in optional_pre_call_checks:
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
self._add_encrypted_content_affinity_check(
enable_global_affinity=(
"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(router=self)
self.optional_callbacks.append(ec_callback)
litellm.logging_callback_manager.add_litellm_callback(ec_callback)
)
# ---------------------------------------------------------------------
# Remaining optional pre-call checks

7
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -39,7 +39,12 @@ class DeploymentAffinityCheck(CustomLogger):
CACHE_KEY_PREFIX = "deployment_affinity:v1"
VALID_FLAGS = frozenset(
{"deployment_affinity", "responses_api_deployment_check", "session_affinity"}
{
"deployment_affinity",
"responses_api_deployment_check",
"session_affinity",
"encrypted_content_affinity",
}
)
def __init__(

36
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -37,7 +37,7 @@ Safe to enable globally:
"""
import time
from typing import TYPE_CHECKING, Any, List, Optional, cast
from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
import httpx
@ -64,17 +64,45 @@ class EncryptedContentAffinityCheck(CustomLogger):
The ``model_id`` is decoded directly from the litellm-encoded item IDs –
no caching or TTL management needed.
Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])``.
Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])`` or
per-model group ``model_group_affinity_config``.
"""
def __init__(self, router: Optional["Router"] = None) -> None:
def __init__(
self,
router: Optional["Router"] = None,
enable_global_affinity: bool = True,
model_group_affinity_config: Optional[Dict[str, List[str]]] = None,
) -> None:
super().__init__()
self.router = router
self.enable_global_affinity = enable_global_affinity
self.model_group_affinity_config: Dict[str, List[str]] = (
model_group_affinity_config or {}
)
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
@staticmethod
def has_model_group_affinity_enabled(
model_group_affinity_config: Optional[Dict[str, List[str]]],
) -> bool:
if not model_group_affinity_config:
return False
return any(
"encrypted_content_affinity" in checks
for checks in model_group_affinity_config.values()
)
def _is_enabled_for_model_group(self, model_group: str) -> bool:
group_checks = self.model_group_affinity_config.get(model_group)
return self.enable_global_affinity or (
group_checks is not None and "encrypted_content_affinity" in group_checks
)
@staticmethod
def _extract_model_id_from_input(request_input: Any) -> Optional[str]:
"""
@ -213,6 +241,8 @@ class EncryptedContentAffinityCheck(CustomLogger):
"""
request_kwargs = request_kwargs or {}
typed_healthy_deployments = cast(List[dict], healthy_deployments)
if not self._is_enabled_for_model_group(model):
return typed_healthy_deployments
# Signal to the response post-processor that encrypted item IDs should be
# encoded in the output of this request. Only set the flag when

57
tests/router_unit_tests/test_router_endpoints.py
View file

@ -1236,3 +1236,60 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
def test_router_model_group_encrypted_content_affinity_callback_registration():
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,
)
model_group = "openai.gpt-5.1-codex"
model_group_affinity_config = {
model_group: ["encrypted_content_affinity"],
}
router = Router(
model_list=[
{
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key",
},
}
],
model_group_affinity_config=model_group_affinity_config,
num_retries=0,
)
try:
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is False
assert (
encrypted_content_callbacks[0].model_group_affinity_config
== model_group_affinity_config
)
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
deployment_callback
)
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is True
assert encrypted_content_callbacks[0].router is router
finally:
router.discard()

188
tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -1471,3 +1471,191 @@ async def test_affinity_does_not_raise_when_boundary_peer_available():
assert result == [peer]
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
@pytest.mark.asyncio
async def test_model_group_affinity_config_enables_encrypted_content_affinity():
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
model_group = "openai.gpt-5.1-codex"
target_deployment = {
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-b"},
}
healthy_deployments = [
{
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-a"},
},
target_deployment,
]
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
"deployment-b", "rs_test"
)
request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {},
}
check = EncryptedContentAffinityCheck(
enable_global_affinity=False,
model_group_affinity_config={
model_group: ["encrypted_content_affinity"],
},
)
filtered = await check.async_filter_deployments(
model=model_group,
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs=request_kwargs,
)
assert filtered == [target_deployment]
assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"]
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
@pytest.mark.asyncio
async def test_model_group_affinity_config_does_not_disable_global_encrypted_content_affinity():
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
model_group = "openai.gpt-5.1-codex"
target_deployment = {
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-b"},
}
healthy_deployments = [
{
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-a"},
},
target_deployment,
]
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
"deployment-b", "rs_test"
)
request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {},
}
check = EncryptedContentAffinityCheck(
enable_global_affinity=True,
model_group_affinity_config={
model_group: ["deployment_affinity"],
},
)
filtered = await check.async_filter_deployments(
model=model_group,
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs=request_kwargs,
)
assert filtered == [target_deployment]
assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"]
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
@pytest.mark.asyncio
async def test_model_group_encrypted_content_affinity_overrides_global_deployment_affinity():
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,
)
model_group = "openai.gpt-5.1-codex"
user_api_key_hash = "test-user-key"
deployment_a = {
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-a",
},
"model_info": {"id": "deployment-a"},
}
deployment_b = {
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-b",
},
"model_info": {"id": "deployment-b"},
}
router = litellm.Router(
model_list=[deployment_a, deployment_b],
optional_pre_call_checks=["deployment_affinity"],
model_group_affinity_config={
model_group: ["encrypted_content_affinity"],
},
num_retries=0,
)
try:
callbacks = router.optional_callbacks or []
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
encrypted_content_callback = next(
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
)
assert callbacks.index(encrypted_content_callback) < callbacks.index(
deployment_callback
)
assert encrypted_content_callback.enable_global_affinity is False
cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=model_group,
user_key=user_api_key_hash,
)
await deployment_callback.cache.async_set_cache(
key=cache_key,
value={"model_id": "deployment-a"},
ttl=60,
)
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
"deployment-b", "rs_test"
)
request_kwargs = {
"input": [
{
"type": "reasoning",
"id": encoded_id,
"encrypted_content": "gAAAAABpnW_yEYmSNEyOG...",
}
],
"metadata": {"user_api_key_hash": user_api_key_hash},
"litellm_metadata": {},
}
after_deployment_affinity = await deployment_callback.async_filter_deployments(
model=model_group,
healthy_deployments=[deployment_a, deployment_b],
messages=None,
request_kwargs=request_kwargs,
)
assert after_deployment_affinity == [deployment_a, deployment_b]
after_encrypted_content_affinity = (
await encrypted_content_callback.async_filter_deployments(
model=model_group,
healthy_deployments=after_deployment_affinity,
messages=None,
request_kwargs=request_kwargs,
)
)
assert after_encrypted_content_affinity == [deployment_b]
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
finally:
router.discard()

250
tests/test_litellm/test_router.py
View file

@ -80,6 +80,256 @@ def test_router_with_model_info_and_model_group():
)
def test_router_model_group_encrypted_content_affinity_callback_registration():
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,
)
model_group = "openai.gpt-5.1-codex"
model_group_affinity_config = {
model_group: ["encrypted_content_affinity"],
}
original_callbacks = list(litellm.callbacks)
litellm.callbacks = []
router = None
try:
router = litellm.Router(
model_list=[
{
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key",
},
}
],
model_group_affinity_config=model_group_affinity_config,
num_retries=0,
)
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is False
assert (
encrypted_content_callbacks[0].model_group_affinity_config
== model_group_affinity_config
)
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
deployment_callback
)
assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (
litellm.callbacks.index(deployment_callback)
)
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is True
assert encrypted_content_callbacks[0].router is router
finally:
if router is not None:
router.discard()
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_encrypted_content_affinity_model_group_config_is_additive():
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
model_group = "openai.gpt-5.1-codex"
target_deployment = {
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-b"},
}
healthy_deployments = [
{
"model_name": model_group,
"litellm_params": {"model": "openai/gpt-5.1-codex"},
"model_info": {"id": "deployment-a"},
},
target_deployment,
]
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
"deployment-b", "rs_test"
)
assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled(
{model_group: ["encrypted_content_affinity"]}
)
assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None)
per_group_check = EncryptedContentAffinityCheck(
enable_global_affinity=False,
model_group_affinity_config={
model_group: ["encrypted_content_affinity"],
},
)
request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {},
}
filtered = await per_group_check.async_filter_deployments(
model=model_group,
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs=request_kwargs,
)
assert filtered == [target_deployment]
assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"]
disabled_check = EncryptedContentAffinityCheck(
enable_global_affinity=False,
model_group_affinity_config={
"other-model-group": ["encrypted_content_affinity"],
},
)
disabled_request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {},
}
unfiltered = await disabled_check.async_filter_deployments(
model=model_group,
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs=disabled_request_kwargs,
)
assert unfiltered == healthy_deployments
assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs[
"litellm_metadata"
]
global_check = EncryptedContentAffinityCheck(
enable_global_affinity=True,
model_group_affinity_config={
model_group: ["deployment_affinity"],
},
)
global_request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {},
}
globally_filtered = await global_check.async_filter_deployments(
model=model_group,
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs=global_request_kwargs,
)
assert globally_filtered == [target_deployment]
assert global_request_kwargs["litellm_metadata"][
"encrypted_content_affinity_enabled"
]
@pytest.mark.asyncio
async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity():
from litellm.responses.utils import ResponsesAPIRequestUtils
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,
)
model_group = "openai.gpt-5.1-codex"
user_api_key_hash = "test-user-key"
deployment_a = {
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-a",
},
"model_info": {"id": "deployment-a"},
}
deployment_b = {
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-b",
},
"model_info": {"id": "deployment-b"},
}
original_callbacks = list(litellm.callbacks)
litellm.callbacks = []
router = None
try:
router = litellm.Router(
model_list=[deployment_a, deployment_b],
model_group_affinity_config={
model_group: [
"deployment_affinity",
"encrypted_content_affinity",
],
},
num_retries=0,
)
callbacks = router.optional_callbacks or []
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
encrypted_content_callback = next(
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
)
assert callbacks.index(encrypted_content_callback) < callbacks.index(
deployment_callback
)
assert litellm.callbacks.index(encrypted_content_callback) < (
litellm.callbacks.index(deployment_callback)
)
cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=model_group,
user_key=user_api_key_hash,
)
await deployment_callback.cache.async_set_cache(
key=cache_key,
value={"model_id": "deployment-a"},
ttl=60,
)
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(
"deployment-b", "rs_test"
)
request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],
"litellm_metadata": {"user_api_key_hash": user_api_key_hash},
}
filtered = await router.async_callback_filter_deployments(
model=model_group,
healthy_deployments=[deployment_a, deployment_b],
messages=None,
parent_otel_span=None,
request_kwargs=request_kwargs,
)
assert filtered == [deployment_b]
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
finally:
if router is not None:
router.discard()
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_arouter_with_tags_and_fallbacks():
"""