mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(complexity_router): session-sticky tier routing to preserve provider prompt cache
The complexity router re-classifies the latest user message on every turn, so a multi-turn conversation can swap models mid-session (haiku on turn 1, opus on turn 2). Every swap lands the conversation on a cold model: no provider KV or prompt cache hit, full input reprocessed at uncached rates. The router's own LLM classifier and semantic query embeddings also fire on every turn. Add session_stickiness to complexity_router_config with three modes. "sticky" pins the first decision per session_id (from the existing x-litellm-session-id header extraction) in the router's DualCache and skips keyword overrides and the classifier entirely on later turns. "sticky_with_escalation" lets the pin ratchet up in severity but never down: downgrades destroy a warm cache to save pennies, while an upgrade breaks cache once for a smarter model; keyword overrides still fire mid-session and can escalate but not demote. "none" keeps today's per-turn behavior and stays the default. The pin stores the tier, not the model name, so a tiers remap between deploys stays meaningful. An enable_deployment_session_affinity opt-in extends model_group_affinity_config with session_affinity for the router's tier target model groups, so the existing DeploymentAffinityCheck pins the same physical deployment within each group and the whole session hits one warm deployment end to end. New entries are seeded with the currently enabled global flags because explicit per-group entries make _get_effective_flags ignore the callback's global flags. The wiring runs after __init__ assigns model_group_affinity_config (set_model_list runs earlier) and again on runtime hot-adds. Pinned decisions log cause=session_pinned and ratchets log cause=session_escalation, consistent with the existing cause= decision format. Concurrent first turns of a session are last-write-wins (DualCache has no SETNX) and self-consistent from the next turn.
This commit is contained in:
Abhimanyu Kapur
parent
0ebcda3027
commit
269ad1f36a
4 changed files with 678 additions and 19 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue