From b8f2c34461f5020084710f2574d562a42d1cd4dc Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 28 Jul 2026 15:07:57 -0700 Subject: [PATCH] feat(complexity_router): capture session payloads for prompt-cache warming Adds the inert capture plane for multi-model cache warming: a CacheWarmingConfig block on the complexity router, a Redis-backed per-session payload store, and a CustomLogger capture hook at async_pre_call_deployment_hook that snapshots the whitelisted request payload (messages/system/tools/tool_choice) per session for warming-enabled auto-routers. Default-off; nothing consumes the records yet --- .../cache_warming/__init__.py | 37 +++ .../cache_warming/capture_hook.py | 250 ++++++++++++++ .../cache_warming/eligibility.py | 22 ++ .../complexity_router/cache_warming/store.py | 215 ++++++++++++ .../complexity_router/cache_warming/types.py | 60 ++++ .../complexity_router/complexity_router.py | 73 ++-- .../complexity_router/config.py | 62 ++++ .../complexity_router/request_metadata.py | 25 ++ .../complexity_router/__init__.py | 0 .../cache_warming/__init__.py | 0 .../cache_warming/test_capture_hook.py | 314 ++++++++++++++++++ .../cache_warming/test_store.py | 240 +++++++++++++ .../cache_warming/test_types.py | 77 +++++ .../router_strategy/test_complexity_router.py | 184 ++++++++++ 14 files changed, 1530 insertions(+), 29 deletions(-) create mode 100644 litellm/router_strategy/complexity_router/cache_warming/__init__.py create mode 100644 litellm/router_strategy/complexity_router/cache_warming/capture_hook.py create mode 100644 litellm/router_strategy/complexity_router/cache_warming/eligibility.py create mode 100644 litellm/router_strategy/complexity_router/cache_warming/store.py create mode 100644 litellm/router_strategy/complexity_router/cache_warming/types.py create mode 100644 litellm/router_strategy/complexity_router/request_metadata.py create mode 100644 tests/test_litellm/router_strategy/complexity_router/__init__.py create mode 100644 tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py create mode 100644 tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture_hook.py create mode 100644 tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py create mode 100644 tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py diff --git a/litellm/router_strategy/complexity_router/cache_warming/__init__.py b/litellm/router_strategy/complexity_router/cache_warming/__init__.py new file mode 100644 index 00000000000..cb28800f2c3 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/__init__.py @@ -0,0 +1,37 @@ +from litellm.router_strategy.complexity_router.cache_warming.capture_hook import ( + ComplexityCacheWarmingCaptureHook, +) +from litellm.router_strategy.complexity_router.cache_warming.eligibility import ( + min_prompt_cache_tokens_for_warm_set, + resolve_warm_models, +) +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_MARKER_KEY, + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, + WARM_FRESHNESS_SLACK_SECONDS, + CacheWarmingAttribution, + CacheWarmingPayload, + CacheWarmingRecord, + compress_payload, + decompress_payload, +) + +__all__ = [ + "CacheWarmingStore", + "ComplexityCacheWarmingCaptureHook", + "min_prompt_cache_tokens_for_warm_set", + "resolve_warm_models", + "CACHE_WARMING_MARKER_KEY", + "CACHE_WARMING_RECORD_SCHEMA_VERSION", + "CACHE_WARMING_REPLAY_MARKER_KEY", + "CACHE_WARMING_REPLAY_TAG", + "WARM_FRESHNESS_SLACK_SECONDS", + "CacheWarmingAttribution", + "CacheWarmingPayload", + "CacheWarmingRecord", + "compress_payload", + "decompress_payload", +] diff --git a/litellm/router_strategy/complexity_router/cache_warming/capture_hook.py b/litellm/router_strategy/complexity_router/cache_warming/capture_hook.py new file mode 100644 index 00000000000..f8384b94336 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/capture_hook.py @@ -0,0 +1,250 @@ +import asyncio +import uuid +import weakref +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Callable, Literal + +from pydantic import TypeAdapter + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.complexity_router.cache_warming.eligibility import ( + min_prompt_cache_tokens_for_warm_set, + resolve_warm_models, +) +from litellm.router_strategy.complexity_router.request_metadata import ( + get_session_id_from_request_kwargs, + get_user_api_key_hash_from_request_kwargs, + iter_metadata_dicts, +) +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_MARKER_KEY, + CACHE_WARMING_REPLAY_MARKER_KEY, + CACHE_WARMING_REPLAY_TAG, + CacheWarmingAttribution, + CacheWarmingPayload, + compress_payload, +) +from litellm.types.utils import CallTypes + +if TYPE_CHECKING: + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + +_CAPTURE_CALL_TYPES = frozenset( + (CallTypes.completion, CallTypes.acompletion, CallTypes.anthropic_messages, CallTypes.aanthropic_messages) +) +_ANTHROPIC_CALL_TYPES = frozenset((CallTypes.anthropic_messages, CallTypes.aanthropic_messages)) +_MAX_UNCOMPRESSED_RATIO = 8 +_ATTRIBUTION_KEYS = ( + "user_api_key", + "user_api_key_hash", + "user_api_key_user_id", + "user_api_key_team_id", + "user_api_key_end_user_id", +) + + +@lru_cache(maxsize=64) +def _warn_privacy_gate_blocked(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming is enabled for auto-router %s but prompt retention is not permitted " + "(store_prompts_in_spend_logs is off or message redaction is active); capture is skipped", + auto_router_model_name, + ) + + +@lru_cache(maxsize=4096) +def _warn_payload_too_large(auto_router_model_name: str, session_id: str) -> None: + verbose_router_logger.warning( + "cache_warming: session %s on auto-router %s exceeds max_payload_bytes; not warming this session", + session_id, + auto_router_model_name, + ) + + +def _capture_allowed(kwargs: Mapping[str, object]) -> bool: + """One consent predicate for one question: may this request's prompt content be + retained? Honors both halves of the operator's stated policy: the redaction + opt-out (turn_off_message_logging, including per-request and header forms) and + the prompt-retention opt-in (store_prompts_in_spend_logs; SDK use without the + proxy consents through cache_warming.enabled itself).""" + from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + if should_redact_message_logging({"litellm_params": kwargs}): # mutable-ok: read-only view for the predicate + return False + try: + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + _should_store_prompts_and_responses_in_spend_logs, # pyright: ignore[reportPrivateUsage] # canonical proxy consent gate; no public accessor exists + ) + except ImportError: + return True + return _should_store_prompts_and_responses_in_spend_logs() + + +def _resolve_stamp(metadata_dicts: Sequence[Mapping[str, object]]) -> "tuple[ComplexityRouter, str] | None": + marker = next( + (value for metadata in metadata_dicts if isinstance(value := metadata.get(CACHE_WARMING_MARKER_KEY), dict)), + None, + ) + if marker is None: + return None + typed = TypeAdapter(dict[str, object]).validate_python(marker) + ref = typed.get("strategy_ref") + routed_model = typed.get("routed_model") + if not isinstance(ref, str) or not isinstance(routed_model, str): + return None + strategy = _WARMING_STRATEGIES.get(ref) + if strategy is None: + verbose_router_logger.debug( + "cache_warming: stamped strategy is no longer registered (config reload between routing and " + "capture); skipping capture for this in-flight turn" + ) + return None + if typed.get("auto_router_model_name") != strategy.model_name: + return None + return (strategy, routed_model) + + +def _is_replay(metadata_dicts: Sequence[Mapping[str, object]]) -> bool: + for metadata in metadata_dicts: + if metadata.get(CACHE_WARMING_REPLAY_MARKER_KEY): + return True + tags = metadata.get("tags") + if isinstance(tags, list) and CACHE_WARMING_REPLAY_TAG in tags: + return True + return False + + +def _gate_and_compress( + payload: CacheWarmingPayload, max_payload_bytes: int, min_tokens: int +) -> "tuple[str, str, int] | Literal['too_large', 'too_small']": + """Size gates plus compression, run off the event loop. The uncompressed bound + runs before any compression so adversarial highly-compressible content cannot + buy unbounded CPU with a small compressed result, and the chars/4 token + estimate deliberately avoids running a real tokenizer over a full multi-turn + conversation in the request path.""" + serialized_chars = len(payload.model_dump_json()) + if serialized_chars > _MAX_UNCOMPRESSED_RATIO * max_payload_bytes: + return "too_large" + token_estimate = serialized_chars // 4 + if token_estimate < min_tokens: + return "too_small" + blob, sha = compress_payload(payload) + if len(blob) > max_payload_bytes: + return "too_large" + return (blob, sha, token_estimate) + + +def _extract_attribution(metadata_dicts: Sequence[Mapping[str, object]]) -> CacheWarmingAttribution: + merged = dict( # mutable-ok: handed straight to pydantic, never retained + (key, str(value)) + for metadata in reversed(metadata_dicts) + for key, value in metadata.items() + if key in _ATTRIBUTION_KEYS and value is not None + ) + return CacheWarmingAttribution(**merged) + + +_WARMING_STRATEGIES: "weakref.WeakValueDictionary[str, ComplexityRouter]" = weakref.WeakValueDictionary() + + +def register_warming_strategy(strategy: "ComplexityRouter") -> str: + """Give a warming-enabled complexity router a per-request-resolvable identity. + + Capture is dispatched through one process-wide stateless hook; each request's + pre-routing stamp carries the UUID returned here and capture resolves it back + to the exact strategy object that routed the request. Replaced strategies + simply fall out of the weak registry, so there is no hook lifecycle to sync + on set_model_list or upsert_deployment and no per-Router hook instances to + dedupe or remove.""" + import litellm + + ref = uuid.uuid4().hex + _WARMING_STRATEGIES[ref] = strategy + litellm.logging_callback_manager.add_litellm_callback(_get_dispatcher()) + return ref + + +@lru_cache(maxsize=1) +def _get_dispatcher() -> "ComplexityCacheWarmingCaptureHook": + return ComplexityCacheWarmingCaptureHook() + + +class ComplexityCacheWarmingCaptureHook(CustomLogger): + def __init__(self, privacy_gate: "Callable[[Mapping[str, object]], bool]" = _capture_allowed) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger.__init__ is legacy-untyped + self.privacy_gate = privacy_gate + + async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: CallTypes | None) -> None: + try: + await self._capture(kwargs, call_type) + except Exception: # noqa: BLE001 # a capture failure must never fail the user's request + verbose_router_logger.exception("cache_warming capture failed; the request continues unaffected") + + async def _capture(self, kwargs: Mapping[str, object], call_type: CallTypes | None) -> None: + if call_type not in _CAPTURE_CALL_TYPES: + return + metadata_dicts = iter_metadata_dicts(kwargs) + if _is_replay(metadata_dicts): + return + resolved = _resolve_stamp(metadata_dicts) + if resolved is None: + return + strategy, routed_model = resolved + if not self.privacy_gate(kwargs): + _warn_privacy_gate_blocked(strategy.model_name) + return + session_id = get_session_id_from_request_kwargs(kwargs) + if session_id is None: + return + payload = self._build_payload(kwargs, call_type, routed_model) + if payload is None: + return + config = strategy.config.cache_warming + warm_models = resolve_warm_models(strategy.config) + gated = await asyncio.to_thread( + _gate_and_compress, payload, config.max_payload_bytes, min_prompt_cache_tokens_for_warm_set(warm_models) + ) + match gated: + case "too_large": + _warn_payload_too_large(strategy.model_name, session_id) + return + case "too_small": + return + case (blob, sha, token_estimate): + pass + store = strategy.get_cache_warming_store() + if store is None: + return + caller_scope = get_user_api_key_hash_from_request_kwargs(kwargs) or "unscoped" + await store.upsert_session( + caller_scope=caller_scope, + session_id=session_id, + payload_compressed=blob, + payload_sha256=sha, + token_estimate=token_estimate, + served_model=routed_model, + attribution=_extract_attribution(metadata_dicts), + ttl_seconds=config.session_ttl_seconds, + max_sessions=config.max_sessions, + ) + + @staticmethod + def _build_payload( + kwargs: Mapping[str, object], call_type: CallTypes | None, routed_model: str + ) -> CacheWarmingPayload | None: + messages = kwargs.get("messages") + if not isinstance(messages, list) or not messages: + return None + is_anthropic_surface = call_type in _ANTHROPIC_CALL_TYPES + return CacheWarmingPayload.model_validate( + { # mutable-ok: pydantic input, never retained + "model": routed_model, + "messages": messages, + "system": kwargs.get("system") if is_anthropic_surface else None, + "tools": kwargs.get("tools"), + "tool_choice": kwargs.get("tool_choice"), + "call_surface": "anthropic_messages" if is_anthropic_surface else "chat_completions", + } + ) diff --git a/litellm/router_strategy/complexity_router/cache_warming/eligibility.py b/litellm/router_strategy/complexity_router/cache_warming/eligibility.py new file mode 100644 index 00000000000..22d0cb300a5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/eligibility.py @@ -0,0 +1,22 @@ +from typing import TYPE_CHECKING + +from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT + +if TYPE_CHECKING: + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + +def resolve_warm_models(config: "ComplexityRouterConfig") -> tuple[str, ...]: + explicit = config.cache_warming.warm_models + if explicit: + return tuple(dict.fromkeys(explicit)) + first_per_tier = (models if isinstance(models, str) else models[0] for models in config.tiers.values() if models) + return tuple(dict.fromkeys(first_per_tier)) + + +def min_prompt_cache_tokens_for_warm_set(warm_models: tuple[str, ...]) -> int: + from litellm.utils import get_prompt_cache_min_tokens + + if not warm_models: + return DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT + return min(get_prompt_cache_min_tokens(model) for model in warm_models) diff --git a/litellm/router_strategy/complexity_router/cache_warming/store.py b/litellm/router_strategy/complexity_router/cache_warming/store.py new file mode 100644 index 00000000000..c98032d14cc --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/store.py @@ -0,0 +1,215 @@ +import time +from collections.abc import Awaitable, Mapping +from functools import lru_cache +from typing import Callable + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_router_logger +from litellm.caching.redis_cache import RedisCache +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingRecord, +) + +_WARMTH_KEY_PREFIX = "complexity_router_cache_warmth:v1" + +_CAPTURE_SCRIPT = """ +local sessions_key = KEYS[1] +local index_key = KEYS[2] +local member = ARGV[1] +local record_json = ARGV[2] +local now = tonumber(ARGV[3]) +local expires_at = tonumber(ARGV[4]) +local max_sessions = tonumber(ARGV[5]) +local expired = redis.call('ZRANGEBYSCORE', index_key, 0, now) +if #expired > 0 then + redis.call('HDEL', sessions_key, unpack(expired)) + redis.call('ZREMRANGEBYSCORE', index_key, 0, now) +end +if not redis.call('ZSCORE', index_key, member) and redis.call('ZCARD', index_key) >= max_sessions then + return 0 +end +redis.call('HSET', sessions_key, member, record_json) +redis.call('ZADD', index_key, expires_at, member) +redis.call('EXPIREAT', sessions_key, math.ceil(expires_at)) +redis.call('EXPIREAT', index_key, math.ceil(expires_at)) +return 1 +""" + +_LIST_LIVE_SESSIONS_SCRIPT = """ +local index_key = KEYS[1] +local now = tonumber(ARGV[1]) +local limit = tonumber(ARGV[2]) +return redis.call('ZRANGEBYSCORE', index_key, '(' .. now, '+inf', 'LIMIT', 0, limit) +""" + +_GET_RECORD_SCRIPT = """ +return redis.call('HGET', KEYS[1], ARGV[1]) +""" + +_MEMBERS_ADAPTER: TypeAdapter[tuple[str | bytes, ...]] = TypeAdapter(tuple[str | bytes, ...]) + + +@lru_cache(maxsize=64) +def _warn_redis_missing(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming is enabled for auto-router %s but the router cache has no Redis; " + "cache warming is inactive until Redis is configured", + auto_router_model_name, + ) + + +@lru_cache(maxsize=64) +def _warn_session_cap_reached(auto_router_model_name: str) -> None: + verbose_router_logger.warning( + "cache_warming: auto-router %s reached max_sessions; new sessions are not captured until " + "existing records expire or go idle", + auto_router_model_name, + ) + + +def _parse_record(raw: object) -> CacheWarmingRecord | None: + if not isinstance(raw, (str, bytes)): + return None + try: + record = CacheWarmingRecord.model_validate_json(raw) + except ValidationError: + return None + if record.schema_version != CACHE_WARMING_RECORD_SCHEMA_VERSION: + return None + return record + + +def _parse_warmth(raw: object) -> float | None: + if isinstance(raw, (int, float)): + return float(raw) + if isinstance(raw, (str, bytes)): + try: + return float(raw) + except ValueError: + return None + return None + + +class CacheWarmingStore: + """Per-session capture state for provider prompt-cache warming. + + Session records and their expiry index live in two hash-tagged keys on one + Redis Cluster slot, so every capture is a single atomic Lua operation that + prunes expired sessions, enforces the max_sessions cap exactly, and writes + the record all-or-nothing. No partial write states exist, so no + compensation or fencing is needed anywhere, and the cap bounds the slot's + footprint by construction. Warmth stamps are plain single-writer keys with + their own TTL. Records are last-writer-wins by design: the latest turn is + the correct replay payload. A script fault raises and fails closed, never + an empty result mistaken for capacity.""" + + def __init__(self, redis_cache: RedisCache | None, auto_router_model_name: str) -> None: + self.redis_cache = redis_cache + self.auto_router_model_name = auto_router_model_name + register = redis_cache.async_register_script if redis_cache is not None else None + self._capture: Callable[..., Awaitable[object]] | None = register(_CAPTURE_SCRIPT) if register else None + self._list_live: Callable[..., Awaitable[object]] | None = ( + register(_LIST_LIVE_SESSIONS_SCRIPT) if register else None + ) + self._get: Callable[..., Awaitable[object]] | None = register(_GET_RECORD_SCRIPT) if register else None + + @staticmethod + def record_key(auto_router_model_name: str, caller_scope: str, session_id: str) -> str: + return f"{caller_scope}:{session_id}" + + @staticmethod + def warmth_key(record_key: str, model_group: str) -> str: + return f"{_WARMTH_KEY_PREFIX}:{record_key}:{model_group}" + + def sessions_key(self) -> str: + return f"{{cache_warm:v1:{self.auto_router_model_name}}}:sessions" + + def index_key(self) -> str: + return f"{{cache_warm:v1:{self.auto_router_model_name}}}:index" + + def _require_redis(self) -> RedisCache | None: + if self.redis_cache is None: + _warn_redis_missing(self.auto_router_model_name) + return None + return self.redis_cache + + async def get_record(self, key: str) -> CacheWarmingRecord | None: + if self._require_redis() is None or self._get is None: + return None + raw = await self._get(keys=[self.sessions_key()], args=[key]) + return _parse_record(raw) + + async def upsert_session( + self, + *, + caller_scope: str, + session_id: str, + payload_compressed: str, + payload_sha256: str, + token_estimate: int, + served_model: str, + attribution: CacheWarmingAttribution, + ttl_seconds: int, + max_sessions: int, + ) -> None: + redis_cache = self._require_redis() + if redis_cache is None or self._capture is None: + return + key = self.record_key(self.auto_router_model_name, caller_scope, session_id) + now = time.time() + record = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed=payload_compressed, + payload_sha256=payload_sha256, + token_estimate=token_estimate, + last_activity=now, + served_model=served_model, + attribution=attribution, + auto_router_model_name=self.auto_router_model_name, + ) + try: + admitted = await self._capture( + keys=[self.sessions_key(), self.index_key()], + args=[key, record.model_dump_json(), now, now + ttl_seconds, max_sessions], + ) + except Exception: # noqa: BLE001 # a capture fault fails closed: no capture beats an uncapped write + verbose_router_logger.warning("cache_warming capture script failed; skipping capture", exc_info=True) + return + if admitted != 1: + _warn_session_cap_reached(self.auto_router_model_name) + return + await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds) + + async def mark_warm_attempt(self, key: str, model_group: str, attempted_at: float, ttl_seconds: int) -> None: + redis_cache = self._require_redis() + if redis_cache is None: + return + await redis_cache.async_set_cache( # pyright: ignore[reportUnknownMemberType] # RedisCache is legacy-untyped + key=self.warmth_key(key, model_group), value=attempted_at, ttl=ttl_seconds + ) + + async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, float]: + redis_cache = self._require_redis() + if redis_cache is None: + return {} # mutable-ok: fresh per-call result, not shared state + return { # mutable-ok: fresh per-call result, not shared state + model_group: stamp + for model_group in model_groups + if ( + stamp := _parse_warmth( + await redis_cache.async_get_cache(self.warmth_key(key, model_group)) # pyright: ignore[reportUnknownMemberType] # RedisCache is legacy-untyped + ) + ) + is not None + } + + async def list_session_keys(self, max_sessions: int) -> tuple[str, ...]: + if self._require_redis() is None or self._list_live is None: + return () + members = _MEMBERS_ADAPTER.validate_python( + await self._list_live(keys=[self.index_key()], args=[time.time(), max_sessions]) + ) + return tuple(member.decode("utf-8") if isinstance(member, bytes) else member for member in members) diff --git a/litellm/router_strategy/complexity_router/cache_warming/types.py b/litellm/router_strategy/complexity_router/cache_warming/types.py new file mode 100644 index 00000000000..970f46c1953 --- /dev/null +++ b/litellm/router_strategy/complexity_router/cache_warming/types.py @@ -0,0 +1,60 @@ +import base64 +import hashlib +import json +import zlib +from collections.abc import Mapping +from typing import Literal + +from pydantic import BaseModel, ConfigDict + +CACHE_WARMING_MARKER_KEY = "_complexity_router_cache_warming" +CACHE_WARMING_REPLAY_MARKER_KEY = "litellm_cache_warming" +CACHE_WARMING_REPLAY_TAG = "litellm_cache_warming" +CACHE_WARMING_RECORD_SCHEMA_VERSION = 1 +WARM_FRESHNESS_SLACK_SECONDS = 60 + + +class CacheWarmingPayload(BaseModel): + model_config = ConfigDict(extra="forbid") + + model: str + messages: tuple[Mapping[str, object], ...] + system: str | tuple[Mapping[str, object], ...] | None = None + tools: tuple[Mapping[str, object], ...] | None = None + tool_choice: str | Mapping[str, object] | None = None + call_surface: Literal["chat_completions", "anthropic_messages"] + + +class CacheWarmingAttribution(BaseModel): + model_config = ConfigDict(extra="forbid") + + user_api_key: str | None = None + user_api_key_hash: str | None = None + user_api_key_user_id: str | None = None + user_api_key_team_id: str | None = None + user_api_key_end_user_id: str | None = None + + +class CacheWarmingRecord(BaseModel): + model_config = ConfigDict(extra="forbid") + + schema_version: int + payload_compressed: str + payload_sha256: str + token_estimate: int + last_activity: float + served_model: str + attribution: CacheWarmingAttribution + auto_router_model_name: str + + +def compress_payload(payload: CacheWarmingPayload) -> tuple[str, str]: + raw = payload.model_dump_json().encode("utf-8") + blob = base64.b64encode(zlib.compress(raw)).decode("ascii") + sha = hashlib.sha256(json.dumps(payload.model_dump(), sort_keys=True).encode("utf-8")).hexdigest() + return blob, sha + + +def decompress_payload(blob_b64: str) -> CacheWarmingPayload: + raw = zlib.decompress(base64.b64decode(blob_b64.encode("ascii"))) + return CacheWarmingPayload.model_validate_json(raw) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index e5268b5107b..47a81a3816f 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -25,6 +25,10 @@ from pydantic import BaseModel from litellm._logging import verbose_router_logger from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.complexity_router.request_metadata import ( + get_session_id_from_request_kwargs, + get_user_api_key_hash_from_request_kwargs, +) from litellm.types.utils import ModelResponse from .config import ( @@ -43,6 +47,7 @@ if TYPE_CHECKING: from litellm.router import Router from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter + from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore from litellm.types.router import PreRoutingHookResponse else: Router = Any @@ -163,6 +168,14 @@ class ComplexityRouter(CustomLogger): else: self.config = ComplexityRouterConfig() + self._cache_warming_ref: str | None = None + if self.config.cache_warming.enabled: + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import ( + register_warming_strategy, + ) + + self._cache_warming_ref = register_warming_strategy(self) + # Override default_model if provided if default_model: self.config.default_model = default_model @@ -897,42 +910,41 @@ class ComplexityRouter(CustomLogger): return user_message, system_prompt - @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: - """Metadata may land on `metadata` or `litellm_metadata` depending on the - endpoint, mirroring DeploymentAffinityCheck's precedence.""" - return [ - metadata - for metadata_key in ("litellm_metadata", "metadata") - if isinstance(metadata := request_kwargs.get(metadata_key), dict) - ] + def _stamp_cache_warming_marker( + self, + request_kwargs: dict, # mutable-ok: stamps marker into request metadata + routed_model: str, + ) -> None: + """Always stamps into litellm_metadata: the plain metadata kwarg IS the + provider-body metadata param on the native Anthropic messages surface, so an + internal marker there would be forwarded upstream and rejected.""" + if not self.config.cache_warming.enabled: + return + from litellm.router_strategy.complexity_router.cache_warming.types import CACHE_WARMING_MARKER_KEY - @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: - """Resolve a client-supplied session_id.""" - for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): - session_id = metadata.get("session_id") - if session_id is not None: - return str(session_id) - return None + metadata = request_kwargs.setdefault("litellm_metadata", {}) + if isinstance(metadata, dict): + metadata[CACHE_WARMING_MARKER_KEY] = { # mutable-ok: stored into live request metadata + "auto_router_model_name": self.model_name, + "routed_model": routed_model, + "strategy_ref": self._cache_warming_ref, + } - @staticmethod - def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: - """Resolve the proxy-derived API key hash, the same trust boundary - DeploymentAffinityCheck uses for its own key-based affinity (not the - client-supplied OpenAI `user` param, which isn't authenticated).""" - for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): - user_key = metadata.get("user_api_key_hash") - if user_key is not None: - return str(user_key) - return None + def get_cache_warming_store(self) -> CacheWarmingStore | None: + if not self.config.cache_warming.enabled: + return None + from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore + + router = self.litellm_router_instance + redis_cache = router.cache.redis_cache if router is not None and router.cache is not None else None + return CacheWarmingStore(redis_cache=redis_cache, auto_router_model_name=self.model_name) def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: # Namespace by the caller's API key hash so two different callers reusing the # same client-supplied session_id can't poison each other's routing pin. Falls # back to "unscoped" only when there's no authenticated caller to scope by # (e.g. direct Router usage without the proxy layer). - caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" + caller_scope = get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" async def async_pre_routing_hook( @@ -964,7 +976,7 @@ class ComplexityRouter(CustomLogger): metadata[RETURN_RAW_MODEL_NAME_METADATA_KEY] = True use_session_affinity = self.config.session_affinity and not self.config.plugins - session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None + session_id = get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None if cache_key is not None: @@ -996,6 +1008,7 @@ class ComplexityRouter(CustomLogger): kwargs_metadata = request_kwargs.setdefault("metadata", {}) if isinstance(kwargs_metadata, dict): kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = routed_model + self._stamp_cache_warming_marker(request_kwargs, routed_model) cause = "session_affinity_escalation" if routed_model != pinned_model else "session_affinity_pin" verbose_router_logger.info( f"ComplexityRouter: routing decision cause={cause}, routed_model={routed_model}" @@ -1019,6 +1032,8 @@ class ComplexityRouter(CustomLogger): value=response.model, ttl=self.config.session_affinity_ttl_seconds, ) + if response is not None: + self._stamp_cache_warming_marker(request_kwargs, response.model) return response async def _classify_and_route( diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 7437138fbb7..48471f5f654 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -248,6 +248,54 @@ class ClassifierLLMConfig(BaseModel): ) +class CacheWarmingConfig(BaseModel): + """Configuration for multi-model provider prompt-cache warming.""" + + enabled: bool = Field( + default=False, + description=( + "Capture each session's latest payload and keep provider prompt caches warm on every " + "tier model via background max_tokens=1 replays, so mid-session tier switches stay " + "cache hits. Requires Redis on the router cache and store_prompts_in_spend_logs=True" + ), + ) + refresh_interval_seconds: int = Field( + default=270, + gt=0, + description="Replay cadence per model; keep under the provider's cache TTL (Anthropic: 5 minutes)", + ) + session_ttl_seconds: int = Field( + default=3600, + gt=0, + description="TTL for the stored session payload record in Redis", + ) + idle_timeout_seconds: int = Field( + default=600, + gt=0, + description="Stop warming a session after this long without a real request; resumes on the next turn", + ) + max_sessions: int = Field( + default=1000, + gt=0, + description=( + "Exact upper bound on concurrently warmed sessions per auto-router, enforced by an atomic Redis " + "admission index; new sessions past the cap are not captured until existing ones expire" + ), + ) + max_payload_bytes: int = Field( + default=1_048_576, + gt=0, + description="Skip capturing sessions whose compressed payload exceeds this size", + ) + warm_models: tuple[str, ...] | None = Field( + default=None, + description=( + "Explicit model groups to keep warm; defaults to the first member of each tier pool. " + "Only Anthropic/Bedrock models that support prompt caching are warmed" + ), + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -399,6 +447,11 @@ class ComplexityRouterConfig(BaseModel): description="TTL for the session affinity pin; refreshed on every cache hit", ) + cache_warming: CacheWarmingConfig = Field( + default_factory=CacheWarmingConfig, + description="Multi-model provider prompt-cache warming; disabled by default", + ) + plugins: list[RoutingPlugin] | None = Field( default=None, description="RoutingPlugin instances that narrow the classified tier's candidate models before selection", @@ -466,6 +519,15 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_cache_warming_adaptive_combo(self) -> "ComplexityRouterConfig": + if self.cache_warming.enabled and self.adaptive: + raise ValueError( + "cache_warming and adaptive=True cannot both be set: adaptive's bandit selection doesn't yet " + "consult warm-cache state. Disable adaptive or disable cache_warming." + ) + return self + # Combined default config DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig() diff --git a/litellm/router_strategy/complexity_router/request_metadata.py b/litellm/router_strategy/complexity_router/request_metadata.py new file mode 100644 index 00000000000..5b4bdbaba40 --- /dev/null +++ b/litellm/router_strategy/complexity_router/request_metadata.py @@ -0,0 +1,25 @@ +from collections.abc import Mapping + + +def iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return tuple( + metadata + for metadata_key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(metadata_key), dict) + ) + + +def get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: + for metadata in iter_metadata_dicts(request_kwargs): + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + return None + + +def get_user_api_key_hash_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: + for metadata in iter_metadata_dicts(request_kwargs): + user_key = metadata.get("user_api_key_hash") + if user_key is not None: + return str(user_key) + return None diff --git a/tests/test_litellm/router_strategy/complexity_router/__init__.py b/tests/test_litellm/router_strategy/complexity_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture_hook.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture_hook.py new file mode 100644 index 00000000000..88112c4796f --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture_hook.py @@ -0,0 +1,314 @@ +import gc +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from litellm.router_strategy.complexity_router.cache_warming.capture_hook import ( + _WARMING_STRATEGIES, + ComplexityCacheWarmingCaptureHook, + _warn_payload_too_large, + _warn_privacy_gate_blocked, +) +from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_MARKER_KEY, + CACHE_WARMING_REPLAY_MARKER_KEY, + decompress_payload, +) +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.types.utils import CallTypes + +from tests.test_litellm.router_strategy.complexity_router.cache_warming.test_store import FakeRedisCache + +LONG_SYSTEM = "All deployment manifests must declare resource ceilings before rollout. " * 200 + + +def _complexity_router(redis: FakeRedisCache | None, **cache_warming_overrides: object) -> ComplexityRouter: + router_instance = MagicMock() + router_instance.cache = SimpleNamespace(redis_cache=redis) + return ComplexityRouter( + model_name="smart-router", + litellm_router_instance=router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-5-mini", "COMPLEX": "claude-sonnet-4-5"}, + "cache_warming": {"enabled": True, **cache_warming_overrides}, + }, + ) + + +def _hook(allow_privacy: bool = True) -> ComplexityCacheWarmingCaptureHook: + return ComplexityCacheWarmingCaptureHook(privacy_gate=lambda _kwargs: allow_privacy) + + +def _kwargs(router: ComplexityRouter, **overrides: object) -> dict: + base: dict = { + "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "messages": [ + {"role": "system", "content": LONG_SYSTEM}, + {"role": "user", "content": "summarize rule 7"}, + ], + "metadata": { + CACHE_WARMING_MARKER_KEY: { + "auto_router_model_name": "smart-router", + "routed_model": "claude-sonnet-4-5", + "strategy_ref": router._cache_warming_ref, + }, + "session_id": "sess-1", + "user_api_key_hash": "hash-1", + "user_api_key": "hash-1", + "user_api_key_team_id": "team-9", + }, + } + return {**base, **overrides} + + +SESSIONS_KEY = "{cache_warm:v1:smart-router}:sessions" + + +def _stored_records(redis: FakeRedisCache) -> list[dict]: + return [json.loads(value) for value in redis.hashes.get(SESSIONS_KEY, {}).values()] + + +@pytest.mark.asyncio +async def test_captures_whitelisted_fields_only_never_credentials(): + redis = FakeRedisCache() + router = _complexity_router(redis) + result = await _hook().async_pre_call_deployment_hook( + _kwargs(router, api_key="sk-live-secret", litellm_params={"api_key": "sk-live-secret"}), + CallTypes.acompletion, + ) + assert result is None + records = _stored_records(redis) + assert len(records) == 1 + payload = decompress_payload(records[0]["payload_compressed"]) + assert payload.model == "claude-sonnet-4-5" + assert payload.call_surface == "chat_completions" + assert "sk-live-secret" not in json.dumps(records[0]) + + +@pytest.mark.asyncio +async def test_skips_unstamped_request(): + redis = FakeRedisCache() + _complexity_router(redis) + await _hook().async_pre_call_deployment_hook( + {"messages": [{"role": "user", "content": "x"}], "metadata": {"session_id": "s"}}, CallTypes.acompletion + ) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_unknown_strategy_ref(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + kwargs["metadata"][CACHE_WARMING_MARKER_KEY]["strategy_ref"] = "deadbeef" + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_when_stamping_strategy_was_replaced_and_collected(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + ref = router._cache_warming_ref + del router + gc.collect() + assert ref is not None and _WARMING_STRATEGIES.get(ref) is None + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_stamp_whose_name_mismatches_resolved_strategy(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + kwargs["metadata"][CACHE_WARMING_MARKER_KEY]["auto_router_model_name"] = "other" + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slot", ["metadata", "litellm_metadata"]) +async def test_skips_replay_marker_in_either_metadata_slot(slot): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + kwargs[slot] = {**kwargs.pop("metadata"), CACHE_WARMING_REPLAY_MARKER_KEY: True} + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_when_no_session_id(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + kwargs["metadata"].pop("session_id") + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_unsupported_call_type(): + redis = FakeRedisCache() + router = _complexity_router(redis) + await _hook().async_pre_call_deployment_hook(_kwargs(router), CallTypes.aembedding) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_privacy_gate_off_blocks_capture_and_warns(): + _warn_privacy_gate_blocked.cache_clear() + redis = FakeRedisCache() + router = _complexity_router(redis) + await _hook(allow_privacy=False).async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_captures_system_and_surface_for_anthropic_messages(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router, system=LONG_SYSTEM, messages=[{"role": "user", "content": "hi there"}]) + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.aanthropic_messages) + records = _stored_records(redis) + assert len(records) == 1 + payload = decompress_payload(records[0]["payload_compressed"]) + assert payload.call_surface == "anthropic_messages" + assert payload.system == LONG_SYSTEM + + +@pytest.mark.asyncio +async def test_chat_surface_ignores_stray_system_kwarg(): + redis = FakeRedisCache() + router = _complexity_router(redis) + await _hook().async_pre_call_deployment_hook(_kwargs(router, system=LONG_SYSTEM), CallTypes.acompletion) + payload = decompress_payload(_stored_records(redis)[0]["payload_compressed"]) + assert payload.system is None + + +@pytest.mark.asyncio +async def test_skips_below_min_prompt_cache_tokens(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router, messages=[{"role": "user", "content": "tiny"}]) + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_skips_oversized_payload(): + _warn_payload_too_large.cache_clear() + redis = FakeRedisCache() + router = _complexity_router(redis, max_payload_bytes=64) + await _hook().async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_captures_attribution_subset(): + redis = FakeRedisCache() + router = _complexity_router(redis) + await _hook().async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + record = _stored_records(redis)[0] + assert record["attribution"]["user_api_key"] == "hash-1" + assert record["attribution"]["user_api_key_team_id"] == "team-9" + assert record["attribution"]["user_api_key_user_id"] is None + + +@pytest.mark.asyncio +async def test_returns_none_and_never_mutates_kwargs(): + redis = FakeRedisCache() + router = _complexity_router(redis) + kwargs = _kwargs(router) + snapshot = json.dumps(kwargs, sort_keys=True, default=str) + result = await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert result is None + assert json.dumps(kwargs, sort_keys=True, default=str) == snapshot + + +@pytest.mark.asyncio +async def test_swallows_store_exceptions(): + class ExplodingRedis(FakeRedisCache): + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + raise RuntimeError("redis down") + + router = _complexity_router(ExplodingRedis()) + result = await _hook().async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + assert result is None + + +@pytest.mark.asyncio +async def test_same_name_routers_capture_only_via_their_own_strategy(): + redis_a = FakeRedisCache() + redis_b = FakeRedisCache() + router_a = _complexity_router(redis_a) + router_b = _complexity_router(redis_b) + await _hook().async_pre_call_deployment_hook(_kwargs(router_a), CallTypes.acompletion) + assert len(_stored_records(redis_a)) == 1 + assert _stored_records(redis_b) == [] + await _hook().async_pre_call_deployment_hook(_kwargs(router_b), CallTypes.acompletion) + assert len(_stored_records(redis_b)) == 1 + + +@pytest.mark.asyncio +async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth(): + redis = FakeRedisCache() + router = _complexity_router(redis) + hook = _hook() + await hook.async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + key = CacheWarmingStore.record_key("smart-router", "hash-1", "sess-1") + first = json.loads(redis.hashes[SESSIONS_KEY][key]) + redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(123.0) + turn2 = _kwargs(router) + turn2["messages"] = turn2["messages"] + [{"role": "user", "content": "and rule 8?"}] + await hook.async_pre_call_deployment_hook(turn2, CallTypes.acompletion) + second = json.loads(redis.hashes[SESSIONS_KEY][key]) + assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == 123.0 + assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")]) > 0 + assert second["payload_sha256"] != first["payload_sha256"] + + +@pytest.mark.asyncio +async def test_skips_highly_compressible_payload_on_uncompressed_bound(): + _warn_payload_too_large.cache_clear() + redis = FakeRedisCache() + router = _complexity_router(redis, max_payload_bytes=1024) + kwargs = _kwargs(router, messages=[{"role": "user", "content": "a" * 20_000}]) + await _hook().async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_global_message_redaction_blocks_capture(): + import litellm + + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import _capture_allowed + + redis = FakeRedisCache() + router = _complexity_router(redis) + hook = ComplexityCacheWarmingCaptureHook(privacy_gate=_capture_allowed) + previous = litellm.turn_off_message_logging + litellm.turn_off_message_logging = True + try: + await hook.async_pre_call_deployment_hook(_kwargs(router), CallTypes.acompletion) + finally: + litellm.turn_off_message_logging = previous + assert redis.hashes.get(SESSIONS_KEY, {}) == {} + + +@pytest.mark.asyncio +async def test_per_request_redaction_header_blocks_capture(): + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import _capture_allowed + + redis = FakeRedisCache() + router = _complexity_router(redis) + hook = ComplexityCacheWarmingCaptureHook(privacy_gate=_capture_allowed) + kwargs = _kwargs(router) + kwargs["metadata"]["headers"] = {"x-litellm-enable-message-redaction": True} + await hook.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + assert redis.hashes.get(SESSIONS_KEY, {}) == {} diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py new file mode 100644 index 00000000000..cf7f0e50a61 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py @@ -0,0 +1,240 @@ +import json +import logging + +import pytest + +from litellm.router_strategy.complexity_router.cache_warming.store import ( + CacheWarmingStore, + _warn_redis_missing, + _warn_session_cap_reached, +) +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CACHE_WARMING_RECORD_SCHEMA_VERSION, + CacheWarmingAttribution, + CacheWarmingRecord, +) + + +class FakeRedisCache: + def __init__(self, namespace: str | None = None) -> None: + self.namespace = namespace + self.data: dict[str, str] = {} + self.ttls: dict[str, int | None] = {} + self.hashes: dict[str, dict[str, str]] = {} + self.zsets: dict[str, dict[str, float]] = {} + + def _namespaced(self, key: str) -> str: + if self.namespace and not key.startswith(self.namespace): + return f"{self.namespace}:{key}" + return key + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + raw = self.data.get(self._namespaced(key)) + return json.loads(raw) if raw is not None else None + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + namespaced = self._namespaced(key) + self.data[namespaced] = json.dumps(value) + ttl = kwargs.get("ttl") + self.ttls[namespaced] = ttl if isinstance(ttl, int) else None + + def async_register_script(self, script: str): + if "HGET" in script: + + async def get_record(keys: list, args: list) -> str | None: + return self.hashes.get(self._namespaced(keys[0]), {}).get(str(args[0])) + + return get_record + if "HSET" in script: + + async def capture(keys: list, args: list) -> int: + sessions = self.hashes.setdefault(self._namespaced(keys[0]), {}) + index = self.zsets.setdefault(self._namespaced(keys[1]), {}) + member, record_json = str(args[0]), str(args[1]) + now, expires_at, max_sessions = float(args[2]), float(args[3]), int(args[4]) + for stale in [m for m, score in index.items() if score <= now]: + del index[stale] + sessions.pop(stale, None) + if member not in index and len(index) >= max_sessions: + return 0 + sessions[member] = record_json + index[member] = expires_at + return 1 + + return capture + + async def list_live(keys: list, args: list) -> list: + index = self.zsets.get(self._namespaced(keys[0]), {}) + now, limit = float(args[0]), int(args[1]) + live = sorted((score, member) for member, score in index.items() if score > now) + return [member.encode("utf-8") for _, member in live[:limit]] + + return list_live + + +def _record_json(**overrides: object) -> str: + base = CacheWarmingRecord( + schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION, + payload_compressed="blob", + payload_sha256="sha", + token_estimate=2048, + last_activity=1000.0, + served_model="sonnet", + attribution=CacheWarmingAttribution(user_api_key="hashed"), + auto_router_model_name="smart-router", + ).model_dump() + return json.dumps({**base, **overrides}) + + +def _store(redis: FakeRedisCache | None) -> CacheWarmingStore: + return CacheWarmingStore(redis_cache=redis, auto_router_model_name="smart-router") + + +async def _upsert(store: CacheWarmingStore, session_id: str = "s1", max_sessions: int = 100) -> None: + await store.upsert_session( + caller_scope="scope", + session_id=session_id, + payload_compressed="blob2", + payload_sha256="sha2", + token_estimate=4096, + served_model="sonnet", + attribution=CacheWarmingAttribution(), + ttl_seconds=1800, + max_sessions=max_sessions, + ) + + +def test_key_shapes_are_scoped_and_hash_tagged(): + assert CacheWarmingStore.record_key("smart-router", "keyhash", "session-1") == "keyhash:session-1" + assert ( + CacheWarmingStore.warmth_key("keyhash:session-1", "opus") + == "complexity_router_cache_warmth:v1:keyhash:session-1:opus" + ) + store = _store(None) + assert store.sessions_key() == "{cache_warm:v1:smart-router}:sessions" + assert store.index_key() == "{cache_warm:v1:smart-router}:index" + assert store.sessions_key().split("}")[0] == store.index_key().split("}")[0] + + +@pytest.mark.asyncio +async def test_upsert_writes_record_and_stamps_served_model_warmth(): + redis = FakeRedisCache() + store = _store(redis) + await _upsert(store) + key = store.record_key("smart-router", "scope", "s1") + stored = await store.get_record(key) + assert stored is not None + assert stored.payload_compressed == "blob2" + assert stored.last_activity > 0 + warmth = await store.get_warmth(key, ("sonnet", "opus")) + assert set(warmth) == {"sonnet"} + assert redis.ttls[store.warmth_key(key, "sonnet")] == 1800 + + +@pytest.mark.asyncio +async def test_cap_enforced_atomically_with_the_record_write(): + _warn_session_cap_reached.cache_clear() + redis = FakeRedisCache() + store = _store(redis) + await _upsert(store, session_id="s1", max_sessions=2) + await _upsert(store, session_id="s2", max_sessions=2) + await _upsert(store, session_id="s3", max_sessions=2) + assert await store.get_record(store.record_key("smart-router", "scope", "s3")) is None + assert len(await store.list_session_keys(max_sessions=10)) == 2 + + +@pytest.mark.asyncio +async def test_existing_session_updates_even_at_cap(): + redis = FakeRedisCache() + store = _store(redis) + await _upsert(store, session_id="s1", max_sessions=1) + await _upsert(store, session_id="s1", max_sessions=1) + stored = await store.get_record(store.record_key("smart-router", "scope", "s1")) + assert stored is not None and stored.payload_compressed == "blob2" + + +@pytest.mark.asyncio +async def test_expired_session_frees_slot_and_record_together(): + redis = FakeRedisCache() + store = _store(redis) + await _upsert(store, session_id="s1", max_sessions=1) + key1 = store.record_key("smart-router", "scope", "s1") + redis.zsets[store.index_key()][key1] = 1.0 + await _upsert(store, session_id="s2", max_sessions=1) + assert await store.get_record(store.record_key("smart-router", "scope", "s2")) is not None + assert await store.get_record(key1) is None + + +@pytest.mark.asyncio +async def test_capture_fault_fails_closed_not_uncapped(): + redis = FakeRedisCache() + store = _store(redis) + + async def exploding_capture(keys: list, args: list) -> int: + raise RuntimeError("cluster moved slot") + + store._capture = exploding_capture + await _upsert(store) + assert redis.hashes == {} and redis.data == {} + + +@pytest.mark.asyncio +async def test_mark_warm_attempt_never_touches_the_record(): + redis = FakeRedisCache() + store = _store(redis) + key = store.record_key("smart-router", "scope", "s1") + redis.hashes[store.sessions_key()] = {key: _record_json()} + record_before = redis.hashes[store.sessions_key()][key] + await store.mark_warm_attempt(key, "opus", attempted_at=999.0, ttl_seconds=3600) + assert redis.hashes[store.sessions_key()][key] == record_before + assert await store.get_warmth(key, ("opus", "sonnet")) == {"opus": 999.0} + + +@pytest.mark.asyncio +async def test_get_record_returns_none_on_schema_version_mismatch(): + redis = FakeRedisCache() + store = _store(redis) + key = store.record_key("smart-router", "scope", "s1") + redis.hashes[store.sessions_key()] = {key: _record_json(schema_version=CACHE_WARMING_RECORD_SCHEMA_VERSION + 1)} + assert await store.get_record(key) is None + + +@pytest.mark.asyncio +async def test_get_record_returns_none_on_validation_error(): + redis = FakeRedisCache() + store = _store(redis) + key = store.record_key("smart-router", "scope", "s1") + redis.hashes[store.sessions_key()] = {key: json.dumps({"schema_version": "not-a-record"})} + assert await store.get_record(key) is None + + +@pytest.mark.asyncio +async def test_get_warmth_skips_corrupt_stamps(): + redis = FakeRedisCache() + store = _store(redis) + key = store.record_key("smart-router", "scope", "s1") + redis.data[store.warmth_key(key, "opus")] = json.dumps("not-a-float") + redis.data[store.warmth_key(key, "sonnet")] = json.dumps(123.5) + assert await store.get_warmth(key, ("opus", "sonnet")) == {"sonnet": 123.5} + + +@pytest.mark.asyncio +async def test_list_session_keys_returns_live_members_decoded(): + redis = FakeRedisCache(namespace="litellm") + store = _store(redis) + await _upsert(store, session_id="s1") + await _upsert(store, session_id="s2") + assert await store.list_session_keys(max_sessions=10) == ("scope:s1", "scope:s2") + + +@pytest.mark.asyncio +async def test_store_noops_without_redis_and_warns_once(caplog): + _warn_redis_missing.cache_clear() + store = CacheWarmingStore(redis_cache=None, auto_router_model_name="warnless-router") + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + assert await store.get_record("k") is None + assert await store.list_session_keys(max_sessions=5) == () + await store.mark_warm_attempt("k", "m", attempted_at=1.0, ttl_seconds=60) + assert await store.get_warmth("k", ("m",)) == {} + warnings = [r for r in caplog.records if "cache warming is inactive" in r.getMessage()] + assert len(warnings) == 1 diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py new file mode 100644 index 00000000000..734c041a1d9 --- /dev/null +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_types.py @@ -0,0 +1,77 @@ +import json + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.cache_warming.types import ( + CacheWarmingPayload, + compress_payload, + decompress_payload, +) + + +def _payload(**overrides: object) -> CacheWarmingPayload: + base: dict[str, object] = { + "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "messages": [ + {"role": "system", "content": "policy manual é中文 " * 50}, + {"role": "user", "content": [{"type": "text", "text": "summarize rule 7"}]}, + ], + "tools": [{"type": "function", "function": {"name": "lookup_rule", "parameters": {"type": "object"}}}], + "call_surface": "chat_completions", + } + return CacheWarmingPayload(**{**base, **overrides}) + + +def test_compress_decompress_roundtrip_preserves_payload(): + payload = _payload(system=[{"type": "text", "text": "cached system block"}], call_surface="anthropic_messages") + blob, _ = compress_payload(payload) + restored = decompress_payload(blob) + assert restored.model_dump() == payload.model_dump() + assert json.dumps(restored.messages) == json.dumps(payload.messages) + + +def test_roundtrip_preserves_message_key_order(): + payload = _payload(messages=[{"role": "user", "content": "hi", "name": "a"}]) + reordered = _payload(messages=[{"name": "a", "content": "hi", "role": "user"}]) + assert json.dumps(decompress_payload(compress_payload(payload)[0]).messages) != json.dumps( + decompress_payload(compress_payload(reordered)[0]).messages + ) + + +def test_decompress_rejects_corrupt_blob(): + with pytest.raises(Exception): + decompress_payload("not-base64-zlib!!") + with pytest.raises(ValidationError): + decompress_payload(compress_and_corrupt()) + + +def compress_and_corrupt() -> str: + import base64 + import zlib + + return base64.b64encode(zlib.compress(json.dumps({"model": "m"}).encode())).decode("ascii") + + +def test_payload_sha256_stable_across_key_order(): + ordered = _payload(messages=[{"role": "user", "content": "hi", "name": "a"}]) + reordered = _payload(messages=[{"name": "a", "content": "hi", "role": "user"}]) + different = _payload(messages=[{"role": "user", "content": "bye", "name": "a"}]) + assert compress_payload(ordered)[1] == compress_payload(reordered)[1] + assert compress_payload(ordered)[1] != compress_payload(different)[1] + + +def test_compression_shrinks_repetitive_payload(): + payload = _payload() + blob, _ = compress_payload(payload) + assert len(blob) < len(payload.model_dump_json()) + + +def test_payload_rejects_unknown_fields(): + with pytest.raises(ValidationError): + CacheWarmingPayload( + model="m", + messages=[], + call_surface="chat_completions", + api_key="sk-should-never-be-here", + ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ef70687bd97..2d6ca515422 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3400,3 +3400,187 @@ class TestEscalationKeywords: messages=[{"role": "user", "content": "LITELLM ESCALATE do better"}], ) assert result.model == "o1-b" # unchanged: no random hop to o1-a / o1-c + + +class TestCacheWarmingConfig: + def test_cache_warming_config_defaults_disabled(self): + config = ComplexityRouterConfig() + assert config.cache_warming.enabled is False + assert config.cache_warming.refresh_interval_seconds == 270 + assert config.cache_warming.session_ttl_seconds == 3600 + assert config.cache_warming.idle_timeout_seconds == 600 + assert config.cache_warming.max_sessions == 1000 + assert config.cache_warming.warm_models is None + + def test_cache_warming_coerces_from_nested_yaml_dict(self): + config = ComplexityRouterConfig( + cache_warming={"enabled": True, "refresh_interval_seconds": 120, "warm_models": ["sonnet", "opus"]} + ) + assert config.cache_warming.enabled is True + assert config.cache_warming.refresh_interval_seconds == 120 + assert config.cache_warming.warm_models == ("sonnet", "opus") + + def test_cache_warming_and_adaptive_mutually_exclusive_raises(self): + with pytest.raises(ValueError, match="cache_warming and adaptive"): + ComplexityRouterConfig( + tiers={"SIMPLE": ["gpt-4o-mini"]}, + adaptive=True, + cache_warming={"enabled": True}, + ) + + def test_cache_warming_disabled_with_adaptive_is_allowed(self): + config = ComplexityRouterConfig(tiers={"SIMPLE": ["gpt-4o-mini"]}, adaptive=True) + assert config.cache_warming.enabled is False + + def test_cache_warming_rejects_nonpositive_intervals(self): + with pytest.raises(ValidationError): + ComplexityRouterConfig(cache_warming={"enabled": True, "refresh_interval_seconds": 0}) + + +class TestCacheWarmingMarkerStamp: + @staticmethod + def _warming_router(mock_router_instance, session_affinity: bool = False) -> ComplexityRouter: + mock_router_instance.cache = DualCache() + return ComplexityRouter( + model_name="stamp-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-5"}, + "session_affinity": session_affinity, + "cache_warming": {"enabled": True}, + }, + ) + + @pytest.mark.asyncio + async def test_classify_path_stamps_marker(self, mock_router_instance): + from litellm.router_strategy.complexity_router.cache_warming.types import CACHE_WARMING_MARKER_KEY + + router = self._warming_router(mock_router_instance) + request_kwargs = {"metadata": {"session_id": "s1"}} + response = await router.async_pre_routing_hook( + model="stamp-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello there"}], + ) + assert response is not None + assert CACHE_WARMING_MARKER_KEY not in request_kwargs["metadata"] + marker = request_kwargs["litellm_metadata"][CACHE_WARMING_MARKER_KEY] + assert marker["auto_router_model_name"] == "stamp-router" + assert marker["routed_model"] == response.model + assert marker["strategy_ref"] == router._cache_warming_ref + + @pytest.mark.asyncio + async def test_affinity_pin_path_stamps_marker(self, mock_router_instance): + from litellm.router_strategy.complexity_router.cache_warming.types import CACHE_WARMING_MARKER_KEY + + router = self._warming_router(mock_router_instance, session_affinity=True) + cache_key = router._get_session_affinity_cache_key("s2", {"metadata": {"session_id": "s2"}}) + await mock_router_instance.cache.async_set_cache(key=cache_key, value="claude-sonnet-4-5") + request_kwargs = {"metadata": {"session_id": "s2"}} + response = await router.async_pre_routing_hook( + model="stamp-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello again"}], + ) + assert response is not None and response.model == "claude-sonnet-4-5" + assert CACHE_WARMING_MARKER_KEY not in request_kwargs["metadata"] + marker = request_kwargs["litellm_metadata"][CACHE_WARMING_MARKER_KEY] + assert marker["auto_router_model_name"] == "stamp-router" + assert marker["routed_model"] == "claude-sonnet-4-5" + assert marker["strategy_ref"] == router._cache_warming_ref + + @pytest.mark.asyncio + async def test_no_stamp_when_cache_warming_disabled(self, mock_router_instance): + from litellm.router_strategy.complexity_router.cache_warming.types import CACHE_WARMING_MARKER_KEY + + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="stamp-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}, + ) + request_kwargs = {"metadata": {"session_id": "s3"}} + await router.async_pre_routing_hook( + model="stamp-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + assert CACHE_WARMING_MARKER_KEY not in request_kwargs["metadata"] + assert CACHE_WARMING_MARKER_KEY not in request_kwargs.get("litellm_metadata", {}) + + @pytest.mark.asyncio + async def test_stamp_lands_in_litellm_metadata_slot_when_present(self, mock_router_instance): + from litellm.router_strategy.complexity_router.cache_warming.types import CACHE_WARMING_MARKER_KEY + + router = self._warming_router(mock_router_instance) + request_kwargs = {"litellm_metadata": {"session_id": "s4"}} + await router.async_pre_routing_hook( + model="stamp-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello there"}], + ) + assert CACHE_WARMING_MARKER_KEY in request_kwargs["litellm_metadata"] + assert "metadata" not in request_kwargs + + +class TestCacheWarmingDispatcherRegistry: + @staticmethod + def _model_list(alias: str, enabled: bool = True) -> List[Dict]: + return [ + { + "model_name": alias, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "cache_warming": {"enabled": enabled}, + }, + }, + }, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini", "api_key": "test"}}, + ] + + @staticmethod + def _dispatchers() -> List: + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import ( + ComplexityCacheWarmingCaptureHook, + ) + + return litellm.logging_callback_manager.get_custom_loggers_for_type(ComplexityCacheWarmingCaptureHook) + + def test_single_dispatcher_across_routers_and_reloads(self): + router_a = Router(model_list=self._model_list("auto-a")) + router_b = Router(model_list=self._model_list("auto-b")) + router_a.set_model_list(router_a.model_list) + router_b.set_model_list(router_b.model_list) + assert len(self._dispatchers()) == 1 + + def test_each_router_strategy_registered_under_its_own_ref(self): + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import _WARMING_STRATEGIES + + router_a = Router(model_list=self._model_list("auto-a")) + router_b = Router(model_list=self._model_list("auto-a")) + strategy_a = router_a.complexity_routers["auto-a"][0].strategy + strategy_b = router_b.complexity_routers["auto-a"][0].strategy + assert strategy_a._cache_warming_ref != strategy_b._cache_warming_ref + assert _WARMING_STRATEGIES.get(strategy_a._cache_warming_ref) is strategy_a + assert _WARMING_STRATEGIES.get(strategy_b._cache_warming_ref) is strategy_b + + def test_disabled_warming_registers_nothing(self): + router = Router(model_list=self._model_list("auto-off", enabled=False)) + strategy = router.complexity_routers["auto-off"][0].strategy + assert strategy._cache_warming_ref is None + + def test_reload_replaces_registration_and_old_ref_goes_stale(self): + import gc + + from litellm.router_strategy.complexity_router.cache_warming.capture_hook import _WARMING_STRATEGIES + + router = Router(model_list=self._model_list("auto-r")) + old_ref = router.complexity_routers["auto-r"][0].strategy._cache_warming_ref + router.set_model_list(self._model_list("auto-r")) + new_strategy = router.complexity_routers["auto-r"][0].strategy + gc.collect() + assert new_strategy._cache_warming_ref != old_ref + assert _WARMING_STRATEGIES.get(new_strategy._cache_warming_ref) is new_strategy + assert _WARMING_STRATEGIES.get(old_ref) is None