diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 6c5c45dc90c..ed55c6de23c 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 140dfd4faf8..56f542361ce 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 b77e3c9c299..eb68414b0db 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 00000000000..6d821cd5f29 --- /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 8ea7a207535..d09d2ca01f5 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 00000000000..3db2192ff06 --- /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]