refactor(redis): migrate complexity session affinity into a typed store

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")

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)
 

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