diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2b8dc67da7a..aa461249ac7 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -1,72 +1,12 @@ -"""Atomic affinity claims shared by deployment and tier-model selection.""" +"""Pod-local affinity policy shared by deployment and tier-model selection.""" -import json from collections.abc import Mapping -from typing import ( - Final, - cast, # noqa: TID251 # Heterogeneous cache results are narrowed only to object, then validated -) +from typing import Final -from pydantic import JsonValue, TypeAdapter, ValidationError - -from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_batch import active_request_redis_batch ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins" -_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) - -_READ_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current ~= false then - redis.call('EXPIRE', KEYS[1], ARGV[1]) -end -return current -""" - -_CLAIM_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if ARGV[3] then - local decoded, stored = pcall(cjson.decode, current) - if ARGV[4] then - local legacy = current - if decoded and type(stored) == 'string' then legacy = stored end - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - if eligible['model_id'] == legacy then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - if decoded and type(stored) == 'table' then - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - local matches = true - for key, value in pairs(eligible) do - if stored[key] ~= value then matches = false; break end - end - for key, _ in pairs(stored) do - if eligible[key] == nil then matches = false; break end - end - if matches then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" - def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" @@ -103,90 +43,3 @@ def claim_affinity_pin_in_memory( winner: Final = existing if existing is not None and eligible_existing else pin_value set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) return winner - - -def _decode_pin(value: str) -> object: - try: - return _PIN_JSON_ADAPTER.validate_json(value) - except ValidationError: - return value - - -async def read_affinity_pin(cache: DualCache, cache_key: str, ttl_seconds: int) -> object: - """Refresh the authoritative pin without writing a previously observed value back.""" - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - try: - read_script: Final = redis_cache.async_register_script(_READ_PIN_SCRIPT) - batch: Final = active_request_redis_batch(redis_cache) - raw: Final = cast( # cast-ok: Redis script returns a serialized pin or a miss; validate before use - object, - await batch.script(_READ_PIN_SCRIPT, read_script, (cache_key,), (int(ttl_seconds),)) - if batch is not None and batch.pending - else await read_script(keys=(cache_key,), args=(int(ttl_seconds),)), - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - value: Final = _decode_pin(decoded) if isinstance(decoded, str) else None - if value is None: - cache.in_memory_cache.delete_cache(cache_key) - else: - set_local_affinity_pin(cache, cache_key, value, ttl_seconds) - except Exception as error: # noqa: BLE001 # Redis faults retain pod-local continuity, unlike a successful miss - verbose_router_logger.debug("Affinity cache: Redis read failed, using pod-local pin. error=%s", error) - else: - return value - local_value: Final = cast( # cast-ok: the heterogeneous local cache is parsed by the affinity caller - object, cache.in_memory_cache.get_cache(cache_key) - ) - if local_value is not None: - set_local_affinity_pin(cache, cache_key, local_value, ttl_seconds) - return local_value - - -async def claim_affinity_pin( - cache: DualCache, - cache_key: str, - pin_value: Mapping[str, str], - ttl_seconds: int, - *, - eligible_values: tuple[Mapping[str, str], ...] | None = None, - preserve_legacy_model_ids: bool = False, -) -> object: - """Return the authoritative first writer, replacing it only when it becomes ineligible. - - Eligible claims refresh the returned winner. Legacy deployment claims only refresh - a matching candidate. Resolve Redis per call because the proxy attaches it lazily. - """ - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - args: Final = ( - json.dumps(dict(pin_value)), - int(ttl_seconds), - *( - (json.dumps(tuple(dict(value) for value in eligible_values)),) - if eligible_values is not None - else () - ), - *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), - ) - raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here - object, await claim_script(keys=(cache_key,), args=args) - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): - return pin_value - winner: Final = _decode_pin(decoded) - set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) - return winner - except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims - verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) - return claim_affinity_pin_in_memory( - cache, - cache_key, - pin_value, - ttl_seconds, - eligible_values=eligible_values, - preserve_legacy_model_ids=preserve_legacy_model_ids, - ) diff --git a/litellm/redis/stores/complexity_tier_model_affinity.py b/litellm/redis/stores/complexity_tier_model_affinity.py new file mode 100644 index 00000000000..18d3057f9cf --- /dev/null +++ b/litellm/redis/stores/complexity_tier_model_affinity.py @@ -0,0 +1,103 @@ +import hashlib +import json +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_CLAIM: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +local decoded, stored = pcall(cjson.decode, current) +if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end +end +redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) +return ARGV[1] +""" + + +def configuration_fingerprint(config_json: str) -> str: + return hashlib.sha256(config_json.encode()).hexdigest() + + +@dataclass(frozen=True, slots=True) +class TierModelIdentity: + model_name: str + configuration: str + caller: str | None + session_id: str + tier: str + replay_fallback: str | None = None + + +def affinity_key(identity: TierModelIdentity) -> str: + parts: Final = (identity.model_name, identity.configuration, identity.caller, identity.session_id, identity.tier) + cache_identity: Final = ( + (*parts, ("replay_fallback", identity.replay_fallback)) if identity.replay_fallback is not None else parts + ) + return "complexity_router_tier_model_affinity:v1:" + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() + + +@dataclass(frozen=True, slots=True) +class TierModelPin: + value: object + + +@dataclass(frozen=True, slots=True) +class UnconfirmedTierModelPin: + pass + + +class ComplexityTierModelAffinityStore(Protocol): + async def claim( + self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int + ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportComplexityTierModelAffinityStore: + transport: RedisTransport + + async def claim( + self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int + ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure: + raw: Final = await self.transport.eval( + _CLAIM, + (affinity_key(identity),), + ( + json.dumps({"model": model}), + int(ttl_seconds), + json.dumps(tuple({"model": value} for value in eligible_models)), + ), + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return UnconfirmedTierModelPin() + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return TierModelPin(_JSON.validate_json(decoded)) + except ValidationError: + return TierModelPin(decoded) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a6adc14dfd9..673c5b167c0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -17,7 +17,6 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter from __future__ import annotations import asyncio -import hashlib import json import random import re @@ -36,7 +35,11 @@ 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, set_local_affinity_pin +from litellm.caching.affinity_cache import ( + ROUTER_SESSION_PINS_TARGET, + claim_affinity_pin_in_memory, + 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, @@ -76,6 +79,18 @@ from litellm.redis.stores.complexity_session_affinity import ( from litellm.redis.stores.complexity_session_affinity import ( pin_value as session_pin_value, ) +from litellm.redis.stores.complexity_tier_model_affinity import ( + ComplexityTierModelAffinityStore, + TierModelIdentity, + TierModelPin, + TransportComplexityTierModelAffinityStore, + UnconfirmedTierModelPin, + configuration_fingerprint, +) +from litellm.redis.stores.complexity_tier_model_affinity import ( + affinity_key as tier_model_affinity_key, +) +from litellm.redis.transport.contract import RedisFailure 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 ( @@ -1515,9 +1530,9 @@ class ComplexityRouter(CustomLogger): case _: self._jev_client: JevClassifierClient | None = jev_client_result - self._tier_affinity_config = hashlib.sha256( - self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() - ).hexdigest() + self._tier_affinity_config = configuration_fingerprint( + self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})) + ) # Checked here rather than on the config model because the deployment's # complexity_router_default_model arrives outside complexity_router_config and is @@ -3125,27 +3140,44 @@ class ComplexityRouter(CustomLogger): if session_id is None: return model caller: Final = DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs) - identity: Final = (self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier)) - cache_identity: Final = ( - (*identity, ("replay_fallback", retained_model)) if retained_model is not None else identity - ) - cache_key: Final = ( - "complexity_router_tier_model_affinity:v1:" - + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() - ) - winner: Final = await claim_affinity_pin( - self.litellm_router_instance.cache, - cache_key, - MappingProxyType({"model": model}), - self.config.session_affinity_ttl_seconds, - eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + identity: Final = TierModelIdentity( + self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier), retained_model ) + winner: Final = await self._claim_tier_model_pin(identity, model, candidates) pinned: Final[object] = winner.get("model") if isinstance(winner, Mapping) else None if not isinstance(pinned, str) or pinned not in candidates: return model self._restamp_adaptive_choice(request_kwargs, model, pinned) return pinned + async def _claim_tier_model_pin( + self, identity: TierModelIdentity, model: str, candidates: tuple[str, ...] + ) -> object: + cache: Final = self.litellm_router_instance.cache + cache_key: Final = tier_model_affinity_key(identity) + ttl: Final = self.config.session_affinity_ttl_seconds + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + store: Final[ComplexityTierModelAffinityStore] = TransportComplexityTierModelAffinityStore( + redis_cache.transport() + ) + result: Final = await store.claim(identity, model, candidates, ttl_seconds=ttl) + match result: + case TierModelPin(value=value): + set_local_affinity_pin(cache, cache_key, value, ttl) + return value + case UnconfirmedTierModelPin(): + return MappingProxyType({"model": model}) + case RedisFailure(): + pass + return claim_affinity_pin_in_memory( + cache, + cache_key, + MappingProxyType({"model": model}), + ttl, + eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + ) + async def _pick_model_for_tier( self, tier: ComplexityTier | str, diff --git a/tests/unit/redis/stores/test_complexity_tier_model_affinity.py b/tests/unit/redis/stores/test_complexity_tier_model_affinity.py new file mode 100644 index 00000000000..f7dcd083ee3 --- /dev/null +++ b/tests/unit/redis/stores/test_complexity_tier_model_affinity.py @@ -0,0 +1,234 @@ +import asyncio +import hashlib +import json +from collections.abc import Iterator +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, + _SessionAffinityPin, # pyright: ignore[reportPrivateUsage] # builds the retained pin the strategy carries between tiers +) +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + local_get, + local_set, + seed_async_client, +) + + +async def _pin( + strategy: ComplexityRouter, + tier: str, + model: str, + candidates: tuple[str, ...], + request_kwargs: dict[str, object], + retained_pin: _SessionAffinityPin | None = None, +) -> object: + return await strategy._pin_model_for_tier(tier, model, candidates, request_kwargs, retained_pin) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _strategy(server: fakeredis.FakeServer | None, *, name: str = "router", simple: str = "tiny") -> 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="tier") + return ComplexityRouter( + model_name=name, + litellm_router_instance=parent, + derive_savings_baseline=False, + complexity_router_config={ + "tiers": {"SIMPLE": [simple, "middle"], "MEDIUM": "middle", "COMPLEX": "large", "REASONING": "top"}, + "session_affinity": True, + "session_affinity_ttl_seconds": 60, + }, + ) + + +def _key(strategy: ComplexityRouter, caller: str | None = "caller", *, fallback: str | None = None) -> str: + config: Final = strategy.config.model_dump_json(include={"tiers": True, "tier_model_configs": True}) + identity: Final = (strategy.model_name, hashlib.sha256(config.encode()).hexdigest(), caller, "same", "SIMPLE") + encoded: Final = json.dumps((*identity, ("replay_fallback", fallback)) if fallback is not None else identity) + return "tier:complexity_router_tier_model_affinity:v1:" + hashlib.sha256(encoded.encode()).hexdigest() + + +def _kwargs(caller: str | None = "caller") -> dict[str, object]: + return {"metadata": {"session_id": "same", "user_api_key_hash": caller}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "name,caller,simple,fallback", + [ + ("router", "caller", "tiny", None), + ("路由", "caller", "tiny", None), + ("router", "另一个", "tiny", None), + ("router", None, "tiny", None), + ("router", "caller", "模型", None), + ("router", "caller", "tiny", "removed"), + ], +) +async def test_identity_and_unicode_wire_match_legacy_hash( + name: str, caller: str | None, simple: str, fallback: str | None +) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server, name=name, simple=simple) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + retained: Final = _SessionAffinityPin(fallback, "SIMPLE") if fallback is not None else None + assert await _pin(strategy, "SIMPLE", "模型", ("模型", "tiny"), _kwargs(caller), retained) == "模型" + assert await raw.get(_key(strategy, caller, fallback=fallback)) == b'{"model": "\\u6a21\\u578b"}' + assert await raw.pttl(_key(strategy, caller, fallback=fallback)) == 60000 + assert await keys_matching(raw, "tier:tier:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'{"model": "tiny"}', "tiny"), + (b'{"model": "removed"}', "middle"), + (b'{"model": "tiny", "extra": true}', "middle"), + (b'"tiny"', "middle"), + (b"tiny", "middle"), + (b"{broken", "middle"), + ], +) +async def test_claim_refreshes_only_complete_eligible_winners(wire: bytes, expected: str) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(strategy), wire, ex=20) + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == expected + assert await raw.get(_key(strategy)) == json.dumps({"model": expected}).encode() + assert await raw.pttl(_key(strategy)) == 60000 + + +@pytest.mark.asyncio +async def test_concurrent_pods_converge_then_replace_expired_pin() -> None: + server: Final = fakeredis.FakeServer() + left: Final = _strategy(server) + right: Final = _strategy(server) + winners: Final = await asyncio.gather( + _pin(left, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()), + _pin(right, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()), + ) + assert winners[0] == winners[1] + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key(left)) == json.dumps({"model": winners[0]}).encode() + with patch("time.time", return_value=1700000061.0): + assert await _pin(left, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle" + assert await raw.get(_key(left)) == b'{"model": "middle"}' + assert await raw.pttl(_key(left)) == 60000 + + +@pytest.mark.asyncio +async def test_no_session_and_eligible_retained_pin_do_not_claim() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), {}) == "tiny" + assert ( + await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs(), _SessionAffinityPin("tiny", "SIMPLE")) + == "tiny" + ) + assert await fakeredis.FakeAsyncRedis(server=server).keys("tier:*") == [] + + +@pytest.mark.asyncio +async def test_late_redis_attachment_replaces_local_winner() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(None) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny" + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(strategy), b'{"model": "middle"}', ex=20) + strategy.litellm_router_instance.cache.redis_cache = fake_redis_cache(server, namespace="tier") + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "middle" + assert local_get(strategy.litellm_router_instance.cache, _key(strategy).removeprefix("tier:")) == { + "model": "middle" + } + + +class LostClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after claim") + + +class ConsumedClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + return 1 + + +@pytest.mark.asyncio +async def test_consumed_claim_returns_candidate_without_installing_local_winner() -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + redis_cache: Final = cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, ConsumedClaimReply(server=server)) + local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20) + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle" + assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == b'{"model": "middle"}' + assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"} + + +@pytest.mark.asyncio +async def test_claim_leaves_pending_request_and_post_call_work_unflushed() -> 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) + 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) + request.set("pending-request", "ready", 60) + post_call.set("pending-post-call", "later", 60) + assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny" + assert await raw.get(_key(strategy)) == b'{"model": "tiny"}' + assert await raw.get("tier:pending-request") is None + assert await raw.get("tier:pending-post-call") is None + await batches.flush_post_call() + await request.flush() + assert await raw.get("tier:pending-request") == b'"ready"' + assert await raw.get("tier:pending-post-call") == b'"later"' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("applied", [False, True]) +async def test_outage_and_unknown_claim_keep_local_winner(applied: bool) -> None: + server: Final = fakeredis.FakeServer() + strategy: Final = _strategy(server) + cache: Final = strategy.litellm_router_instance.cache + redis_cache: Final = cache.redis_cache + assert redis_cache is not None + local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20) + if applied: + seed_async_client(redis_cache, LostClaimReply(server=server)) + else: + server.connected = False + assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "tiny" + server.connected = True + assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == ( + b'{"model": "middle"}' if applied else None + ) + with patch("time.time", return_value=1700000021.0): + assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"} diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ab168b2a950..98a53def5b4 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial from types import MappingProxyType -from typing import Dict, Final, List, Literal +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -39,7 +39,6 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_strategy.auto_router.auto_router_model_naming import ( CUSTOMIZATION_CAPABILITY, - GATED_AUTO_ROUTER_CAPABILITIES, HEURISTIC_V2_CAPABILITY, count_capability_routers, ) @@ -128,7 +127,7 @@ def mock_router_instance(): @pytest.fixture -def basic_config() -> Dict: +def basic_config() -> dict: """Basic configuration with tier mappings.""" return { "tiers": { @@ -1498,7 +1497,6 @@ class TestSingletonMutation: def test_default_config_not_mutated(self, mock_router_instance): """Test that creating routers without config doesn't mutate defaults.""" from litellm.router_strategy.complexity_router.config import ( - DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ComplexityRouterConfig, ) @@ -1554,7 +1552,7 @@ class TestKeywordFalsePositives: tier, score, signals = complexity_router.classify(prompt) # 'entry' contains 'try' but should not trigger code detection # Note: 'application' might trigger something, but 'try' should not - pass # Just ensure no crash; false positive check is the main goal + # Just ensure no crash; false positive check is the main goal def test_error_not_in_terrorism(self, complexity_router): """'error' should not match in 'terrorism'.""" @@ -1696,17 +1694,19 @@ class TestRouterComplexityDeploymentMethods: def test_custom_default_resolution_preserves_explicit_precedence(self, tier_models, explicit, expected): from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig - config: Final = ComplexityRouterConfig.model_validate({ - "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"}, - "tier_definitions": [ - {"name": "CUSTOM", "description": "Custom work"}, - {"name": "OTHER", "description": "Other work"}, - ], - "fallback_tier": " CUSTOM ", - "classifier_type": "llm", - "classifier_llm_config": {"model": "classifier"}, - "default_model": "configured-model", - }) + config: Final = ComplexityRouterConfig.model_validate( + { + "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"}, + "tier_definitions": [ + {"name": "CUSTOM", "description": "Custom work"}, + {"name": "OTHER", "description": "Other work"}, + ], + "fallback_tier": " CUSTOM ", + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier"}, + "default_model": "configured-model", + } + ) assert config.resolve_default_model(explicit) == expected inferred: Final = config.model_copy(update={"default_model": None}) assert inferred.resolve_default_model() == "custom-model" @@ -2142,9 +2142,7 @@ class TestRouterComplexityDeploymentMethods: @pytest.mark.parametrize("classifier", ("heuristic_first", "hybrid")) @pytest.mark.parametrize("held_capability", ("heuristic_v2", "custom_prompt")) - def test_v2_chain_registration_enforces_each_claimed_allowance( - self, classifier: str, held_capability: str - ) -> None: + def test_v2_chain_registration_enforces_each_claimed_allowance(self, classifier: str, held_capability: str) -> None: held: Final = ( self._router_row("existing", "existing-id", "heuristic_v2") if held_capability == "heuristic_v2" @@ -2160,7 +2158,8 @@ class TestRouterComplexityDeploymentMethods: "local_heuristic": "heuristic_v2", **( {"heuristic_first_max_tier": "MEDIUM"} - if classifier == "heuristic_first" else {"hybrid_boundary_margin": 0.1} + if classifier == "heuristic_first" + else {"hybrid_boundary_margin": 0.1} ), "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "Grade by difficulty"}, "tiers": {"SIMPLE": "gpt-4o-mini"}, @@ -2785,7 +2784,7 @@ def _wrapped_reply(shape: str, verdict: str) -> str: @pytest.fixture -def llm_classifier_config() -> Dict: +def llm_classifier_config() -> dict: """Config with an LLM-based classifier wired to a 'haiku-classifier' model.""" return { "tiers": { @@ -2839,7 +2838,7 @@ class TestLLMClassifierConfig: ) -CAPABILITY_TIERS: Dict[str, str] = { +CAPABILITY_TIERS: dict[str, str] = { "SIMPLE": "efficient-model", "REASONING": "capable-model", } @@ -3415,7 +3414,7 @@ class TestCapabilityClassifier: mock_router_instance.acompletion.assert_awaited_once() -CUSTOM_TIER_LABELS: Dict[str, str] = { +CUSTOM_TIER_LABELS: dict[str, str] = { "SIMPLE": "Cheap", "MEDIUM": "Standard", "COMPLEX": "Premium", @@ -4856,7 +4855,7 @@ class TestRouterPreRoutingAliasOverrides: reach the outbound request even though the tier deployment is what actually gets called.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="smart-router", @@ -4889,7 +4888,7 @@ class TestRouterPreRoutingAliasOverrides: {"model_name": "gpt-5-mini", "litellm_params": {"model": "openai/gpt-5-mini"}}, ] ) - request_kwargs: Dict = {"reasoning_effort": "low"} + request_kwargs: dict = {"reasoning_effort": "low"} deployment = await router.async_get_available_deployment( model="smart-router", @@ -4900,7 +4899,7 @@ class TestRouterPreRoutingAliasOverrides: assert deployment["model_name"] == "gpt-5-mini" assert request_kwargs["reasoning_effort"] == "xhigh" - def _make_effort_pinned_router(self, tier_litellm_params: Dict) -> Router: + def _make_effort_pinned_router(self, tier_litellm_params: dict) -> Router: return Router( model_list=[ { @@ -4951,7 +4950,7 @@ class TestRouterPreRoutingAliasOverrides: carrier precedence over the reasoning_effort alias, so the pin only reaches the wire if those carriers are dropped at the merge.""" router = self._make_effort_pinned_router({"reasoning_effort": "xhigh"}) - request_kwargs: Dict = dict(client_carriers) + request_kwargs: dict = dict(client_carriers) await router.async_get_available_deployment( model="smart-router", @@ -4989,7 +4988,7 @@ class TestRouterPreRoutingAliasOverrides: }, ] ) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment_for_pass_through( model="smart-router", @@ -5002,15 +5001,15 @@ class TestRouterPreRoutingAliasOverrides: assert "output_config" not in request_kwargs def test_drop_client_effort_carriers_helper_edge_shapes(self): - no_pin: Dict = {"thinking": {"type": "adaptive"}} + no_pin: dict = {"thinking": {"type": "adaptive"}} Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) assert no_pin == {"thinking": {"type": "adaptive"}} - non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3} + non_dict_carriers: dict = {"output_config": "max", "reasoning": 3} Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) assert non_dict_carriers == {"output_config": "max", "reasoning": 3} - effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} + effort_only: dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} Router._pop_effort_from_nested_carrier(effort_only, "output_config") Router._pop_effort_from_nested_carrier(effort_only, "reasoning") assert effort_only == {} @@ -5038,7 +5037,7 @@ class TestRouterPreRoutingAliasOverrides: {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, ] ) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment( model="smart-router", @@ -5053,7 +5052,7 @@ class TestRouterPreRoutingAliasOverrides: @pytest.mark.asyncio async def test_client_effort_carriers_survive_when_tier_pins_no_effort(self): router = self._make_effort_pinned_router({"temperature": 0.2}) - request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} + request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}} await router.async_get_available_deployment( model="smart-router", @@ -5098,7 +5097,7 @@ class TestRouterPreRoutingAliasOverrides: ] ) real_get_llm_provider = litellm.get_llm_provider - copilot_resolutions: List = [] + copilot_resolutions: list = [] def _guarded(*args, **kwargs): target = str(kwargs.get("model") or (args[0] if args else "")) + str( @@ -5110,7 +5109,7 @@ class TestRouterPreRoutingAliasOverrides: return real_get_llm_provider(*args, **kwargs) monkeypatch.setattr(litellm, "get_llm_provider", _guarded) - request_kwargs: Dict = {} + request_kwargs: dict = {} deployment = await router.async_get_available_deployment( model="smart-router", @@ -5178,7 +5177,7 @@ class TestRouterPreRoutingAliasOverrides: test_router_init_only_params_are_never_sent_to_a_provider for the guard on that downstream filter.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} await router.async_pre_routing_hook( model="smart-router", @@ -5232,7 +5231,7 @@ class TestRouterPreRoutingAliasOverrides: """A value the caller already passed for this request takes precedence over the alias's configured default.""" router = self._make_router() - request_kwargs: Dict = {"drop_params": False} + request_kwargs: dict = {"drop_params": False} await router.async_pre_routing_hook( model="smart-router", @@ -5247,7 +5246,7 @@ class TestRouterPreRoutingAliasOverrides: """A plain (non-router-alias) model name is not affected by the alias-override merge at all.""" router = self._make_router() - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="gpt-4o-mini", @@ -5282,7 +5281,7 @@ class TestRouterPreRoutingAliasOverrides: router.set_model_list(model_list) assert "smart-router" in router.adaptive_routers - request_kwargs: Dict = {} + request_kwargs: dict = {} await router.async_pre_routing_hook( model="smart-router", request_kwargs=request_kwargs, @@ -5344,7 +5343,7 @@ class TestRouterPreRoutingSharedAliasName: else [self._marker_entry(), self._plain_entry()] ) router = Router(model_list=[*shared_name_entries, self._tier_entry()]) - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="gpt4o", @@ -5374,7 +5373,7 @@ class TestRouterPreRoutingSharedAliasName: }, } router = Router(model_list=[marker_with_connection_params, self._tier_entry()]) - request_kwargs: Dict = {} + request_kwargs: dict = {} result = await router.async_pre_routing_hook( model="smart", @@ -5414,7 +5413,7 @@ class TestRouterPreRoutingSharedAliasName: ] ) - us_kwargs: Dict = {"metadata": {"tags": ["us"]}} + us_kwargs: dict = {"metadata": {"tags": ["us"]}} us_result = await router.async_pre_routing_hook( model="smart", request_kwargs=us_kwargs, @@ -5423,7 +5422,7 @@ class TestRouterPreRoutingSharedAliasName: assert us_result is not None and us_result.model == "gpt-us" assert us_kwargs["drop_params"] is True - cn_kwargs: Dict = {"metadata": {"tags": ["cn"]}} + cn_kwargs: dict = {"metadata": {"tags": ["cn"]}} cn_result = await router.async_pre_routing_hook( model="smart", request_kwargs=cn_kwargs, @@ -5597,7 +5596,7 @@ class TestRouterPreRoutingSharedAliasName: await asyncio.sleep(0.01) return await healthy_deployments(*args, **kwargs) - sent: Dict[str, str | None] = {} + sent: dict[str, str | None] = {} async def record(**kwargs): sent[kwargs["model"]] = kwargs.get("aws_region_name") @@ -5682,7 +5681,7 @@ class TestAdaptiveSoftFloors: return router @pytest.fixture - def hybrid_config(self) -> Dict: + def hybrid_config(self) -> dict: return { "adaptive": True, "adaptive_weights": {"quality": 0.7, "cost": 0.3}, @@ -5712,7 +5711,7 @@ class TestAdaptiveSoftFloors: }, }, ) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} with patch( "litellm.router_strategy.complexity_router.complexity_router.random.choice", @@ -5811,7 +5810,7 @@ class TestAdaptiveSoftFloors: assert adaptive is not None for model in ("cheap", "premium"): adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=6.0, beta=5.0) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} with patch( "litellm.router_strategy.adaptive_router.bandit.thompson_sample", @@ -5832,7 +5831,7 @@ class TestAdaptiveSoftFloors: litellm_router_instance=adaptive_router_instance, complexity_router_config=hybrid_config, ) - request_kwargs: Dict = {"metadata": {}} + request_kwargs: dict = {"metadata": {}} result = await cr.async_pre_routing_hook( model="hybrid", request_kwargs=request_kwargs, @@ -5854,7 +5853,7 @@ class TestLexicalKeywordTierRules: """Test deterministic (literal) keyword_tier_rules overrides.""" @pytest.fixture - def rule_config(self, basic_config) -> Dict: + def rule_config(self, basic_config) -> dict: return { **basic_config, "keyword_tier_rules": [ @@ -5994,7 +5993,7 @@ class TestLexicalKeywordTierRules: class TestCjkKeywordTierRules: """CJK keyword_tier_rules must fire mid-sentence, where regex word boundaries cannot.""" - def _router(self, mock_router_instance, basic_config, keywords: List[str]) -> ComplexityRouter: + def _router(self, mock_router_instance, basic_config, keywords: list[str]) -> ComplexityRouter: return ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, @@ -6069,7 +6068,7 @@ class TestCjkKeywordTierRules: assert complexity_router._keyword_matches("appelle l' api maintenant", "api") is True -def _make_embedding_response(vectors: List[List[float]]) -> "litellm.EmbeddingResponse": +def _make_embedding_response(vectors: list[list[float]]) -> "litellm.EmbeddingResponse": return litellm.EmbeddingResponse( model="fake-embed", data=[{"embedding": vec, "index": idx, "object": "embedding"} for idx, vec in enumerate(vectors)], @@ -6088,22 +6087,22 @@ class FakeEmbeddingRouter: _CLUSTER_MARKERS = ("k8s", "kube", "container", "cluster", "orchestrat") def __init__(self): - self.async_embedding_calls: List[List[str]] = [] - self.async_embedding_kwargs: List[Dict] = [] + self.async_embedding_calls: list[list[str]] = [] + self.async_embedding_kwargs: list[dict] = [] # Every embedded batch (sync route-index build AND async query), so tests can count # builds independently of which embedding path the library happens to use. - self.embedded_batches: List[List[str]] = [] + self.embedded_batches: list[list[str]] = [] # Thread ids of the synchronous (route-index build) embedding calls, so a test can # assert the build is offloaded off the event-loop thread. - self.sync_embedding_thread_ids: List[int] = [] + self.sync_embedding_thread_ids: list[int] = [] - def _vectors(self, docs: List[str]) -> List[List[float]]: + def _vectors(self, docs: list[str]) -> list[list[float]]: return [ [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] for doc in docs ] @staticmethod - def _as_list(text) -> List[str]: + def _as_list(text) -> list[str]: return text if isinstance(text, list) else [text] def embedding(self, input, model, **kwargs): @@ -6598,7 +6597,7 @@ class _StubEncoder: """Minimal stand-in for LiteLLMRouterEncoder.aencode_queries, capturing the kwargs it was called with.""" def __init__(self): - self.aencode_queries_calls: List[Dict] = [] + self.aencode_queries_calls: list[dict] = [] async def aencode_queries(self, docs, **kwargs): self.aencode_queries_calls.append(kwargs) @@ -7198,11 +7197,11 @@ class TestSessionAffinity: SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}] @pytest.fixture - def session_affinity_config(self, basic_config) -> Dict: + def session_affinity_config(self, basic_config) -> dict: return {**basic_config, "session_affinity": True} @staticmethod - def _request_kwargs(session_id: str) -> Dict: + def _request_kwargs(session_id: str) -> dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio @@ -7609,9 +7608,9 @@ class TestSessionAffinity: assert recovered.model == latest["model"] assert cache.in_memory_cache.get_cache(key) == latest assert read_started.is_set() - assert tuple( - tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list - ) == (("MGET", "EVALSHA"),) * (2 if batch_pending else 0) + assert tuple(tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list) == ( + ("MGET", "EVALSHA"), + ) * (2 if batch_pending else 0) finally: for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict): if client_key in cached_key: @@ -7737,8 +7736,12 @@ class TestChildSessionTierPins: "model_info": {"id": f"{group}-{index}"}, } for group, index in ( - ("simple", 0), ("simple", 1), ("reasoning", 0), - ("reasoning", 1), ("reasoning-peer", 0), ("reasoning-peer", 1), + ("simple", 0), + ("simple", 1), + ("reasoning", 0), + ("reasoning", 1), + ("reasoning-peer", 0), + ("reasoning-peer", 1), ) ], ], @@ -7754,9 +7757,9 @@ class TestChildSessionTierPins: metadata_variable_name_for_route, ) - route: Final = { - "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages" - }[surface] + route: Final = {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"}[ + surface + ] bucket: Final = metadata_variable_name_for_route(route) data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}} LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( @@ -7770,54 +7773,74 @@ class TestChildSessionTierPins: return data, bucket async def _route( - self, router: Router, surface: str, session_id: str | None, prompt: str | None, - expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True, + self, + router: Router, + surface: str, + session_id: str | None, + prompt: str | None, + expected_tier: str, + expected_cause: str, + caller: str = "caller-a", + child: bool = True, expected_deployment: str | None = None, ) -> str: from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck data, bucket = self._normalized_request(surface, session_id, caller, child) turn: Final[dict[str, object]] = ( - {"role": "user", "content": prompt} if prompt is not None else - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} - if surface == "messages" else - {"type": "function_call_output", "call_id": "call_1", "output": "done"} - if surface == "responses" else - {"role": "tool", "tool_call_id": "call_1", "content": "done"} + {"role": "user", "content": prompt} + if prompt is not None + else {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} + if surface == "messages" + else {"type": "function_call_output", "call_id": "call_1", "output": "done"} + if surface == "responses" + else {"role": "tool", "tool_call_id": "call_1", "content": "done"} ) tool_call: Final = ( {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": "{}"} - if surface == "responses" else - {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}]} - if surface == "messages" else TestClassificationMode.TOOL_CALL_1 + if surface == "responses" + else { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}], + } + if surface == "messages" + else TestClassificationMode.TOOL_CALL_1 ) turns: Final = [turn] if prompt is not None else [TestClassificationMode.SIMPLE_ASK, tool_call, turn] if surface == "responses": data["input"] = turns selected: Final = await router.async_get_available_deployment( - model="child-tier-router", request_kwargs=data, + model="child-tier-router", + request_kwargs=data, messages=None if surface == "responses" else turns, input=turns if surface == "responses" else None, ) metadata: Final = data[bucket] assert isinstance(metadata, dict) assert (metadata["session_id"], metadata["trace_id"], data["litellm_session_id"], data["litellm_trace_id"]) == ( - "parent-session", "parent-session", "parent-session", "parent-session" + "parent-session", + "parent-session", + "parent-session", + "parent-session", ) assert (metadata["routing_decision"]["tier"], metadata["routing_decision"]["cause"]) == ( - expected_tier, expected_cause + expected_tier, + expected_cause, ) affinity: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) if expected_deployment is not None: filtered: Final = await affinity.async_filter_deployments( - model=selected["model_name"], healthy_deployments=router.get_model_list(selected["model_name"]) or [], - messages=None, request_kwargs=data, + model=selected["model_name"], + healthy_deployments=router.get_model_list(selected["model_name"]) or [], + messages=None, + request_kwargs=data, ) assert [deployment["model_info"]["id"] for deployment in filtered] == [expected_deployment] assert selected["model_info"]["id"] == expected_deployment await affinity.async_pre_call_deployment_hook( kwargs={ - **data, "model_info": selected["model_info"], + **data, + "model_info": selected["model_info"], bucket: {**metadata, "deployment_model_name": selected["model_name"]}, }, call_type=None, @@ -7845,13 +7868,29 @@ class TestChildSessionTierPins: (None, None, "SIMPLE", "heuristic_scorer", "caller-a", True), ): await self._route( - router, surface, session, prompt, tier, cause, caller, child, + router, + surface, + session, + prompt, + tier, + cause, + caller, + child, expected_deployment=parent if tier == "REASONING" and caller == "caller-a" else None, ) - assert await self._route( - router, surface, "parent-session", None, "REASONING", "user_turn_continuation", - child=False, expected_deployment=parent, - ) == parent + assert ( + await self._route( + router, + surface, + "parent-session", + None, + "REASONING", + "user_turn_continuation", + child=False, + expected_deployment=parent, + ) + == parent + ) finally: router.discard() @@ -7862,10 +7901,18 @@ class TestChildSessionTierPins: parent: Final = await self._route( router, "chat", "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False ) - assert await self._route( - router, "chat", "child-session-a", "Hello!", "REASONING", "session_affinity_pin", - expected_deployment=parent, - ) == parent + assert ( + await self._route( + router, + "chat", + "child-session-a", + "Hello!", + "REASONING", + "session_affinity_pin", + expected_deployment=parent, + ) + == parent + ) finally: router.discard() @@ -8804,7 +8851,7 @@ class TestEscalationKeywords: one step higher so a user can force a stronger model when unhappy with results.""" @staticmethod - def _request_kwargs(session_id: str) -> Dict: + def _request_kwargs(session_id: str) -> dict: return {"metadata": {"session_id": session_id}} def test_default_escalation_keyword(self, complexity_router): @@ -9140,7 +9187,7 @@ class TestEscalationKeywords: assert result.model == "o1-b" # unchanged: no random hop to o1-a / o1-c -def _stalled_tool_history(repeats: int = 3) -> List[Dict]: +def _stalled_tool_history(repeats: int = 3) -> list[dict]: """`repeats` identical bash tool calls in a row, the automatic counterpart to a user typing an escalation keyword: the assistant, not the human, is the one stuck.""" return [ @@ -9692,7 +9739,7 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: # Mirror function_setup: it copies `litellm_metadata` by value into # litellm_params AFTER the router hook has run, so the copy must carry # the decision. Reading the stash any earlier would lose it. - litellm_params: Dict = {} + litellm_params: dict = {} if "metadata" in request_kwargs: litellm_params["metadata"] = request_kwargs["metadata"] if isinstance(request_kwargs.get("litellm_metadata"), dict): @@ -9708,11 +9755,15 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: assert persisted is not None, f"routing_decision dropped for {expected_bucket}" assert persisted["router_model_name"] == "smart-router" if classifier_type == "heuristic_v2": - assert persisted["heuristic_v2_forecast"] == request_kwargs[expected_bucket]["routing_decision"][ - "heuristic_v2_forecast" - ] + assert ( + persisted["heuristic_v2_forecast"] + == request_kwargs[expected_bucket]["routing_decision"]["heuristic_v2_forecast"] + ) assert set(persisted["heuristic_v2_forecast"]["probabilities"]) == { - "SIMPLE", "MEDIUM", "COMPLEX", "REASONING" + "SIMPLE", + "MEDIUM", + "COMPLEX", + "REASONING", } else: assert "heuristic_v2_forecast" not in persisted @@ -9747,7 +9798,7 @@ class TestRoutingDecisionIsPerAttempt: @pytest.mark.asyncio async def test_fallback_to_plain_model_group_clears_the_earlier_decision(self, seed, bucket): router = Router(model_list=self.MODEL_LIST) - request_kwargs: Dict = dict(seed) + request_kwargs: dict = dict(seed) messages = [{"role": "user", "content": "Hello!"}] await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs, messages=messages) @@ -9767,7 +9818,7 @@ class TestRoutingDecisionIsPerAttempt: Skipping the write there would drop provenance on a successfully routed request with no error, so the shared bucket owner replaces the value.""" router = Router(model_list=self.MODEL_LIST) - request_kwargs: Dict = {"litellm_metadata": unusable_bucket} + request_kwargs: dict = {"litellm_metadata": unusable_bucket} response = await router.async_pre_routing_hook( model="smart-router", @@ -9788,7 +9839,7 @@ class TestRecordRoutingDecision: DECISION = {"router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini"} def test_none_clears_a_previous_decision_from_both_buckets(self): - request_kwargs: Dict = { + request_kwargs: dict = { "metadata": {"routing_decision": self.DECISION, "keep": 1}, "litellm_metadata": {"routing_decision": self.DECISION}, } @@ -9798,7 +9849,7 @@ class TestRecordRoutingDecision: assert request_kwargs["metadata"]["keep"] == 1 def test_none_creates_no_bucket_on_a_request_that_had_none(self): - request_kwargs: Dict = {} + request_kwargs: dict = {} Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) assert request_kwargs == {} @@ -9957,7 +10008,7 @@ class TestRedactedLoggingDropsPromptText: MESSAGES = [{"role": "user", "content": "LITELLM ESCALATE please deploy to k8s now"}] - async def _decision(self, request_kwargs: Dict) -> Dict: + async def _decision(self, request_kwargs: dict) -> dict: router = Router(model_list=self.MODEL_LIST) response = await router.async_pre_routing_hook( model="smart-router", request_kwargs=request_kwargs, messages=self.MESSAGES @@ -10026,7 +10077,7 @@ class TestRedactedLoggingDropsPromptText: @pytest.mark.asyncio async def test_redaction_via_request_header_is_honored(self): - request_kwargs: Dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}} + request_kwargs: dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}} decision = await self._decision(request_kwargs) assert "matched_keyword" not in decision assert decision["cause"] == "literal_keyword_match" @@ -10042,8 +10093,10 @@ class TestRedactedLoggingDropsPromptText: { "classifier_type": "capability", "capability_classifier_config": { - "efficient_tier": "SIMPLE", "capable_tier": "REASONING", - "base_threshold": 0.5, "threshold_step": 0.1, + "efficient_tier": "SIMPLE", + "capable_tier": "REASONING", + "base_threshold": 0.5, + "threshold_step": 0.1, }, }, "capability_classifier_fallback", @@ -10052,8 +10105,10 @@ class TestRedactedLoggingDropsPromptText: { "classifier_type": "llm_v2", "llm_v2_config": { - "efficient_profile": "Routine tasks", "capable_profile": "Complex tasks", - "harness": "Read files and run checks", "max_quality_gap": 0.05, + "efficient_profile": "Routine tasks", + "capable_profile": "Complex tasks", + "harness": "Read files and run checks", + "max_quality_gap": 0.05, }, }, "llm_v2_fallback", @@ -10064,36 +10119,46 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction( overrides: Mapping[str, object], expected_cause: str, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - router: Final = Router(model_list=[ - { - "model_name": "auto", - "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "cheap", "REASONING": "capable"}, - "default_model": "cheap", - "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000}, - **overrides, + router: Final = Router( + model_list=[ + { + "model_name": "auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "cheap", "REASONING": "capable"}, + "default_model": "cheap", + "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000}, + **overrides, + }, }, }, - }, - { - "model_name": "classifier", - "litellm_params": { - "model": "openai/test-classifier", "api_key": "test-key", - "api_base": "https://classifier.test/v1", + { + "model_name": "classifier", + "litellm_params": { + "model": "openai/test-classifier", + "api_key": "test-key", + "api_base": "https://classifier.test/v1", + }, }, - }, - {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}}, - {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}}, - ]) + {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}}, + {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}}, + ] + ) with respx.mock: upstream: Final = respx.post("https://classifier.test/v1/chat/completions").respond( 200, json={ - "id": "classifier-result", "object": "chat.completion", "model": "test-classifier", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"}, - "finish_reason": "stop"}], + "id": "classifier-result", + "object": "chat.completion", + "model": "test-classifier", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"}, + "finish_reason": "stop", + } + ], }, ) metadata: Final[dict[str, object]] = {"headers": {"x-litellm-enable-message-redaction": True}} @@ -10110,11 +10175,21 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction( assert "PRIVATE_INVALID_REPLY" not in json.dumps(decision) assert "signals" not in decision if overrides["classifier_type"] == "llm": - upstream.respond(200, json={ - "id": "classifier-recovered", "object": "chat.completion", "model": "test-classifier", - "choices": [{"index": 0, "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'}, - "finish_reason": "stop"}], - }) + upstream.respond( + 200, + json={ + "id": "classifier-recovered", + "object": "chat.completion", + "model": "test-classifier", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'}, + "finish_reason": "stop", + } + ], + }, + ) recovered: Final = await router.async_pre_routing_hook( model="auto", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] ) @@ -12528,7 +12603,7 @@ class TestClassifierFallbackChoice: }, ) mock_router_instance.cache = DualCache() - request_kwargs: Dict = {"metadata": {"session_id": "session-flaky"}} + request_kwargs: dict = {"metadata": {"session_id": "session-flaky"}} mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) first = await router.async_pre_routing_hook( @@ -12567,7 +12642,7 @@ class TestClassifierFallbackChoice: }, ) mock_router_instance.cache = DualCache() - request_kwargs: Dict = {"metadata": {"session_id": "session-steady"}} + request_kwargs: dict = {"metadata": {"session_id": "session-steady"}} mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) first = await router.async_pre_routing_hook( @@ -13045,7 +13120,7 @@ class TestClassificationRubrics: assert config.classifier_llm_config.system_prompt == "Grade the data sensitivity of the request." -def _custom_tier_config(**overrides) -> Dict: +def _custom_tier_config(**overrides) -> dict: """A valid operator-defined tier set: two built-in names plus one custom tier.""" return { "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-20250514", "SECURITY_REVIEW": "o1-preview"}, @@ -14069,8 +14144,8 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout complexity_router_config=config, ) request_kwargs = {"metadata": {"session_id": "masked-params-session"}} + mock_router_instance.cache = DualCache() if route == "session": - mock_router_instance.cache = DualCache() await mock_router_instance.cache.async_set_cache( key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs), value={"model": "opus", "tier": "REASONING"}, @@ -14816,7 +14891,9 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), - complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}), + complexity_router_config=_tier_config( + tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"} + ), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14923,9 +15000,7 @@ class TestContextWindowEscalation: else MappingProxyType({}) ) raw_config: Final = RequestComplexityRouterConfig.model_validate( - MappingProxyType( - {"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting} - ) + MappingProxyType({"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting}) ) config: Final = ( RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json()) @@ -15086,7 +15161,7 @@ class TestContextWindowEscalation: complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}), ) real_get_llm_provider = litellm.get_llm_provider - copilot_resolutions: List = [] + copilot_resolutions: list = [] def _guarded(*args, **kwargs): target = str(kwargs.get("model") or (args[0] if args else "")) + str( @@ -15171,9 +15246,7 @@ class TestModalityRouting: self.BASE_VISION, ) original: Final = await router.aclassify("What color is this?") - result: Final = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE - ) + result: Final = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE) assert original.heuristic_v2_forecast is not None assert result is not None and result.routing_decision is not None @@ -17367,16 +17440,18 @@ class TestMaxTokensFromTierModel: tier_litellm_params={tier_carrier: 4321}, max_tokens_from_tier_model="clamp", marker_params={marker_carrier: 2000}, - simple_deployments=[{ - "model_name": "small", - "litellm_params": { - "model": "openai/gpt-5.6", - "api_key": "test-only", - "api_base": "https://tier.test/v1", - surface_key: 9000, - }, - "model_info": {"max_output_tokens": 16384}, - }], + simple_deployments=[ + { + "model_name": "small", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-only", + "api_base": "https://tier.test/v1", + surface_key: 9000, + }, + "model_info": {"max_output_tokens": 16384}, + } + ], ) with respx.mock(assert_all_mocked=True) as upstream: @@ -17426,10 +17501,14 @@ class TestMaxTokensFromTierModel: max_tokens_from_tier_model="clamp", marker_params={marker_carrier: 2000}, tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None, - simple_deployments=None if known_ceiling else [{ - "model_name": "small", - "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, - }], + simple_deployments=None + if known_ceiling + else [ + { + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + } + ], ) request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap} await router.async_get_available_deployment( @@ -17437,14 +17516,15 @@ class TestMaxTokensFromTierModel: ) first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} await router.async_get_available_deployment( - model="smart-router", request_kwargs=request, + model="smart-router", + request_kwargs=request, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], ) expected_first: Final = ( {"max_tokens": tier_cap or caller_cap or 8192} - if known_ceiling or tier_cap else - {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} + if known_ceiling or tier_cap + else {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} ) assert first == expected_first assert request["max_tokens"] == (caller_cap or 64000) @@ -17464,10 +17544,12 @@ class TestMaxTokensFromTierModel: router: Final = self._router( max_tokens_from_tier_model=mode, marker_params={marker_carrier: 2000}, - simple_deployments=[{ - "model_name": "small", - "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, - }], + simple_deployments=[ + { + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + } + ], ) sent: Final = await self._routed(router) @@ -17696,7 +17778,11 @@ class TestMaxTokensFromTierModel: ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}), ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}), ( - {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}}, + { + "metadata": { + CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True} + } + }, {"max_output_tokens": 5}, ), ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None), diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index d317f0d5741..deb1bf97538 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -11,7 +11,7 @@ from fakeredis.aioredis import FakeConnection from redis.asyncio import BlockingConnectionPool, Redis import litellm -from litellm.caching.affinity_cache import claim_affinity_pin +from litellm.caching.affinity_cache import claim_affinity_pin_in_memory from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -901,7 +901,6 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): assert second == "our-deployment" -@pytest.mark.asyncio @pytest.mark.parametrize( ("stored", "expected"), [ @@ -914,13 +913,13 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): (None, {"model": "second"}), ], ) -async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: +def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) clock.return_value = 105.0 - winner: Final = await claim_affinity_pin( + winner: Final = claim_affinity_pin_in_memory( cache, "tier-pin", {"model": "second"}, @@ -936,24 +935,21 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored assert cache.in_memory_cache.get_cache("tier-pin") is None -@pytest.mark.asyncio -async def test_concurrent_eligible_claims_return_one_winner() -> None: +def test_eligible_local_claims_keep_first_winner() -> None: cache: Final = DualCache() candidates: Final = ({"model": "first"}, {"model": "second"}) - winners: Final = await asyncio.gather( - *( - claim_affinity_pin( - cache, - "tier-pin", - candidates[index % 2], - 30, - eligible_values=candidates, - ) - for index in range(20) + winners: Final = tuple( + claim_affinity_pin_in_memory( + cache, + "tier-pin", + candidates[index % 2], + 30, + eligible_values=candidates, ) + for index in range(20) ) - assert winners == [{"model": "first"}] * 20 + assert winners == ({"model": "first"},) * 20 assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"}