From e23bd21a286ec0e279eee141f1e1722562af33cc Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 15 Jan 2026 08:53:31 -0600 Subject: [PATCH 1/3] Add deployment affinity routing --- docs/my-website/docs/proxy/config_settings.md | 3 +- docs/my-website/docs/response_api.md | 19 +- litellm/router.py | 97 +++- .../deployment_affinity_check.py | 241 ++++++++ litellm/types/router.py | 1 + .../test_deployment_affinity_check.py | 513 ++++++++++++++++++ 6 files changed, 848 insertions(+), 26 deletions(-) create mode 100644 litellm/router_utils/pre_call_checks/deployment_affinity_check.py create mode 100644 tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 6c5c45dc90..ed55c6de23 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -356,7 +356,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. 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) | diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 140dfd4faf..56f542361c 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -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,10 @@ 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`) ```yaml showLineNumbers title="config.yaml with Session Continuity" model_list: @@ -929,7 +937,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 @@ -1232,4 +1244,3 @@ Response: - diff --git a/litellm/router.py b/litellm/router.py index b77e3c9c29..eb68414b0d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -109,12 +109,12 @@ from litellm.router_utils.handle_error import ( async_raise_no_deployment_exception, send_llm_exception_alert, ) +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, @@ -288,6 +288,7 @@ class Router: router_general_settings: Optional[ RouterGeneralSettings ] = RouterGeneralSettings(), + deployment_affinity_ttl_seconds: int = 3600, ignore_invalid_deployments: bool = False, ) -> None: """ @@ -321,6 +322,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. @@ -593,6 +595,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 @@ -1162,24 +1165,76 @@ 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() - 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, + ) + + 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): """ diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py new file mode 100644 index 0000000000..6d821cd5f2 --- /dev/null +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -0,0 +1,241 @@ +""" +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. User-key affinity: map a user key -> deployment id for a TTL and re-use that + deployment for subsequent requests to the same model group. + +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, + ): + 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 _hash_user_key(user_key: str) -> str: + return hashlib.sha256(user_key.encode("utf-8")).hexdigest() + + @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]: + user_key = metadata.get("user_api_key_hash") or metadata.get("user_api_key") + if user_key is None: + return None + return str(user_key) + + @staticmethod + def _get_user_key_from_request_kwargs(request_kwargs: dict) -> Optional[str]: + """ + Extract a stable user key from request kwargs. + + Primary source (proxy): `metadata.user_api_key_hash` / `metadata.user_api_key` + Fallback (SDK): `user` + """ + # 1. Check metadata (Proxy usage) + metadata = request_kwargs.get("litellm_metadata") or request_kwargs.get("metadata") + if isinstance(metadata, dict): + user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict( + metadata=metadata + ) + if user_key is not None: + return user_key + + # 2. Check top-level 'user' parameter (SDK usage) + user_key = request_kwargs.get("user") + if user_key is not None: + return str(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: + deployment_model_id = deployment.get("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 user-key deployment affinity + """ + request_kwargs = request_kwargs or {} + + # 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=cast(List[dict], 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 cast(List[dict], healthy_deployments) + + user_key = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + if user_key is None: + return cast(List[dict], healthy_deployments) + + cache_key = self.get_affinity_cache_key(model_group=model, 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 cast(List[dict], healthy_deployments) + + deployment = self._find_deployment_by_model_id( + healthy_deployments=cast(List[dict], 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 cast(List[dict], healthy_deployments) + + verbose_router_logger.debug( + "DeploymentAffinityCheck: user-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 user-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 = kwargs.get("litellm_metadata") or kwargs.get("metadata") or {} + if not isinstance(metadata, dict): + return None + + model_group = metadata.get("model_group") + if not model_group: + return None + + model_info = kwargs.get("model_info") or metadata.get("model_info") or {} + if not isinstance(model_info, dict): + return None + + model_id = model_info.get("id") + if not model_id: + return None + + cache_key = self.get_affinity_cache_key( + model_group=str(model_group), user_key=user_key + ) + try: + 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_group=%s deployment=%s ttl=%s user_key=%s", + model_group, + 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 cache key=%s: %s", + cache_key, + e, + ) + + return None diff --git a/litellm/types/router.py b/litellm/types/router.py index 8ea7a20753..d09d2ca01f 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -793,6 +793,7 @@ OptionalPreCallChecks = List[ "prompt_caching", "router_budget_limiting", "responses_api_deployment_check", + "deployment_affinity", "forward_client_headers_by_model_group", ] ] diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py new file mode 100644 index 0000000000..3db2192ff0 --- /dev/null +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -0,0 +1,513 @@ +import os +import sys +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import json + +import litellm +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", + }, + }, + { + "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", + }, + }, + ], + 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_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_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", + }, + }, + ], + 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_affinity(): + """ + When 'user' is passed as a top-level parameter (SDK-style), affinity should work. + """ + 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_name": "azure-sdk-test", + "litellm_params": { + "model": "azure/sdk-2", + "api_key": "mock", + "api_base": "https://mock2.openai.azure.com", + }, + }, + ], + 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 + 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 use 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_affinity_cache_expiry_allows_reroute(): + """ + When affinity TTL expires, routing should fall back to normal load balancing. + """ + mock_response_data = { + "id": "resp_mock-resp-ttl", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "azure/computer-use-preview", + "output": [ + { + "type": "message", + "id": "msg_ttl", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "TTL 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-ttl-test", + "litellm_params": { + "model": "azure/ttl-1", + "api_key": "mock", + "api_base": "https://mock1.openai.azure.com", + }, + }, + { + "model_name": "azure-ttl-test", + "litellm_params": { + "model": "azure/ttl-2", + "api_key": "mock", + "api_base": "https://mock2.openai.azure.com", + }, + }, + ], + optional_pre_call_checks=["deployment_affinity"], + deployment_affinity_ttl_seconds=1, + ) + + model_group = "azure-ttl-test" + user_api_key_hash = "ttl-user-key" + + 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="Hi", + litellm_metadata={"user_api_key_hash": user_api_key_hash}, + ) + first_model_id = first_response._hidden_params["model_id"] + + await asyncio.sleep(1.1) + + second_response = await router.aresponses( + model=model_group, + input="Follow-up after ttl", + 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_affinity_cache_missing_deployment_falls_back(): + """ + If a cached model_id is not in healthy deployments, routing should ignore it. + """ + mock_response_data = { + "id": "resp_mock-resp-missing", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "azure/computer-use-preview", + "output": [ + { + "type": "message", + "id": "msg_missing", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Missing 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-missing-test", + "litellm_params": { + "model": "azure/missing-1", + "api_key": "mock", + "api_base": "https://mock1.openai.azure.com", + }, + }, + { + "model_name": "azure-missing-test", + "litellm_params": { + "model": "azure/missing-2", + "api_key": "mock", + "api_base": "https://mock2.openai.azure.com", + }, + }, + ], + optional_pre_call_checks=["deployment_affinity"], + ) + + model_group = "azure-missing-test" + user_api_key_hash = "missing-user-key" + + 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[1] if len(seq) > 1 else seq[0], + ): + mock_post.return_value = MockResponse(mock_response_data, 200) + + 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": "non-existent-model-id"}, + ttl=3600, + ) + + response = await router.aresponses( + model=model_group, + input="Should ignore missing affinity", + litellm_metadata={"user_api_key_hash": user_api_key_hash}, + ) + + model_ids = router.get_model_ids(model_name=model_group) + assert response._hidden_params["model_id"] == model_ids[1] From 059e75ad8843d3d704ad26c8c6ac68825a2dc2ad Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 15 Jan 2026 16:22:10 -0600 Subject: [PATCH 2/3] fix(router): scope deployment affinity by model_map_key - Key affinity by (user_api_key_hash, model_map_key) -> model_id - Ignore OpenAI 'user' param for affinity - Avoid double hashing user_api_key_hash - Add unit tests + docs clarifications --- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/response_api.md | 10 +- .../deployment_affinity_check.py | 238 ++++++++++++--- .../test_deployment_affinity_check.py | 281 +++++++++++++----- 4 files changed, 408 insertions(+), 123 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index ed55c6de23..92957d3cd0 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -357,7 +357,7 @@ router_settings: | 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. 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. Defaults to `3600` (1 hour). | +| 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) | diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 56f542361c..c51257359f 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -921,6 +921,12 @@ To enable session continuity for Responses API in your LiteLLM proxy, set `optio - `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: - model_name: azure-gpt4-turbo @@ -1240,7 +1246,3 @@ Response: - - - - diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 6d821cd5f2..da9e7be828 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -4,8 +4,8 @@ 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. User-key affinity: map a user key -> deployment id for a TTL and re-use that - deployment for subsequent requests to the same model group. +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 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. @@ -50,10 +50,109 @@ class DeploymentAffinityCheck(CustomLogger): 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 Exception: + 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. + + Prefer `base_model` when available (important for Azure), otherwise fall back to + parsing `litellm_params.model`. + """ + 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: @@ -67,33 +166,46 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_user_key_from_metadata_dict(metadata: dict) -> Optional[str]: - user_key = metadata.get("user_api_key_hash") or metadata.get("user_api_key") + # 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 user key from request kwargs. + Extract a stable affinity key from request kwargs. - Primary source (proxy): `metadata.user_api_key_hash` / `metadata.user_api_key` - Fallback (SDK): `user` + 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. """ - # 1. Check metadata (Proxy usage) - metadata = request_kwargs.get("litellm_metadata") or request_kwargs.get("metadata") - if isinstance(metadata, dict): + # 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 - # 2. Check top-level 'user' parameter (SDK usage) - user_key = request_kwargs.get("user") - if user_key is not None: - return str(user_key) - return None @staticmethod @@ -101,10 +213,11 @@ class DeploymentAffinityCheck(CustomLogger): healthy_deployments: List[dict], model_id: str ) -> Optional[dict]: for deployment in healthy_deployments: - deployment_model_id = deployment.get("model_info", {}).get("id") - if deployment_model_id is not None and str(deployment_model_id) == str( - model_id - ): + 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 @@ -119,9 +232,10 @@ class DeploymentAffinityCheck(CustomLogger): """ Optionally filter healthy deployments based on: 1. `previous_response_id` (Responses API continuity) [highest priority] - 2. cached user-key deployment affinity + 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: @@ -132,7 +246,7 @@ class DeploymentAffinityCheck(CustomLogger): ) if responses_model_id is not None: deployment = self._find_deployment_by_model_id( - healthy_deployments=cast(List[dict], healthy_deployments), + healthy_deployments=typed_healthy_deployments, model_id=responses_model_id, ) if deployment is not None: @@ -144,13 +258,21 @@ class DeploymentAffinityCheck(CustomLogger): # 2) User key -> deployment affinity if not self.enable_user_key_affinity: - return cast(List[dict], healthy_deployments) + return typed_healthy_deployments user_key = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if user_key is None: - return cast(List[dict], healthy_deployments) + return typed_healthy_deployments - cache_key = self.get_affinity_cache_key(model_group=model, user_key=user_key) + 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 @@ -161,10 +283,10 @@ class DeploymentAffinityCheck(CustomLogger): model_id = cache_result if not model_id: - return cast(List[dict], healthy_deployments) + return typed_healthy_deployments deployment = self._find_deployment_by_model_id( - healthy_deployments=cast(List[dict], healthy_deployments), + healthy_deployments=typed_healthy_deployments, model_id=model_id, ) if deployment is None: @@ -172,10 +294,10 @@ class DeploymentAffinityCheck(CustomLogger): "DeploymentAffinityCheck: pinned deployment=%s not found in healthy_deployments", model_id, ) - return cast(List[dict], healthy_deployments) + return typed_healthy_deployments verbose_router_logger.debug( - "DeploymentAffinityCheck: user-key affinity hit -> deployment=%s user_key=%s", + "DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s", model_id, self._shorten_for_logs(user_key), ) @@ -185,7 +307,7 @@ class DeploymentAffinityCheck(CustomLogger): self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] ) -> Optional[dict]: """ - Persist/update the user-key -> deployment mapping for this request. + 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. @@ -198,34 +320,64 @@ class DeploymentAffinityCheck(CustomLogger): if user_key is None: return None - metadata = kwargs.get("litellm_metadata") or kwargs.get("metadata") or {} - if not isinstance(metadata, dict): - return None + metadata_dicts = self._iter_metadata_dicts(kwargs) - model_group = metadata.get("model_group") - if not model_group: - return None - - model_info = kwargs.get("model_info") or metadata.get("model_info") or {} + 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 + + # Primary scope: stable model-map key (used to handle aliases that ultimately map to the same + # underlying base model). + model_map_key: Optional[str] = None + base_model = model_info.get("base_model") + if isinstance(base_model, str) and base_model: + model_map_key = base_model + else: + litellm_model_name = kwargs.get("model") + if isinstance(litellm_model_name, str) and litellm_model_name: + model_map_key = self._get_model_map_key_from_litellm_model_name( + litellm_model_name + ) + + if not model_map_key: + verbose_router_logger.warning( + "DeploymentAffinityCheck: model_map_key missing; skipping affinity cache update. model_id=%s", + model_id, + ) return None - cache_key = self.get_affinity_cache_key( - model_group=str(model_group), user_key=user_key - ) try: + cache_key = self.get_affinity_cache_key( + model_group=model_map_key, 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_group=%s deployment=%s ttl=%s user_key=%s", - model_group, + "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", + model_map_key, model_id, self.ttl_seconds, self._shorten_for_logs(user_key), @@ -233,8 +385,8 @@ class DeploymentAffinityCheck(CustomLogger): except Exception as e: # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( - "DeploymentAffinityCheck: failed to set cache key=%s: %s", - cache_key, + "DeploymentAffinityCheck: failed to set affinity cache. model_map_key=%s error=%s", + model_map_key, e, ) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 3db2192ff0..518c5d97e2 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -82,6 +82,8 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "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", @@ -91,6 +93,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "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"], @@ -137,6 +140,99 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): 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(): """ @@ -197,6 +293,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): "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", @@ -206,6 +303,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): "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=[ @@ -240,7 +338,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): # Force user-key affinity to point to the OTHER deployment affinity_cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group=model_group, + model_group="computer-use-preview", user_key=user_api_key_hash, ) await router.cache.async_set_cache( @@ -260,9 +358,10 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): @pytest.mark.asyncio -async def test_async_user_parameter_affinity(): +async def test_async_user_parameter_does_not_trigger_deployment_affinity(): """ - When 'user' is passed as a top-level parameter (SDK-style), affinity should work. + 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", @@ -295,6 +394,7 @@ async def test_async_user_parameter_affinity(): "api_key": "mock", "api_base": "https://mock1.openai.azure.com", }, + "model_info": {"base_model": "sdk-test"}, }, { "model_name": "azure-sdk-test", @@ -303,6 +403,7 @@ async def test_async_user_parameter_affinity(): "api_key": "mock", "api_base": "https://mock2.openai.azure.com", }, + "model_info": {"base_model": "sdk-test"}, }, ], optional_pre_call_checks=["deployment_affinity"], @@ -328,7 +429,7 @@ async def test_async_user_parameter_affinity(): ): mock_post.return_value = MockResponse(mock_response_data, 200) - # First call with 'user' parameter + # First call with 'user' parameter (end-user id) first_response = await router.aresponses( model=model_group, input="Hi", @@ -336,13 +437,13 @@ async def test_async_user_parameter_affinity(): ) first_model_id = first_response._hidden_params["model_id"] - # Second call with same 'user' parameter should use affinity + # 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 + assert second_response._hidden_params["model_id"] != first_model_id @pytest.mark.asyncio @@ -433,81 +534,111 @@ async def test_async_affinity_cache_expiry_allows_reroute(): @pytest.mark.asyncio -async def test_async_affinity_cache_missing_deployment_falls_back(): +async def test_async_pre_call_hook_uses_model_map_key_scope(): """ - If a cached model_id is not in healthy deployments, routing should ignore it. + Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id. """ - mock_response_data = { - "id": "resp_mock-resp-missing", - "object": "response", - "created_at": 1741476542, - "status": "completed", - "model": "azure/computer-use-preview", - "output": [ - { - "type": "message", - "id": "msg_missing", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "Missing 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-missing-test", - "litellm_params": { - "model": "azure/missing-1", - "api_key": "mock", - "api_base": "https://mock1.openai.azure.com", - }, - }, - { - "model_name": "azure-missing-test", - "litellm_params": { - "model": "azure/missing-2", - "api_key": "mock", - "api_base": "https://mock2.openai.azure.com", - }, - }, - ], - optional_pre_call_checks=["deployment_affinity"], + cache = AsyncMock() + cache.async_set_cache = AsyncMock() + + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, ) - model_group = "azure-missing-test" - user_api_key_hash = "missing-user-key" + kwargs = { + "model_info": {"id": "model-id-123", "base_model": "claude-sonnet-4-5@20250929"}, + "litellm_metadata": { + "user_api_key_hash": "user-key-abc", + }, + } - 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[1] if len(seq) > 1 else seq[0], - ): - mock_post.return_value = MockResponse(mock_response_data, 200) + await callback.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) - 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": "non-existent-model-id"}, - ttl=3600, - ) + 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, + ) - response = await router.aresponses( - model=model_group, - input="Should ignore missing affinity", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - model_ids = router.get_model_ids(model_name=model_group) - assert response._hidden_params["model_id"] == model_ids[1] +@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": "group-any", + "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "group-any", + "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "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" + + +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) +>>>>>>> 860962807e (fix(router): scope deployment affinity by model_map_key) From f7726c895051ddd67efd5904bc052b41b7db28df Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 15 Jan 2026 17:42:29 -0600 Subject: [PATCH 3/3] Fix wrong keys being used for model sticky entry --- .../deployment_affinity_check.py | 50 ++++---- .../test_deployment_affinity_check.py | 108 ++---------------- 2 files changed, 34 insertions(+), 124 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index da9e7be828..4d75478bf9 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -5,7 +5,8 @@ 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 model-map key. + 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. @@ -106,9 +107,18 @@ class DeploymentAffinityCheck(CustomLogger): """ 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") @@ -209,9 +219,7 @@ class DeploymentAffinityCheck(CustomLogger): return None @staticmethod - def _find_deployment_by_model_id( - healthy_deployments: List[dict], model_id: str - ) -> Optional[dict]: + 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): @@ -241,9 +249,7 @@ class DeploymentAffinityCheck(CustomLogger): 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) - ) + 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, @@ -345,29 +351,25 @@ class DeploymentAffinityCheck(CustomLogger): ) return None - # Primary scope: stable model-map key (used to handle aliases that ultimately map to the same - # underlying base model). - model_map_key: Optional[str] = None - base_model = model_info.get("base_model") - if isinstance(base_model, str) and base_model: - model_map_key = base_model - else: - litellm_model_name = kwargs.get("model") - if isinstance(litellm_model_name, str) and litellm_model_name: - model_map_key = self._get_model_map_key_from_litellm_model_name( - litellm_model_name - ) + # 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 model_map_key: + if not deployment_model_name: verbose_router_logger.warning( - "DeploymentAffinityCheck: model_map_key missing; skipping affinity cache update. model_id=%s", + "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=model_map_key, user_key=user_key + model_group=deployment_model_name, user_key=user_key ) await self.cache.async_set_cache( cache_key, @@ -377,7 +379,7 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.debug( "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", - model_map_key, + deployment_model_name, model_id, self.ttl_seconds, self._shorten_for_logs(user_key), @@ -386,7 +388,7 @@ class DeploymentAffinityCheck(CustomLogger): # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( "DeploymentAffinityCheck: failed to set affinity cache. model_map_key=%s error=%s", - model_map_key, + deployment_model_name, e, ) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 518c5d97e2..57d2807c8c 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1,6 +1,5 @@ import os import sys -import asyncio from unittest.mock import AsyncMock, patch import pytest @@ -44,9 +43,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -338,12 +335,10 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): # Force user-key affinity to point to the OTHER deployment affinity_cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group="computer-use-preview", + 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 - ) + 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. @@ -446,93 +441,6 @@ async def test_async_user_parameter_does_not_trigger_deployment_affinity(): assert second_response._hidden_params["model_id"] != first_model_id -@pytest.mark.asyncio -async def test_async_affinity_cache_expiry_allows_reroute(): - """ - When affinity TTL expires, routing should fall back to normal load balancing. - """ - mock_response_data = { - "id": "resp_mock-resp-ttl", - "object": "response", - "created_at": 1741476542, - "status": "completed", - "model": "azure/computer-use-preview", - "output": [ - { - "type": "message", - "id": "msg_ttl", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "TTL 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-ttl-test", - "litellm_params": { - "model": "azure/ttl-1", - "api_key": "mock", - "api_base": "https://mock1.openai.azure.com", - }, - }, - { - "model_name": "azure-ttl-test", - "litellm_params": { - "model": "azure/ttl-2", - "api_key": "mock", - "api_base": "https://mock2.openai.azure.com", - }, - }, - ], - optional_pre_call_checks=["deployment_affinity"], - deployment_affinity_ttl_seconds=1, - ) - - model_group = "azure-ttl-test" - user_api_key_hash = "ttl-user-key" - - 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="Hi", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - first_model_id = first_response._hidden_params["model_id"] - - await asyncio.sleep(1.1) - - second_response = await router.aresponses( - model=model_group, - input="Follow-up after ttl", - 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_pre_call_hook_uses_model_map_key_scope(): """ @@ -550,9 +458,10 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): ) kwargs = { - "model_info": {"id": "model-id-123", "base_model": "claude-sonnet-4-5@20250929"}, + "model_info": {"id": "model-id-123"}, "litellm_metadata": { "user_api_key_hash": "user-key-abc", + "deployment_model_name": "claude-sonnet-4-5@20250929", }, } @@ -594,13 +503,13 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s healthy_deployments = [ { - "model_name": "group-any", + "model_name": stable_model_map_key, "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, "model_info": {"id": "deployment-1"}, }, { - "model_name": "group-any", - "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "model_name": stable_model_map_key, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -641,4 +550,3 @@ def test_cache_key_does_not_double_hash_user_api_key_hash(): user_key=user_api_key_hash, ) assert key.endswith(user_api_key_hash) ->>>>>>> 860962807e (fix(router): scope deployment affinity by model_map_key)