tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-16 14:23:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #19143 from emerzon/affinity_callback

Browse Source 
Add deployment affinity routing callback
This commit is contained in:
Sameer Kankute
2026-02-18 22:21:02 +05:30
committed by GitHub
parent 45d3d1adfd 3dd55a7b61
commit a7da7b53f0
7 changed files with 1169 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/config_settings.md
+2 -1
View File
@@ -358,7 +358,8 @@ router_settings:
| redis_url | str | URL for Redis server. **Known performance issue with Redis URL.** |
| cache_responses | boolean | Flag to enable caching LLM Responses, if cache set under `router_settings`. If true, caches responses. Defaults to False. |
| router_general_settings | RouterGeneralSettings | [SDK-Only] Router general settings - contains optimizations like 'async_only_mode'. [Docs](../routing.md#router-general-settings) |
| optional_pre_call_checks | List[str] | List of pre-call checks to add to the router. Currently supported: 'router_budget_limiting', 'prompt_caching' |
| optional_pre_call_checks | List[str] | List of pre-call checks to add to the router. Supported: `router_budget_limiting`, `prompt_caching`, `responses_api_deployment_check`, `deployment_affinity`, `forward_client_headers_by_model_group` |
| deployment_affinity_ttl_seconds | int | TTL (seconds) for user-key → deployment affinity mapping when `deployment_affinity` is enabled (configured at Router init / proxy startup). Defaults to `3600` (1 hour). |
| ignore_invalid_deployments | boolean | If true, ignores invalid deployments. Default for proxy is True - to prevent invalid models from blocking other models from being loaded. |
| search_tools | List[SearchToolTypedDict] | List of search tool configurations for Search API integration. Each tool specifies a search_tool_name and litellm_params with search_provider, api_key, api_base, etc. [Further Docs](../search.md) |
| guardrail_list | List[GuardrailTypedDict] | List of guardrail configurations for guardrail load balancing. Enables load balancing across multiple guardrail deployments with the same guardrail_name. [Further Docs](./guardrails/guardrail_load_balancing.md) |
docs/my-website/docs/response_api.md
+21 -8
View File
@@ -884,7 +884,12 @@ router = litellm.Router(
},
},
],
optional_pre_call_checks=["responses_api_deployment_check"],
# `responses_api_deployment_check` ensures Requests with `previous_response_id`
# are routed to the same deployment. `deployment_affinity` adds sticky sessions
# for requests without `previous_response_id` (useful for implicit caching).
optional_pre_call_checks=["responses_api_deployment_check", "deployment_affinity"],
# Optional (default is 3600 seconds / 1 hour)
deployment_affinity_ttl_seconds=3600,
)
# Initial request
@@ -911,7 +916,16 @@ follow_up = await router.aresponses(
#### 1. Setup session continuity on proxy config.yaml
To enable session continuity for Responses API in your LiteLLM proxy, set `optional_pre_call_checks: ["responses_api_deployment_check"]` in your proxy config.yaml.
To enable session continuity for Responses API in your LiteLLM proxy, set `optional_pre_call_checks` in your proxy config.yaml.
- `responses_api_deployment_check`: high priority routing when `previous_response_id` is provided
- `deployment_affinity`: sticky sessions based on user key (applies even without `previous_response_id`)
Notes:
- User-key affinity is keyed on `metadata.user_api_key_hash` (the API key hash). The OpenAI `user` request parameter is an end-user identifier and is intentionally not used for deployment affinity.
- `user_api_key_hash` is already SHA-256, and is used as-is (no double hashing).
- Affinity is scoped by a stable model identifier (the model-map key, e.g. `model_map_information.model_map_key`) so model aliases map to the same stickiness bucket.
- The mapping TTL is controlled by `deployment_affinity_ttl_seconds` (configured on Router init / proxy startup).
```yaml showLineNumbers title="config.yaml with Session Continuity"
model_list:
@@ -929,7 +943,11 @@ model_list:
api_base: https://endpoint2.openai.azure.com
router_settings:
optional_pre_call_checks: ["responses_api_deployment_check"]
optional_pre_call_checks:
- responses_api_deployment_check
- deployment_affinity
# Optional (default is 3600 seconds / 1 hour)
deployment_affinity_ttl_seconds: 3600
```
#### 2. Use the OpenAI Python SDK to make requests to LiteLLM Proxy
@@ -1356,8 +1374,3 @@ Response:
litellm/router.py
+78 -23
View File
@@ -113,12 +113,12 @@ from litellm.router_utils.handle_error import (
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
ModelRateLimitingCheck,
)
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
PromptCachingDeploymentCheck,
)
from litellm.router_utils.pre_call_checks.responses_api_deployment_check import (
ResponsesApiDeploymentCheck,
)
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
increment_deployment_failures_for_current_minute,
increment_deployment_successes_for_current_minute,
@@ -293,6 +293,7 @@ class Router:
router_general_settings: Optional[
RouterGeneralSettings
] = RouterGeneralSettings(),
deployment_affinity_ttl_seconds: int = 3600,
ignore_invalid_deployments: bool = False,
) -> None:
"""
@@ -326,6 +327,7 @@ class Router:
routing_strategy_args (dict): Additional args for latency-based routing. Defaults to {}.
alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None.
provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None.
deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600.
ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error.
Returns:
Router: An instance of the litellm.Router class.
@@ -604,6 +606,7 @@ class Router:
litellm.failure_callback = [self.deployment_callback_on_failure]
self.routing_strategy_args = routing_strategy_args
self.provider_budget_config = provider_budget_config
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.router_budget_logger: Optional[RouterBudgetLimiting] = None
if RouterBudgetLimiting.should_init_router_budget_limiter(
model_list=model_list, provider_budget_config=self.provider_budget_config
@@ -1184,26 +1187,78 @@ class Router:
def add_optional_pre_call_checks(
self, optional_pre_call_checks: Optional[OptionalPreCallChecks]
):
if optional_pre_call_checks is not None:
for pre_call_check in optional_pre_call_checks:
_callback: Optional[CustomLogger] = None
if pre_call_check == "prompt_caching":
_callback = PromptCachingDeploymentCheck(cache=self.cache)
elif pre_call_check == "router_budget_limiting":
_callback = RouterBudgetLimiting(
dual_cache=self.cache,
provider_budget_config=self.provider_budget_config,
model_list=self.model_list,
)
elif pre_call_check == "responses_api_deployment_check":
_callback = ResponsesApiDeploymentCheck()
elif pre_call_check == "enforce_model_rate_limits":
_callback = ModelRateLimitingCheck(dual_cache=self.cache)
if _callback is not None:
if self.optional_callbacks is None:
self.optional_callbacks = []
self.optional_callbacks.append(_callback)
litellm.logging_callback_manager.add_litellm_callback(_callback)
if optional_pre_call_checks is None:
return
# ---------------------------------------------------------------------
# Unified deployment affinity (session stickiness)
# ---------------------------------------------------------------------
enable_user_key_affinity = "deployment_affinity" in optional_pre_call_checks
enable_responses_api_affinity = (
"responses_api_deployment_check" in optional_pre_call_checks
)
if enable_user_key_affinity or enable_responses_api_affinity:
if self.optional_callbacks is None:
self.optional_callbacks = []
existing_affinity_callback: Optional[DeploymentAffinityCheck] = None
for cb in self.optional_callbacks:
if isinstance(cb, DeploymentAffinityCheck):
existing_affinity_callback = cb
break
if existing_affinity_callback is not None:
existing_affinity_callback.enable_user_key_affinity = (
existing_affinity_callback.enable_user_key_affinity
or enable_user_key_affinity
)
existing_affinity_callback.enable_responses_api_affinity = (
existing_affinity_callback.enable_responses_api_affinity
or enable_responses_api_affinity
)
existing_affinity_callback.ttl_seconds = (
self.deployment_affinity_ttl_seconds
)
else:
affinity_callback = DeploymentAffinityCheck(
cache=self.cache,
ttl_seconds=self.deployment_affinity_ttl_seconds,
enable_user_key_affinity=enable_user_key_affinity,
enable_responses_api_affinity=enable_responses_api_affinity,
)
self.optional_callbacks.append(affinity_callback)
litellm.logging_callback_manager.add_litellm_callback(
affinity_callback
)
# ---------------------------------------------------------------------
# Remaining optional pre-call checks
# ---------------------------------------------------------------------
for pre_call_check in optional_pre_call_checks:
_callback: Optional[CustomLogger] = None
if pre_call_check in (
"deployment_affinity",
"responses_api_deployment_check",
):
continue
if pre_call_check == "prompt_caching":
_callback = PromptCachingDeploymentCheck(cache=self.cache)
elif pre_call_check == "router_budget_limiting":
_callback = RouterBudgetLimiting(
dual_cache=self.cache,
provider_budget_config=self.provider_budget_config,
model_list=self.model_list,
)
elif pre_call_check == "enforce_model_rate_limits":
_callback = ModelRateLimitingCheck(dual_cache=self.cache)
if _callback is None:
continue
if self.optional_callbacks is None:
self.optional_callbacks = []
self.optional_callbacks.append(_callback)
litellm.logging_callback_manager.add_litellm_callback(_callback)
def print_deployment(self, deployment: dict):
"""
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
+396
View File
@@ -0,0 +1,396 @@
"""
Unified deployment affinity (session stickiness) for the Router.
Features (independently enable-able):
1. Responses API continuity: when a `previous_response_id` is provided, route to the
deployment that generated the original response (highest priority).
2. API-key affinity: map an API key hash -> deployment id for a TTL and re-use that
deployment for subsequent requests to the same router deployment model name
(alias-safe, aligns to `model_map_information.model_map_key`).
This is designed to support "implicit prompt caching" scenarios (no explicit cache_control),
where routing to a consistent deployment is still beneficial.
"""
import hashlib
from typing import Any, Dict, List, Optional, cast
from typing_extensions import TypedDict
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes
class DeploymentAffinityCacheValue(TypedDict):
model_id: str
class DeploymentAffinityCheck(CustomLogger):
"""
Router deployment affinity callback.
NOTE: This is a Router-only callback intended to be wired through
`Router(optional_pre_call_checks=[...])`.
"""
CACHE_KEY_PREFIX = "deployment_affinity:v1"
def __init__(
self,
cache: DualCache,
ttl_seconds: int,
enable_user_key_affinity: bool,
enable_responses_api_affinity: bool,
):
super().__init__()
self.cache = cache
self.ttl_seconds = ttl_seconds
self.enable_user_key_affinity = enable_user_key_affinity
self.enable_responses_api_affinity = enable_responses_api_affinity
@staticmethod
def _looks_like_sha256_hex(value: str) -> bool:
if len(value) != 64:
return False
try:
int(value, 16)
except ValueError:
return False
return True
@staticmethod
def _hash_user_key(user_key: str) -> str:
"""
Hash user identifiers before storing them in cache keys.
This avoids putting raw API keys / user identifiers into Redis keys (and therefore
into logs/metrics), while keeping the cache key stable and a fixed length.
"""
# If the proxy already provides a stable SHA-256 (e.g. `metadata.user_api_key_hash`),
# keep it as-is to avoid double-hashing and to make correlation/debugging possible.
if DeploymentAffinityCheck._looks_like_sha256_hex(user_key):
return user_key.lower()
return hashlib.sha256(user_key.encode("utf-8")).hexdigest()
@staticmethod
def _get_model_map_key_from_litellm_model_name(litellm_model_name: str) -> Optional[str]:
"""
Best-effort derivation of a stable "model map key" for affinity scoping.
The intent is to align with `standard_logging_payload.model_map_information.model_map_key`,
which is typically the base model identifier (stable across deployments/endpoints).
Notes:
- When the model name is in "provider/model" format, the provider prefix is stripped.
- For Azure, the string after "azure/" is commonly an *Azure deployment name*, which may
differ across instances. If `base_model` is not explicitly set, we skip deriving a
model-map key from the model string to avoid generating unstable keys.
"""
if not litellm_model_name:
return None
if "/" not in litellm_model_name:
return litellm_model_name
provider_prefix, remainder = litellm_model_name.split("/", 1)
if provider_prefix == "azure":
return None
return remainder
@staticmethod
def _get_model_map_key_from_deployment(deployment: dict) -> Optional[str]:
"""
Derive a stable model-map key from a router deployment dict.
Primary source: `deployment.model_name` (Router's canonical group name after
alias resolution). This is stable across provider-specific deployments (e.g.,
Azure/Vertex/Bedrock for the same logical model) and aligns with
`model_map_information.model_map_key` in standard logging.
Prefer `base_model` when available (important for Azure), otherwise fall back to
parsing `litellm_params.model`.
"""
model_name = deployment.get("model_name")
if isinstance(model_name, str) and model_name:
return model_name
model_info = deployment.get("model_info")
if isinstance(model_info, dict):
base_model = model_info.get("base_model")
if isinstance(base_model, str) and base_model:
return base_model
litellm_params = deployment.get("litellm_params")
if isinstance(litellm_params, dict):
base_model = litellm_params.get("base_model")
if isinstance(base_model, str) and base_model:
return base_model
litellm_model_name = litellm_params.get("model")
if isinstance(litellm_model_name, str) and litellm_model_name:
return DeploymentAffinityCheck._get_model_map_key_from_litellm_model_name(
litellm_model_name
)
return None
@staticmethod
def _get_stable_model_map_key_from_deployments(
healthy_deployments: List[dict],
) -> Optional[str]:
"""
Only use model-map key scoping when it is stable across the deployment set.
This prevents accidentally keying on per-deployment identifiers like Azure deployment
names (when `base_model` is not configured).
"""
if not healthy_deployments:
return None
keys: List[str] = []
for deployment in healthy_deployments:
key = DeploymentAffinityCheck._get_model_map_key_from_deployment(deployment)
if key is None:
return None
keys.append(key)
unique_keys = set(keys)
if len(unique_keys) != 1:
return None
return keys[0]
@staticmethod
def _shorten_for_logs(value: str, keep: int = 8) -> str:
if len(value) <= keep:
return value
return f"{value[:keep]}..."
@classmethod
def get_affinity_cache_key(cls, model_group: str, user_key: str) -> str:
hashed_user_key = cls._hash_user_key(user_key=user_key)
return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}"
@staticmethod
def _get_user_key_from_metadata_dict(metadata: dict) -> Optional[str]:
# NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the
# OpenAI `user` parameter, which is an end-user identifier).
user_key = metadata.get("user_api_key_hash")
if user_key is None:
return None
return str(user_key)
@staticmethod
def _iter_metadata_dicts(request_kwargs: dict) -> List[dict]:
"""
Return all metadata dicts available on the request.
Depending on the endpoint, Router may populate `metadata` or `litellm_metadata`.
Users may also send one or both, so we check both (rather than using `or`).
"""
metadata_dicts: List[dict] = []
for key in ("litellm_metadata", "metadata"):
md = request_kwargs.get(key)
if isinstance(md, dict):
metadata_dicts.append(md)
return metadata_dicts
@staticmethod
def _get_user_key_from_request_kwargs(request_kwargs: dict) -> Optional[str]:
"""
Extract a stable affinity key from request kwargs.
Source (proxy): `metadata.user_api_key_hash`
Note: the OpenAI `user` parameter is an end-user identifier and is intentionally
not used for deployment affinity.
"""
# Check metadata dicts (Proxy usage)
for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(
metadata=metadata
)
if user_key is not None:
return user_key
return None
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: List[dict], model_id: str) -> Optional[dict]:
for deployment in healthy_deployments:
model_info = deployment.get("model_info")
if not isinstance(model_info, dict):
continue
deployment_model_id = model_info.get("id")
if deployment_model_id is not None and str(deployment_model_id) == str(model_id):
return deployment
return None
async def async_filter_deployments(
self,
model: str,
healthy_deployments: List,
messages: Optional[List[AllMessageValues]],
request_kwargs: Optional[dict] = None,
parent_otel_span: Optional[Span] = None,
) -> List[dict]:
"""
Optionally filter healthy deployments based on:
1. `previous_response_id` (Responses API continuity) [highest priority]
2. cached API-key deployment affinity
"""
request_kwargs = request_kwargs or {}
typed_healthy_deployments = cast(List[dict], healthy_deployments)
# 1) Responses API continuity (high priority)
if self.enable_responses_api_affinity:
previous_response_id = request_kwargs.get("previous_response_id")
if previous_response_id is not None:
responses_model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id(str(previous_response_id))
if responses_model_id is not None:
deployment = self._find_deployment_by_model_id(
healthy_deployments=typed_healthy_deployments,
model_id=responses_model_id,
)
if deployment is not None:
verbose_router_logger.debug(
"DeploymentAffinityCheck: previous_response_id pinning -> deployment=%s",
responses_model_id,
)
return [deployment]
# 2) User key -> deployment affinity
if not self.enable_user_key_affinity:
return typed_healthy_deployments
user_key = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if user_key is None:
return typed_healthy_deployments
stable_model_map_key = self._get_stable_model_map_key_from_deployments(
healthy_deployments=typed_healthy_deployments
)
if stable_model_map_key is None:
return typed_healthy_deployments
cache_key = self.get_affinity_cache_key(
model_group=stable_model_map_key, user_key=user_key
)
cache_result = await self.cache.async_get_cache(key=cache_key)
model_id: Optional[str] = None
if isinstance(cache_result, dict):
model_id = cast(Optional[str], cache_result.get("model_id"))
elif isinstance(cache_result, str):
# Backwards / safety: allow raw string values.
model_id = cache_result
if not model_id:
return typed_healthy_deployments
deployment = self._find_deployment_by_model_id(
healthy_deployments=typed_healthy_deployments,
model_id=model_id,
)
if deployment is None:
verbose_router_logger.debug(
"DeploymentAffinityCheck: pinned deployment=%s not found in healthy_deployments",
model_id,
)
return typed_healthy_deployments
verbose_router_logger.debug(
"DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s",
model_id,
self._shorten_for_logs(user_key),
)
return [deployment]
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
"""
Persist/update the API-key -> deployment mapping for this request.
Why pre-call?
- LiteLLM runs async success callbacks via a background logging worker for performance.
- We want affinity to be immediately available for subsequent requests.
"""
if not self.enable_user_key_affinity:
return None
user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
if user_key is None:
return None
metadata_dicts = self._iter_metadata_dicts(kwargs)
model_info = kwargs.get("model_info")
if not isinstance(model_info, dict):
model_info = None
if model_info is None:
for metadata in metadata_dicts:
maybe_model_info = metadata.get("model_info")
if isinstance(maybe_model_info, dict):
model_info = maybe_model_info
break
if model_info is None:
# Router sets `model_info` after selecting a deployment. If it's missing, this is
# likely a non-router call or a call path that doesn't support affinity.
return None
model_id = model_info.get("id")
if not model_id:
verbose_router_logger.warning(
"DeploymentAffinityCheck: model_id missing; skipping affinity cache update."
)
return None
# Scope affinity by the Router deployment model name (alias-safe, consistent across
# heterogeneous providers, and matches standard logging's `model_map_key`).
deployment_model_name: Optional[str] = None
for metadata in metadata_dicts:
maybe_deployment_model_name = metadata.get("deployment_model_name")
if isinstance(maybe_deployment_model_name, str) and maybe_deployment_model_name:
deployment_model_name = maybe_deployment_model_name
break
if not deployment_model_name:
verbose_router_logger.warning(
"DeploymentAffinityCheck: deployment_model_name missing; skipping affinity cache update. model_id=%s",
model_id,
)
return None
try:
cache_key = self.get_affinity_cache_key(
model_group=deployment_model_name, user_key=user_key
)
await self.cache.async_set_cache(
cache_key,
DeploymentAffinityCacheValue(model_id=str(model_id)),
ttl=self.ttl_seconds,
)
verbose_router_logger.debug(
"DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
self._shorten_for_logs(user_key),
)
except Exception as e:
# Non-blocking: affinity is a best-effort optimization.
verbose_router_logger.debug(
"DeploymentAffinityCheck: failed to set affinity cache. model_map_key=%s error=%s",
deployment_model_name,
e,
)
return None
litellm/router_utils/pre_call_checks/responses_api_deployment_check.py
+12
View File
@@ -10,6 +10,7 @@ This is different from the normal behavior of the router, which does not have ro
If previous_response_id is provided, route to the deployment that returned the previous response
"""
import warnings
from typing import List, Optional
from litellm.integrations.custom_logger import CustomLogger, Span
@@ -18,6 +19,17 @@ from litellm.types.llms.openai import AllMessageValues
class ResponsesApiDeploymentCheck(CustomLogger):
def __init__(self) -> None:
super().__init__()
warnings.warn(
(
"ResponsesApiDeploymentCheck is deprecated. "
"Use DeploymentAffinityCheck(enable_responses_api_affinity=True) instead."
),
DeprecationWarning,
stacklevel=2,
)
async def async_filter_deployments(
self,
model: str,
litellm/types/router.py
+1
View File
@@ -802,6 +802,7 @@ OptionalPreCallChecks = List[
"prompt_caching",
"router_budget_limiting",
"responses_api_deployment_check",
"deployment_affinity",
"forward_client_headers_by_model_group",
"enforce_model_rate_limits",
]
tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
+659
View File
@@ -0,0 +1,659 @@
import asyncio
import os
import sys
from unittest.mock import AsyncMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import json
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
self.headers = {}
def json(self):
return self._json_data
@pytest.mark.asyncio
async def test_async_user_key_affinity_routes_to_same_deployment():
"""
When deployment_affinity is enabled, subsequent requests from the same user key
should route to the same deployment (even if the routing strategy would pick another).
"""
mock_response_data = {
"id": "resp_mock-resp-123",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [
{
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "Hello there!", "annotations": []}],
}
],
"parallel_tool_calls": True,
"usage": {
"input_tokens": 5,
"output_tokens": 10,
"total_tokens": 15,
"output_tokens_details": {"reasoning_tokens": 0},
},
"text": {"format": {"type": "text"}},
"error": None,
"incomplete_details": None,
"instructions": None,
"metadata": {},
"temperature": 1.0,
"tool_choice": "auto",
"tools": [],
"top_p": 1.0,
"max_output_tokens": None,
"previous_response_id": None,
"reasoning": {"effort": None, "summary": None},
"truncation": "disabled",
"user": None,
}
router = litellm.Router(
model_list=[
{
"model_name": "azure-computer-use-preview",
"litellm_params": {
"model": "azure/computer-use-preview-1",
"api_key": "mock-api-key-1",
"api_version": "mock-api-version",
"api_base": "https://mock-endpoint-1.openai.azure.com",
},
# Required for stable affinity scoping across multiple Azure deployments
"model_info": {"base_model": "computer-use-preview"},
},
{
"model_name": "azure-computer-use-preview",
"litellm_params": {
"model": "azure/computer-use-preview-2",
"api_key": "mock-api-key-2",
"api_version": "mock-api-version-2",
"api_base": "https://mock-endpoint-2.openai.azure.com",
},
"model_info": {"base_model": "computer-use-preview"},
},
],
optional_pre_call_checks=["deployment_affinity"],
)
model_group = "azure-computer-use-preview"
user_api_key_hash = "test-user-key-1"
# Deterministic routing: first selection uses seq[0], second selection attempts seq[1]
# unless the list has been filtered to length=1 by deployment affinity.
choice_calls = {"count": 0}
def deterministic_choice(seq):
choice_calls["count"] += 1
if choice_calls["count"] == 1:
return seq[0]
return seq[1] if len(seq) > 1 else seq[0]
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post, patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=deterministic_choice,
):
mock_post.return_value = MockResponse(mock_response_data, 200)
first_response = await router.aresponses(
model=model_group,
input="Hello, how are you?",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
first_model_id = first_response._hidden_params["model_id"]
# If affinity works, second request should be pinned to the same deployment
# even though deterministic_choice would pick the other deployment when len(seq)>1.
second_response = await router.aresponses(
model=model_group,
input="Follow-up question",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
assert second_response._hidden_params["model_id"] == first_model_id
@pytest.mark.asyncio
async def test_async_user_key_affinity_routes_with_model_group_alias():
"""
When Router model_group_alias is used, the requested model group (alias) can differ
from the internally-routed model group. Deployment affinity should still stick.
"""
mock_response_data = {
"id": "resp_mock-resp-alias",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [
{
"type": "message",
"id": "msg_alias",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "Alias Response"}],
}
],
"parallel_tool_calls": True,
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
"text": {"format": {"type": "text"}},
"error": None,
"previous_response_id": None,
}
canonical_model_group = "azure-computer-use-preview"
alias_model_group = "azure-computer-use-preview-alias"
user_api_key_hash = "test-user-key-alias"
router = litellm.Router(
model_list=[
{
"model_name": canonical_model_group,
"litellm_params": {
"model": "azure/computer-use-preview-1",
"api_key": "mock-api-key-1",
"api_version": "mock-api-version",
"api_base": "https://mock-endpoint-1.openai.azure.com",
},
"model_info": {"base_model": "computer-use-preview"},
},
{
"model_name": canonical_model_group,
"litellm_params": {
"model": "azure/computer-use-preview-2",
"api_key": "mock-api-key-2",
"api_version": "mock-api-version-2",
"api_base": "https://mock-endpoint-2.openai.azure.com",
},
"model_info": {"base_model": "computer-use-preview"},
},
],
model_group_alias={alias_model_group: canonical_model_group},
optional_pre_call_checks=["deployment_affinity"],
)
choice_calls = {"count": 0}
def deterministic_choice(seq):
choice_calls["count"] += 1
if choice_calls["count"] == 1:
return seq[0]
return seq[1] if len(seq) > 1 else seq[0]
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post, patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=deterministic_choice,
):
mock_post.return_value = MockResponse(mock_response_data, 200)
first_response = await router.aresponses(
model=alias_model_group,
input="Hello",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
first_model_id = first_response._hidden_params["model_id"]
second_response = await router.aresponses(
model=alias_model_group,
input="Follow-up",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
assert second_response._hidden_params["model_id"] == first_model_id
@pytest.mark.asyncio
async def test_async_previous_response_id_priority_over_user_key_affinity():
"""
If both deployment_affinity and responses_api_deployment_check are enabled,
`previous_response_id` routing should take priority over user-key affinity.
"""
mock_response_data = {
"id": "resp_mock-resp-456",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [
{
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "I'm doing well, thank you for asking!",
"annotations": [],
}
],
}
],
"parallel_tool_calls": True,
"usage": {
"input_tokens": 10,
"output_tokens": 20,
"total_tokens": 30,
"output_tokens_details": {"reasoning_tokens": 0},
},
"text": {"format": {"type": "text"}},
"error": None,
"incomplete_details": None,
"instructions": None,
"metadata": {},
"temperature": 1.0,
"tool_choice": "auto",
"tools": [],
"top_p": 1.0,
"max_output_tokens": None,
"previous_response_id": None,
"reasoning": {"effort": None, "summary": None},
"truncation": "disabled",
"user": None,
}
router = litellm.Router(
model_list=[
{
"model_name": "azure-computer-use-preview",
"litellm_params": {
"model": "azure/computer-use-preview-1",
"api_key": "mock-api-key-1",
"api_version": "mock-api-version",
"api_base": "https://mock-endpoint-1.openai.azure.com",
},
"model_info": {"base_model": "computer-use-preview"},
},
{
"model_name": "azure-computer-use-preview",
"litellm_params": {
"model": "azure/computer-use-preview-2",
"api_key": "mock-api-key-2",
"api_version": "mock-api-version-2",
"api_base": "https://mock-endpoint-2.openai.azure.com",
},
"model_info": {"base_model": "computer-use-preview"},
},
],
optional_pre_call_checks=[
"deployment_affinity",
"responses_api_deployment_check",
],
)
model_group = "azure-computer-use-preview"
user_api_key_hash = "test-user-key-1"
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post, patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=lambda seq: seq[0],
):
mock_post.return_value = MockResponse(mock_response_data, 200)
first_response = await router.aresponses(
model=model_group,
input="Hello, how are you?",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
first_model_id = first_response._hidden_params["model_id"]
first_response_id = first_response.id
all_model_ids = router.get_model_ids(model_name=model_group)
other_model_id = next(mid for mid in all_model_ids if mid != first_model_id)
# Force user-key affinity to point to the OTHER deployment
affinity_cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=model_group,
user_key=user_api_key_hash,
)
await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600)
# Even though user-key affinity points elsewhere, previous_response_id should pin
# to the deployment that created the original response.
follow_up = await router.aresponses(
model=model_group,
input="Follow-up question",
truncation="auto",
previous_response_id=first_response_id,
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
assert follow_up._hidden_params["model_id"] == first_model_id
@pytest.mark.asyncio
async def test_async_user_parameter_does_not_trigger_deployment_affinity():
"""
The OpenAI `user` parameter identifies the *end-user* (not the API key), and should
not be used as an affinity key.
"""
mock_response_data = {
"id": "resp_mock-resp-sdk",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [
{
"type": "message",
"id": "msg_sdk",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "SDK Response"}],
}
],
"parallel_tool_calls": True,
"usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10},
"text": {"format": {"type": "text"}},
"error": None,
"previous_response_id": None,
}
router = litellm.Router(
model_list=[
{
"model_name": "azure-sdk-test",
"litellm_params": {
"model": "azure/sdk-1",
"api_key": "mock",
"api_base": "https://mock1.openai.azure.com",
},
"model_info": {"base_model": "sdk-test"},
},
{
"model_name": "azure-sdk-test",
"litellm_params": {
"model": "azure/sdk-2",
"api_key": "mock",
"api_base": "https://mock2.openai.azure.com",
},
"model_info": {"base_model": "sdk-test"},
},
],
optional_pre_call_checks=["deployment_affinity"],
)
model_group = "azure-sdk-test"
user_id = "sdk-user-123"
choice_calls = {"count": 0}
def deterministic_choice(seq):
choice_calls["count"] += 1
if choice_calls["count"] == 1:
return seq[0]
return seq[1] if len(seq) > 1 else seq[0]
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post, patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=deterministic_choice,
):
mock_post.return_value = MockResponse(mock_response_data, 200)
# First call with 'user' parameter (end-user id)
first_response = await router.aresponses(
model=model_group,
input="Hi",
user=user_id,
)
first_model_id = first_response._hidden_params["model_id"]
# Second call with same 'user' parameter should NOT be pinned by affinity
second_response = await router.aresponses(
model=model_group,
input="Follow-up",
user=user_id,
)
assert second_response._hidden_params["model_id"] != first_model_id
@pytest.mark.asyncio
async def test_async_pre_call_hook_uses_model_map_key_scope():
"""
Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id.
"""
cache = AsyncMock()
cache.async_set_cache = AsyncMock()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
)
kwargs = {
"model_info": {"id": "model-id-123"},
"litellm_metadata": {
"user_api_key_hash": "user-key-abc",
"deployment_model_name": "claude-sonnet-4-5@20250929",
},
}
await callback.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None)
expected_cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group="claude-sonnet-4-5@20250929",
user_key="user-key-abc",
)
cache.async_set_cache.assert_called_once_with(
expected_cache_key,
{"model_id": "model-id-123"},
ttl=123,
)
@pytest.mark.asyncio
async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_scope():
"""
When a stable model-map key can be derived from the deployment set, affinity should
be scoped to that key (this helps stickiness across aliases).
This is intentionally tested at the callback level (not via Router), to validate the
cache key selection logic deterministically.
"""
user_key = "user-key-abc"
stable_model_map_key = "claude-sonnet-4-5@20250929"
cache = AsyncMock()
cache.async_get_cache = AsyncMock()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
)
healthy_deployments = [
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"},
"model_info": {"id": "deployment-1"},
},
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"},
"model_info": {"id": "deployment-2"},
},
]
expected_cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=stable_model_map_key,
user_key=user_key,
)
async def get_cache_side_effect(*, key: str):
if key == expected_cache_key:
return {"model_id": "deployment-2"}
return None
cache.async_get_cache.side_effect = get_cache_side_effect
filtered = await callback.async_filter_deployments(
model="some-router-model-group",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={"metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"}},
parent_otel_span=None,
)
assert len(filtered) == 1
assert filtered[0]["model_info"]["id"] == "deployment-2"
@pytest.mark.asyncio
async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unhealthy():
"""
If affinity cache points to a deployment that's no longer healthy, callback should
return all healthy deployments so router can pick an available one.
"""
user_key = "user-key-unhealthy"
stable_model_map_key = "claude-sonnet-4-5@20250929"
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model_id": "stale-deployment"})
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
)
healthy_deployments = [
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"},
"model_info": {"id": "deployment-1"},
},
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"},
"model_info": {"id": "deployment-2"},
},
]
filtered = await callback.async_filter_deployments(
model="some-router-model-group",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={"metadata": {"user_api_key_hash": user_key}},
parent_otel_span=None,
)
assert filtered == healthy_deployments
@pytest.mark.asyncio
async def test_async_user_key_affinity_ttl_expiry_allows_reroute():
"""
After affinity TTL expires, cached pinning should no longer filter deployments.
"""
callback = DeploymentAffinityCheck(
cache=DualCache(),
ttl_seconds=1,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
)
user_key = "ttl-user-key"
stable_model_map_key = "claude-sonnet-4-5@20250929"
healthy_deployments = [
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"},
"model_info": {"id": "deployment-1"},
},
{
"model_name": stable_model_map_key,
"litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"},
"model_info": {"id": "deployment-2"},
},
]
await callback.async_pre_call_deployment_hook(
kwargs={
"model_info": {"id": "deployment-1"},
"metadata": {
"user_api_key_hash": user_key,
"deployment_model_name": stable_model_map_key,
},
},
call_type=None,
)
pinned = await callback.async_filter_deployments(
model="some-router-model-group",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={"metadata": {"user_api_key_hash": user_key}},
parent_otel_span=None,
)
assert len(pinned) == 1
assert pinned[0]["model_info"]["id"] == "deployment-1"
await asyncio.sleep(1.2)
after_ttl_expiry = await callback.async_filter_deployments(
model="some-router-model-group",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={"metadata": {"user_api_key_hash": user_key}},
parent_otel_span=None,
)
assert after_ttl_expiry == healthy_deployments
def test_cache_key_does_not_double_hash_user_api_key_hash():
"""
Proxy typically provides `metadata.user_api_key_hash` as a SHA-256 hex string.
The affinity cache key should not hash it again.
"""
user_api_key_hash = "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1"
key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group="any-model-group",
user_key=user_api_key_hash,
)
assert key.endswith(user_api_key_hash)