diff --git a/litellm/router.py b/litellm/router.py index 6e773a06c7f..232df718294 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -701,6 +701,8 @@ class Router: self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + self._apply_complexity_router_session_affinity() + if self.alerting_config is not None: self._initialize_alerting() @@ -7571,6 +7573,79 @@ class Router: f"Complexity-router deployment {deployment.model_name} already exists. Please use a different model name." ) self.complexity_routers[deployment.model_name] = complexity_router + if hasattr(self, "model_group_affinity_config"): + self._apply_complexity_router_session_affinity() + + def _apply_complexity_router_session_affinity(self) -> None: + """Enable deployment-level session affinity for the tier target model groups of + complexity routers that opted in via + session_stickiness.enable_deployment_session_affinity. + + The complexity router's pre-routing hook rewrites the request's model to a tier + target group before DeploymentAffinityCheck.async_filter_deployments runs, so + affinity must be configured on the tier groups (not the auto-router group) for + the same physical deployment to serve a whole session. + + Runs after __init__ has assigned model_group_affinity_config (guarded at the + init_complexity_router_deployment callsite because set_model_list runs earlier + in __init__), and again on runtime hot-adds. Explicit per-group entries make + _get_effective_flags ignore the callback's global flags for that group, so new + entries are seeded with the currently-enabled global flags; user-provided + entries are extended, never replaced. + """ + opted_in_groups = { + group + for complexity_router in self.complexity_routers.values() + if complexity_router.config.session_stickiness.enable_deployment_session_affinity + for group in (*complexity_router.config.tiers.values(), complexity_router.config.default_model) + if group + } + if not opted_in_groups: + return + + existing_affinity_callback = next( + (cb for cb in (self.optional_callbacks or []) if isinstance(cb, DeploymentAffinityCheck)), + None, + ) + global_flag_seed = tuple( + flag + for flag, enabled in ( + ("deployment_affinity", getattr(existing_affinity_callback, "enable_user_key_affinity", False)), + ( + "responses_api_deployment_check", + getattr(existing_affinity_callback, "enable_responses_api_affinity", False), + ), + ("session_affinity", getattr(existing_affinity_callback, "enable_session_id_affinity", False)), + ) + if enabled + ) + + current_config = self.model_group_affinity_config or {} + merged_config = { + **current_config, + **{ + group: list(dict.fromkeys([*current_config.get(group, global_flag_seed), "session_affinity"])) + for group in opted_in_groups + }, + } + self.model_group_affinity_config = merged_config + + if existing_affinity_callback is not None: + existing_affinity_callback.model_group_affinity_config = merged_config + return + + if self.optional_callbacks is None: + self.optional_callbacks = [] + affinity_callback = DeploymentAffinityCheck( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=merged_config, + ) + self.optional_callbacks.append(affinity_callback) + litellm.logging_callback_manager.add_litellm_callback(affinity_callback) def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 11719b8a18f..4ebfbb2ebaf 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -660,6 +660,64 @@ class ComplexityRouter(CustomLogger): return user_message, system_prompt + SESSION_TIER_CACHE_KEY_PREFIX = "complexity_router:session_tier:v1" + + @classmethod + def get_session_tier_cache_key(cls, model_name: str, session_id: str) -> str: + return f"{cls.SESSION_TIER_CACHE_KEY_PREFIX}:{model_name}:{session_id}" + + @staticmethod + def _get_session_id(request_kwargs: Dict) -> "str | None": + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + return DeploymentAffinityCheck._get_session_id_from_request_kwargs(request_kwargs) + + async def _get_pinned_tier(self, session_id: str) -> "ComplexityTier | None": + """Read the session's pinned tier from the router cache. + + An unparseable stored value (e.g. from an older/newer version) is treated as + no pin rather than failing the request. + """ + cache = getattr(self.litellm_router_instance, "cache", None) + if cache is None: + return None + raw = await cache.async_get_cache(self.get_session_tier_cache_key(self.model_name, session_id)) + if raw is None: + return None + try: + return ComplexityTier(str(raw)) + except ValueError: + return None + + async def _pin_session_tier(self, session_id: str, tier: ComplexityTier) -> None: + cache = getattr(self.litellm_router_instance, "cache", None) + if cache is None: + return + ttl = self.config.session_stickiness.ttl_seconds + await cache.async_set_cache( + self.get_session_tier_cache_key(self.model_name, session_id), + tier.value, + ttl=ttl, + ) + verbose_router_logger.debug(f"ComplexityRouter: pinned session {session_id} to tier={tier.value} ttl={ttl}") + + async def _decide_tier( + self, + user_message: str, + system_prompt: "str | None", + request_kwargs: Dict, + ) -> Tuple[ComplexityTier, str, str]: + """Run keyword overrides then the classifier; returns (tier, cause, extra_log).""" + override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs) + if override_tier is not None: + cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match" + return override_tier, cause, "" + + tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs) + return tier, "complexity_scorer", f"score={score:.3f}, signals={signals}, " + async def async_pre_routing_hook( self, model: str, @@ -675,6 +733,10 @@ class ComplexityRouter(CustomLogger): Supports chat completions (messages), Responses API (input), and other formats via the guardrail translation handler dispatch. + With session_stickiness enabled, the decision is pinned per session_id so + multi-turn conversations keep hitting the same tier (preserving provider-side + prompt/KV cache) and skip classifier/embedding work where possible. + Args: model: The original model name requested. request_kwargs: The request kwargs. @@ -695,6 +757,7 @@ class ComplexityRouter(CustomLogger): # Determine whether the original request used messages directly has_original_messages = messages is not None and len(messages) > 0 + response_messages = messages if has_original_messages else None user_message, system_prompt = self._extract_user_message_and_system_prompt(resolved_messages) @@ -702,31 +765,36 @@ class ComplexityRouter(CustomLogger): verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") return PreRoutingHookResponse( model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), - messages=messages if has_original_messages else None, + messages=response_messages, ) - override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs) - if override_tier is not None: - routed_model = self.get_model_for_tier(override_tier) - cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match" + stickiness = self.config.session_stickiness + session_id = self._get_session_id(request_kwargs) if stickiness.mode != "none" else None + pinned_tier = await self._get_pinned_tier(session_id) if session_id is not None else None + + if pinned_tier is not None and (stickiness.mode == "sticky" or pinned_tier is TIER_SEVERITY_ORDER[-1]): + routed_model = self.get_model_for_tier(pinned_tier) verbose_router_logger.info( - f"ComplexityRouter: routing decision cause={cause}, " - f"tier={override_tier.value}, routed_model={routed_model}" - ) - return PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, + f"ComplexityRouter: routing decision cause=session_pinned, " + f"tier={pinned_tier.value}, routed_model={routed_model}" ) + return PreRoutingHookResponse(model=routed_model, messages=response_messages) + + tier, cause, extra_log = await self._decide_tier(user_message, system_prompt, request_kwargs) + + if session_id is not None and pinned_tier is not None: + if TIER_SEVERITY_ORDER.index(tier) > TIER_SEVERITY_ORDER.index(pinned_tier): + await self._pin_session_tier(session_id, tier) + cause, extra_log = "session_escalation", "" + else: + tier, cause, extra_log = pinned_tier, "session_pinned", "" + elif session_id is not None: + await self._pin_session_tier(session_id, tier) - tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs) routed_model = self.get_model_for_tier(tier) - verbose_router_logger.info( - f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, " - f"score={score:.3f}, signals={signals}, routed_model={routed_model}" + f"ComplexityRouter: routing decision cause={cause}, tier={tier.value}, " + f"{extra_log}routed_model={routed_model}" ) - return PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, - ) + return PreRoutingHookResponse(model=routed_model, messages=response_messages) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 125de6f7489..6aff2c369e1 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -241,6 +241,36 @@ class ClassifierLLMConfig(BaseModel): ) +class SessionStickinessConfig(BaseModel): + """Session-sticky tier routing to preserve provider-side prompt/KV cache across turns. + + Modes: + - "none": classify every turn (default; today's behavior). + - "sticky": the first decision pins the session; later turns skip keyword overrides + and the classifier entirely and reuse the pinned tier. + - "sticky_with_escalation": the pinned tier can only ratchet up in severity, never + down. Keyword overrides and the classifier still run (unless already at the top + tier) so a complexity jump mid-session escalates; downgrades are suppressed. + """ + + mode: Literal["none", "sticky", "sticky_with_escalation"] = Field( + default="none", + description="Session stickiness mode for tier decisions", + ) + ttl_seconds: int = Field( + default=3600, + gt=0, + description="How long a session's pinned tier lives after its last write", + ) + enable_deployment_session_affinity: bool = Field( + default=False, + description=( + "Also enable deployment-level session_affinity for this router's tier target " + "model groups, so the same physical deployment serves the whole session" + ), + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -333,6 +363,12 @@ class ComplexityRouterConfig(BaseModel): description="Minimum cosine similarity for a semantic keyword match", ) + # Session-sticky tier routing + session_stickiness: SessionStickinessConfig = Field( + default_factory=SessionStickinessConfig, + description="Pin tier decisions per session to preserve provider prompt/KV cache", + ) + model_config = ConfigDict(extra="allow") # Allow additional fields @model_validator(mode="after") @@ -351,6 +387,16 @@ class ComplexityRouterConfig(BaseModel): raise ValueError("keyword_tier_rules must be non-empty when semantic_keyword_matching is enabled") return self + @model_validator(mode="after") + def _validate_session_stickiness(self) -> "ComplexityRouterConfig": + stickiness = self.session_stickiness + if stickiness.enable_deployment_session_affinity and stickiness.mode == "none": + raise ValueError( + "enable_deployment_session_affinity requires session_stickiness.mode to be " + "'sticky' or 'sticky_with_escalation'" + ) + return self + # Combined default config DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index f47c19b2baa..a7e76aa7897 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2106,3 +2106,473 @@ class TestRoutingDecisionCauseLogging: assert "score=" in router_log_capture.text assert "cause=literal_keyword_match" not in router_log_capture.text assert "cause=semantic_keyword_match" not in router_log_capture.text + + +class TestSessionStickinessConfigValidation: + """SessionStickinessConfig defaults and cross-field validation.""" + + def test_defaults(self): + config = ComplexityRouterConfig() + assert config.session_stickiness.mode == "none" + assert config.session_stickiness.ttl_seconds == 3600 + assert config.session_stickiness.enable_deployment_session_affinity is False + + def test_deployment_affinity_requires_stickiness(self): + with pytest.raises(ValidationError, match="enable_deployment_session_affinity"): + ComplexityRouterConfig( + session_stickiness={"mode": "none", "enable_deployment_session_affinity": True} + ) + + def test_deployment_affinity_allowed_with_sticky_mode(self): + config = ComplexityRouterConfig( + session_stickiness={"mode": "sticky", "enable_deployment_session_affinity": True} + ) + assert config.session_stickiness.enable_deployment_session_affinity is True + + def test_non_positive_ttl_rejected(self): + with pytest.raises(ValidationError): + ComplexityRouterConfig(session_stickiness={"mode": "sticky", "ttl_seconds": 0}) + + def test_unknown_mode_rejected(self): + with pytest.raises(ValidationError): + ComplexityRouterConfig(session_stickiness={"mode": "always"}) + + +def _sticky_router(mock_router_instance, basic_config: Dict, mode: str) -> ComplexityRouter: + from litellm.caching.caching import DualCache + + mock_router_instance.cache = DualCache() + return ComplexityRouter( + model_name="test-sticky-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "session_stickiness": {"mode": mode}}, + ) + + +def _session_kwargs(session_id: str, metadata_key: str = "metadata") -> Dict: + return {metadata_key: {"session_id": session_id}} + + +SIMPLE_MESSAGES = [{"role": "user", "content": "What is the capital of France?"}] +COMPLEX_MESSAGES = [ + { + "role": "user", + "content": ( + "Refactor this distributed microservice architecture to optimize database " + "query throughput, implement async request handling in python, and debug " + "the authentication api endpoint errors step by step" + ), + } +] + + +class TestSessionStickyRouting: + """mode=sticky pins the first decision and skips all classification work afterwards.""" + + @pytest.mark.asyncio + async def test_mode_none_never_touches_cache(self, mock_router_instance, basic_config): + mock_router_instance.cache = MagicMock() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-none"), + messages=SIMPLE_MESSAGES, + ) + assert result is not None + mock_router_instance.cache.async_get_cache.assert_not_called() + mock_router_instance.cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_no_session_id_falls_back_to_per_turn(self, mock_router_instance, basic_config): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=SIMPLE_MESSAGES, + ) + assert result is not None + assert result.model == "gpt-4o-mini" + cached = await mock_router_instance.cache.async_get_cache( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", "") + ) + assert cached is None + + @pytest.mark.asyncio + async def test_first_turn_classifies_and_pins_with_ttl(self, mock_router_instance, basic_config): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + set_cache_mock = AsyncMock() + with patch.object(mock_router_instance.cache, "async_set_cache", set_cache_mock): + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-1"), + messages=SIMPLE_MESSAGES, + ) + assert result is not None + assert result.model == "gpt-4o-mini" + set_cache_mock.assert_awaited_once_with( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", "sess-1"), + ComplexityTier.SIMPLE.value, + ttl=3600, + ) + + @pytest.mark.asyncio + async def test_second_turn_skips_classifier_and_overrides(self, mock_router_instance, basic_config): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-2"), + messages=SIMPLE_MESSAGES, + ) + with ( + patch.object(router, "aclassify", new_callable=AsyncMock) as classify_mock, + patch.object( + router, "_resolve_keyword_tier_override", new_callable=AsyncMock + ) as override_mock, + ): + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-2"), + messages=COMPLEX_MESSAGES, + ) + assert result is not None + assert result.model == "gpt-4o-mini" + classify_mock.assert_not_awaited() + override_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_keyword_override_on_first_turn_shapes_the_pin( + self, mock_router_instance, basic_config + ): + from litellm.caching.caching import DualCache + + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-sticky-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_stickiness": {"mode": "sticky"}, + "keyword_tier_rules": [{"keywords": ["wire transfer"], "tier": "REASONING"}], + }, + ) + first = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-3"), + messages=[{"role": "user", "content": "please process this wire transfer"}], + ) + assert first is not None + assert first.model == "o1-preview" + second = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-3"), + messages=SIMPLE_MESSAGES, + ) + assert second is not None + assert second.model == "o1-preview" + + @pytest.mark.asyncio + async def test_invalid_cached_value_falls_through_and_overwrites( + self, mock_router_instance, basic_config + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + key = ComplexityRouter.get_session_tier_cache_key("test-sticky-router", "sess-4") + await mock_router_instance.cache.async_set_cache(key, "BOGUS") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-4"), + messages=SIMPLE_MESSAGES, + ) + assert result is not None + assert result.model == "gpt-4o-mini" + assert await mock_router_instance.cache.async_get_cache(key) == ComplexityTier.SIMPLE.value + + @pytest.mark.asyncio + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) + async def test_session_id_read_from_both_metadata_dicts( + self, mock_router_instance, basic_config, metadata_key + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + session_id = f"sess-{metadata_key}" + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs(session_id, metadata_key), + messages=SIMPLE_MESSAGES, + ) + cached = await mock_router_instance.cache.async_get_cache( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", session_id) + ) + assert cached == ComplexityTier.SIMPLE.value + + +class TestSessionEscalation: + """mode=sticky_with_escalation ratchets the pinned tier up, never down.""" + + async def _seed_pin(self, mock_router_instance, tier: ComplexityTier, session_id: str) -> None: + await mock_router_instance.cache.async_set_cache( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", session_id), + tier.value, + ) + + @pytest.mark.asyncio + async def test_classified_below_pin_stays_pinned_without_write( + self, mock_router_instance, basic_config + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky_with_escalation") + await self._seed_pin(mock_router_instance, ComplexityTier.MEDIUM, "sess-e1") + set_cache_mock = AsyncMock() + with patch.object(mock_router_instance.cache, "async_set_cache", set_cache_mock): + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-e1"), + messages=SIMPLE_MESSAGES, + ) + assert result is not None + assert result.model == "gpt-4o" + set_cache_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_classified_above_pin_escalates_and_rewrites( + self, mock_router_instance, basic_config + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky_with_escalation") + await self._seed_pin(mock_router_instance, ComplexityTier.SIMPLE, "sess-e2") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-e2"), + messages=COMPLEX_MESSAGES, + ) + assert result is not None + assert result.model != "gpt-4o-mini" + cached = await mock_router_instance.cache.async_get_cache( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", "sess-e2") + ) + assert cached != ComplexityTier.SIMPLE.value + + @pytest.mark.asyncio + async def test_max_tier_pin_skips_all_classification_work( + self, mock_router_instance, basic_config + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky_with_escalation") + await self._seed_pin(mock_router_instance, ComplexityTier.REASONING, "sess-e3") + with ( + patch.object(router, "aclassify", new_callable=AsyncMock) as classify_mock, + patch.object( + router, "_resolve_keyword_tier_override", new_callable=AsyncMock + ) as override_mock, + ): + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-e3"), + messages=SIMPLE_MESSAGES, + ) + assert result is not None + assert result.model == "o1-preview" + classify_mock.assert_not_awaited() + override_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_keyword_override_escalates_a_pin(self, mock_router_instance, basic_config): + from litellm.caching.caching import DualCache + + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-sticky-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_stickiness": {"mode": "sticky_with_escalation"}, + "keyword_tier_rules": [{"keywords": ["wire transfer"], "tier": "REASONING"}], + }, + ) + await self._seed_pin(mock_router_instance, ComplexityTier.MEDIUM, "sess-e4") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-e4"), + messages=[{"role": "user", "content": "please process this wire transfer"}], + ) + assert result is not None + assert result.model == "o1-preview" + + @pytest.mark.asyncio + async def test_keyword_override_cannot_demote_a_pin(self, mock_router_instance, basic_config): + from litellm.caching.caching import DualCache + + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-sticky-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_stickiness": {"mode": "sticky_with_escalation"}, + "keyword_tier_rules": [{"keywords": ["hello"], "tier": "SIMPLE"}], + }, + ) + await self._seed_pin(mock_router_instance, ComplexityTier.COMPLEX, "sess-e5") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-e5"), + messages=[{"role": "user", "content": "hello there"}], + ) + assert result is not None + assert result.model == "claude-sonnet-4-20250514" + + +class TestSessionStickyCauseLogging: + """Pinned and escalated decisions must be distinguishable in the info log.""" + + @pytest.fixture + def router_log_capture(self, caplog): + caplog.set_level(logging.INFO, logger="LiteLLM Router") + verbose_router_logger.addHandler(caplog.handler) + try: + yield caplog + finally: + verbose_router_logger.removeHandler(caplog.handler) + + @pytest.mark.asyncio + async def test_pinned_decision_logs_session_pinned( + self, mock_router_instance, basic_config, router_log_capture + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky") + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-log1"), + messages=SIMPLE_MESSAGES, + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-log1"), + messages=SIMPLE_MESSAGES, + ) + assert "routing decision cause=session_pinned" in router_log_capture.text + assert "routed_model=gpt-4o-mini" in router_log_capture.text + + @pytest.mark.asyncio + async def test_escalation_logs_session_escalation( + self, mock_router_instance, basic_config, router_log_capture + ): + router = _sticky_router(mock_router_instance, basic_config, "sticky_with_escalation") + await mock_router_instance.cache.async_set_cache( + ComplexityRouter.get_session_tier_cache_key("test-sticky-router", "sess-log2"), + ComplexityTier.SIMPLE.value, + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=_session_kwargs("sess-log2"), + messages=COMPLEX_MESSAGES, + ) + assert "routing decision cause=session_escalation" in router_log_capture.text + + +class TestDeploymentAffinityAutoEnable: + """enable_deployment_session_affinity wires session_affinity for the tier target groups.""" + + @staticmethod + def _router(complexity_router_config: Dict, **router_kwargs) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + { + "model_name": "auto-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": complexity_router_config, + }, + }, + ], + **router_kwargs, + ) + + _TIERS = { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "gpt-4o", + "REASONING": "gpt-4o", + } + + def test_opted_in_config_enables_affinity_for_tier_groups(self): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = self._router( + { + "tiers": self._TIERS, + "session_stickiness": { + "mode": "sticky", + "enable_deployment_session_affinity": True, + }, + } + ) + assert router.model_group_affinity_config is not None + for group in ("gpt-4o-mini", "gpt-4o"): + assert "session_affinity" in router.model_group_affinity_config[group] + affinity_callbacks = [ + cb for cb in (router.optional_callbacks or []) if isinstance(cb, DeploymentAffinityCheck) + ] + assert len(affinity_callbacks) == 1 + assert affinity_callbacks[0].enable_session_id_affinity is False + assert affinity_callbacks[0].enable_user_key_affinity is False + + def test_opt_out_is_a_full_noop(self): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = self._router( + {"tiers": self._TIERS, "session_stickiness": {"mode": "sticky"}} + ) + assert router.model_group_affinity_config is None + assert not any( + isinstance(cb, DeploymentAffinityCheck) for cb in (router.optional_callbacks or []) + ) + + def test_user_affinity_entry_is_extended_not_replaced(self): + router = self._router( + { + "tiers": self._TIERS, + "session_stickiness": { + "mode": "sticky", + "enable_deployment_session_affinity": True, + }, + }, + model_group_affinity_config={"gpt-4o": ["deployment_affinity"]}, + ) + assert router.model_group_affinity_config is not None + assert set(router.model_group_affinity_config["gpt-4o"]) == { + "deployment_affinity", + "session_affinity", + } + + def test_existing_affinity_callback_is_reused_and_seeded(self): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = self._router( + { + "tiers": self._TIERS, + "session_stickiness": { + "mode": "sticky", + "enable_deployment_session_affinity": True, + }, + }, + optional_pre_call_checks=["deployment_affinity"], + ) + affinity_callbacks = [ + cb for cb in (router.optional_callbacks or []) if isinstance(cb, DeploymentAffinityCheck) + ] + assert len(affinity_callbacks) == 1 + assert router.model_group_affinity_config is not None + for group in ("gpt-4o-mini", "gpt-4o"): + assert set(router.model_group_affinity_config[group]) == { + "deployment_affinity", + "session_affinity", + } + assert affinity_callbacks[0].model_group_affinity_config == router.model_group_affinity_config