mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
refactor(redis): migrate complexity tier model affinity into a typed store
This commit is contained in:
Devin AI
parent
1e294f8b3b
commit
1bbd3c2121
6 changed files with 662 additions and 358 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
103
litellm/redis/stores/complexity_tier_model_affinity.py
Normal file
103
litellm/redis/stores/complexity_tier_model_affinity.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
234
tests/unit/redis/stores/test_complexity_tier_model_affinity.py
Normal file
234
tests/unit/redis/stores/test_complexity_tier_model_affinity.py
Normal file
|
|
@ -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"}
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue