From 24964fdea057376768b206d9d039d126f01903fd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:49 +0000 Subject: [PATCH] refactor(redis): migrate complexity session affinity into a typed store --- .../stores/complexity_session_affinity.py | 82 +++++++ .../complexity_router/complexity_router.py | 109 +++++++-- .../test_complexity_session_affinity.py | 230 ++++++++++++++++++ 3 files changed, 397 insertions(+), 24 deletions(-) create mode 100644 litellm/redis/stores/complexity_session_affinity.py create mode 100644 tests/unit/redis/stores/test_complexity_session_affinity.py diff --git a/litellm/redis/stores/complexity_session_affinity.py b/litellm/redis/stores/complexity_session_affinity.py new file mode 100644 index 00000000000..703b6b27c96 --- /dev/null +++ b/litellm/redis/stores/complexity_session_affinity.py @@ -0,0 +1,82 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_READ: Final = """ +local current = redis.call('GET', KEYS[1]) +if current ~= false then + redis.call('EXPIRE', KEYS[1], ARGV[1]) +end +return current +""" + + +def affinity_key(model_name: str, caller_scope: str, session_id: str) -> str: + return f"complexity_router_session_affinity:v1:{model_name}:{caller_scope}:{session_id}" + + +def pin_value(model: str, tier: str | None) -> Mapping[str, str | None]: + return {"model": model, "tier": tier} + + +@dataclass(frozen=True, slots=True) +class ComplexitySessionPin: + value: object + + +class ComplexitySessionAffinityStore(Protocol): + async def read_and_refresh( + self, model_name: str, caller_scope: str, session_id: str, *, ttl_seconds: int + ) -> ComplexitySessionPin | RedisFailure: ... + + async def overwrite( + self, model_name: str, caller_scope: str, session_id: str, model: str, tier: str | None, *, ttl_seconds: int + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportComplexitySessionAffinityStore: + transport: RedisTransport + coalesce_reads: bool = False + + async def read_and_refresh( + self, model_name: str, caller_scope: str, session_id: str, *, ttl_seconds: int + ) -> ComplexitySessionPin | RedisFailure: + key: Final = affinity_key(model_name, caller_scope, session_id) + batch: Final = self.transport.request_batch() if self.coalesce_reads else None + raw: Final = ( + await batch.eval(_READ, (key,), (int(ttl_seconds),)) + if batch is not None + else await self.transport.eval(_READ, (key,), (int(ttl_seconds),)) + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return ComplexitySessionPin(None) + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return ComplexitySessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return ComplexitySessionPin(decoded) + + async def overwrite( + self, model_name: str, caller_scope: str, session_id: str, model: str, tier: str | None, *, ttl_seconds: int + ) -> None | RedisFailure: + result: Final = await self.transport.set_seconds( + affinity_key(model_name, caller_scope, session_id), + json.dumps(pin_value(model, tier)), + ttl=timedelta(seconds=int(ttl_seconds)), + ) + if isinstance(result, RedisFailure): + return result + return None if result else RedisFailure(kind="decode", message="Unexpected session overwrite reply") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 30f1021251f..a6adc14dfd9 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -36,7 +36,8 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, read_affinity_pin +from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, set_local_affinity_pin +from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import ( COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, EMPTY_MAPPING, @@ -64,6 +65,17 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG +from litellm.redis.stores.complexity_session_affinity import ( + ComplexitySessionAffinityStore, + ComplexitySessionPin, + TransportComplexitySessionAffinityStore, +) +from litellm.redis.stores.complexity_session_affinity import ( + affinity_key as session_affinity_key, +) +from litellm.redis.stores.complexity_session_affinity import ( + pin_value as session_pin_value, +) from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1388,7 +1400,7 @@ def _parse_session_affinity_pin(value: object, active_tiers: tuple[str, ...]) -> def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]: tier_value: Final = _tier_name(tier) if tier is not None else None - return {"model": model, "tier": tier_value} + return session_pin_value(model, tier_value) @dataclass(frozen=True, slots=True) @@ -4412,7 +4424,7 @@ class ComplexityRouter(CustomLogger): return self._get_session_id_from_request_kwargs(request_kwargs) @staticmethod - def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_user_api_key_hash_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: """Resolve the proxy-derived API key hash, the same trust boundary DeploymentAffinityCheck uses for its own key-based affinity (not the client-supplied OpenAI `user` param, which isn't authenticated).""" @@ -4422,13 +4434,69 @@ class ComplexityRouter(CustomLogger): return str(user_key) return None - def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: + def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: Mapping[str, object]) -> str: # Namespace by the caller's API key hash so two different callers reusing the # same client-supplied session_id can't poison each other's routing pin. Falls # back to "unscoped" only when there's no authenticated caller to scope by # (e.g. direct Router usage without the proxy layer). caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" - return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + return session_affinity_key(self.model_name, caller_scope, session_id) + + async def _read_session_affinity_pin(self, session_id: str, request_kwargs: Mapping[str, object]) -> object: + cache: Final = self.litellm_router_instance.cache + cache_key: Final = self._get_session_affinity_cache_key(session_id, request_kwargs) + ttl: Final = self.config.session_affinity_ttl_seconds + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + batch: Final = active_request_redis_batch(redis_cache) + store: Final[ComplexitySessionAffinityStore] = TransportComplexitySessionAffinityStore( + redis_cache.transport(), coalesce_reads=batch is not None and batch.pending > 0 + ) + result: Final = await store.read_and_refresh( + self.model_name, + self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped", + session_id, + ttl_seconds=ttl, + ) + if isinstance(result, ComplexitySessionPin): + if result.value is None: + cache.in_memory_cache.delete_cache(cache_key) + else: + set_local_affinity_pin(cache, cache_key, result.value, ttl) + return result.value + local: Final = TypeAdapter(object).validate_python(cache.in_memory_cache.get_cache(cache_key)) + if local is not None: + set_local_affinity_pin(cache, cache_key, local, ttl) + return local + + async def _store_session_affinity_pin( + self, + session_id: str, + request_kwargs: Mapping[str, object], + model: str, + tier: ComplexityTier | str | None, + ) -> None: + cache: Final = self.litellm_router_instance.cache + ttl: Final = self.config.session_affinity_ttl_seconds + await cache.async_set_cache( + key=self._get_session_affinity_cache_key(session_id, request_kwargs), + value=_session_affinity_cache_value(model, tier), + ttl=ttl, + local_only=True, + ) + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + store: Final[ComplexitySessionAffinityStore] = TransportComplexitySessionAffinityStore( + redis_cache.transport() + ) + await store.overwrite( + self.model_name, + self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped", + session_id, + model, + _tier_name(tier) if tier is not None else None, + ttl_seconds=ttl, + ) @property def _uses_tier_pin(self) -> bool: @@ -4456,7 +4524,7 @@ class ComplexityRouter(CustomLogger): async def async_pre_routing_hook( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -4496,7 +4564,6 @@ class ComplexityRouter(CustomLogger): use_session_affinity: Final = self._uses_tier_pin session_id: Final = self._get_tier_pin_session_id(request_kwargs) if use_session_affinity else None - cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human # ask falls through and re-classifies. session_affinity restores pin-first for asks too. @@ -4507,10 +4574,8 @@ class ComplexityRouter(CustomLogger): or codex_compaction ) - if cache_key is not None and pin_replay_allowed: - pinned_value: Final = await read_affinity_pin( - self.litellm_router_instance.cache, cache_key, self.config.session_affinity_ttl_seconds - ) + if session_id is not None and pin_replay_allowed: + pinned_value: Final = await self._read_session_affinity_pin(session_id, request_kwargs) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: user_message: Final = ( @@ -4608,10 +4673,8 @@ class ComplexityRouter(CustomLogger): else floor_model ) if escalated: - await self.litellm_router_instance.cache.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value(session_model, resolved_pin_tier), - ttl=self.config.session_affinity_ttl_seconds, + await self._store_session_affinity_pin( + session_id, request_kwargs, session_model, resolved_pin_tier ) if self.config.adaptive: from litellm.router_strategy.adaptive_router.config import ( @@ -4705,19 +4768,17 @@ class ComplexityRouter(CustomLogger): # the hard floor constrained its pick, so pinning it would carry a plan-mode-shaped # choice past plan mode's exit. No sentinel turn writes the pin, whatever its cause. pinnable: Final = ( - cache_key is not None + session_id is not None and response is not None and _decision_is_pinnable(response.routing_decision) and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None ) - if pinnable and cache_key is not None and response is not None: - await self.litellm_router_instance.cache.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value( - response.model, - response.routing_decision.get("tier") if response.routing_decision is not None else None, - ), - ttl=self.config.session_affinity_ttl_seconds, + if pinnable and session_id is not None and response is not None: + await self._store_session_affinity_pin( + session_id, + request_kwargs, + response.model, + response.routing_decision.get("tier") if response.routing_decision is not None else None, ) return self._with_session_deployment_affinity(response) diff --git a/tests/unit/redis/stores/test_complexity_session_affinity.py b/tests/unit/redis/stores/test_complexity_session_affinity.py new file mode 100644 index 00000000000..2e73c230f04 --- /dev/null +++ b/tests/unit/redis/stores/test_complexity_session_affinity.py @@ -0,0 +1,230 @@ +import json +from collections.abc import Iterator, Mapping +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm import Router +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.types.router import PreRoutingHookResponse +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + legacy_get, + local_get, + local_set, + seed_async_client, +) + + +async def _store_pin( + strategy: ComplexityRouter, session_id: str, request_kwargs: Mapping[str, object], model: str, tier: str | None +) -> None: + await strategy._store_session_affinity_pin(session_id, request_kwargs, model, tier) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +async def _read_pin(strategy: ComplexityRouter, session_id: str, request_kwargs: Mapping[str, object]) -> object: + return await strategy._read_session_affinity_pin(session_id, request_kwargs) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +async def _pre_routing( + strategy: ComplexityRouter, *, model: str, request_kwargs: dict[str, object], messages: list[dict[str, object]] +) -> PreRoutingHookResponse | None: + response: PreRoutingHookResponse | None = await strategy.async_pre_routing_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook API is untyped + model=model, request_kwargs=request_kwargs, messages=messages + ) + return response + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _strategy(server: fakeredis.FakeServer | None = None) -> ComplexityRouter: + parent: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/test"}, "model_info": {"id": name}} + for name in ("tiny", "middle", "large", "top") + ] + ) + if server is not None: + parent.cache.redis_cache = fake_redis_cache(server, namespace="complexity") + return ComplexityRouter( + model_name="router", + litellm_router_instance=parent, + derive_savings_baseline=False, + complexity_router_config={ + "tiers": {"SIMPLE": "tiny", "MEDIUM": "middle", "COMPLEX": "large", "REASONING": "top"}, + "session_affinity": True, + "session_affinity_ttl_seconds": 60, + }, + ) + + +def _key(caller: str = "caller") -> str: + return f"complexity:complexity_router_session_affinity:v1:router:{caller}:same" + + +def _kwargs(caller: str = "caller") -> dict[str, object]: + return {"metadata": {"session_id": "same", "user_api_key_hash": caller}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tier", [None, "自定义"]) +async def test_overwrite_preserves_wire_bytes_scope_and_ttl(tier: str | None) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await _store_pin(strategy, "same", _kwargs(), "模型", tier) + expected: Final = json.dumps({"model": "模型", "tier": tier}).encode() + assert await raw.get(_key()) == expected + assert await raw.pttl(_key()) == 60000 + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + assert await legacy_get(redis_cache, _key().removeprefix("complexity:")) == {"model": "模型", "tier": tier} + await _store_pin(strategy, "same", _kwargs(), "replacement", None) + await _store_pin(strategy, "same", _kwargs("second"), "other", None) + assert await raw.get(_key()) == b'{"model": "replacement", "tier": null}' + assert await raw.get(_key("second")) == b'{"model": "other", "tier": null}' + assert await keys_matching(raw, "complexity:complexity:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0, -1, 0.5, 1.9]) +async def test_overwrite_preserves_legacy_integer_seconds_and_invalid_expiry(ttl: float) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + strategy.config = strategy.config.model_copy(update={"session_affinity_ttl_seconds": ttl}) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"old"', ex=20) + await _store_pin(strategy, "same", _kwargs(), "tiny", None) + assert await raw.get(_key()) == (b'{"model": "tiny", "tier": null}' if int(ttl) > 0 else b'"old"') + assert await raw.pttl(_key()) == (int(ttl) * 1000 if int(ttl) > 0 else 20000) + + +@pytest.mark.asyncio +async def test_nonpositive_read_refresh_returns_old_pin_but_deletes_shared_key() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + strategy.config = strategy.config.model_copy(update={"session_affinity_ttl_seconds": 0}) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"tiny"', ex=20) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + with patch("time.time", return_value=1700000000.001): + assert await raw.get(_key()) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'"tiny"', "tiny"), + (b"tiny", "tiny"), + (b'{"model": "tiny", "tier": null}', {"model": "tiny", "tier": None}), + (b"malformed{", "malformed{"), + ], +) +async def test_read_refreshes_without_reencoding_legacy_values(wire: bytes, expected: object) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + assert await _read_pin(strategy, "same", _kwargs()) == expected + assert await raw.get(_key()) == wire + assert await raw.pttl(_key()) == 60000 + + +@pytest.mark.asyncio +async def test_successful_miss_invalidates_but_outage_refreshes_local_fallback() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + logical: Final = _key().removeprefix("complexity:") + local_set(cache, logical, "stale", ttl=20) + assert await _read_pin(strategy, "same", _kwargs()) is None + assert local_get(cache, logical) is None + local_set(cache, logical, "local", ttl=20) + server.connected = False + assert await _read_pin(strategy, "same", _kwargs()) == "local" + with patch("time.time", return_value=1700000021.0): + assert local_get(cache, logical) == "local" + + +@pytest.mark.asyncio +async def test_attached_redis_supersedes_local_only_session() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy() + await _store_pin(strategy, "same", _kwargs(), "local", None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"shared"', ex=20) + strategy.litellm_router_instance.cache.redis_cache = fake_redis_cache(server, namespace="complexity") + assert await _read_pin(strategy, "same", _kwargs()) == "shared" + + +class LostWriteReply(fakeredis.FakeAsyncRedis): + async def set(self, *args: object, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # passes the write through fakeredis + await super().set(*args, **kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after overwrite") + + +@pytest.mark.asyncio +async def test_unknown_overwrite_keeps_the_value_already_written_locally() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, LostWriteReply(server=server)) + await _store_pin(strategy, "same", _kwargs(), "tiny", None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key()) == b'{"model": "tiny", "tier": null}' + assert local_get(strategy.litellm_router_instance.cache, _key().removeprefix("complexity:")) == { + "model": "tiny", + "tier": None, + } + + +@pytest.mark.asyncio +async def test_classification_and_pin_escalation_both_overwrite_shared_session() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + first: Final = await _pre_routing( + strategy, model="router", request_kwargs=_kwargs(), messages=[{"role": "user", "content": "Hi"}] + ) + assert first is not None and first.model == "tiny" + assert await raw.get(_key()) == b'{"model": "tiny", "tier": "SIMPLE"}' + second: Final = await _pre_routing( + strategy, model="router", request_kwargs=_kwargs(), messages=[{"role": "user", "content": "LITELLM ESCALATE"}] + ) + assert second is not None and second.model == "middle" + assert await raw.get(_key()) == b'{"model": "middle", "tier": "MEDIUM"}' + + +@pytest.mark.asyncio +async def test_reads_join_only_pending_request_work_not_post_call_writes() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache + assert redis_cache is not None + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), b'"tiny"', ex=20) + with request_redis_batch_scope(post_call_deadline=3600) as batches: + request: Final = batches.batch(redis_cache) + post_call: Final = batches.post_call(redis_cache) + post_call.set("unrelated", "later", 60) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + assert request.flushes == 0 + request.set("before-read", "ready", 60) + assert await _read_pin(strategy, "same", _kwargs()) == "tiny" + assert request.flushes == 1 + assert await raw.get("complexity:before-read") == b'"ready"' + assert await raw.get("complexity:unrelated") is None + assert post_call.pending == 1 + await batches.flush_post_call() + assert await raw.pttl(_key()) == 60000