From a3a886b40dbc8038ec799c69332146494815e098 Mon Sep 17 00:00:00 2001 From: Tin Date: Fri, 9 Oct 2026 22:54:13 -0700 Subject: [PATCH] fix(router): preserve cache history across switches and skipped turns --- .../llms/anthropic/prompt_cache_prediction.py | 15 +++- litellm/llms/prompt_cache_estimation.py | 31 +++++++- .../proxy/hooks/autorouter_baseline_cache.py | 69 +++++++++++----- .../spend_tracking/baseline_accounting.py | 70 ++++++++++++++-- .../spend/test_baseline_accounting.py | 29 +++++++ .../test_anthropic_prompt_cache_prediction.py | 19 +++++ .../unit/llms/test_prompt_cache_estimation.py | 13 ++- .../hooks/test_autorouter_baseline_cache.py | 79 +++++++++++++++---- .../test_baseline_accounting.py | 75 +++++++++++++++++- 9 files changed, 351 insertions(+), 49 deletions(-) diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 634eb264b38..e7adfc4aca5 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -15,6 +15,7 @@ from pydantic import ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, Valid import litellm from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.anthropic.common_utils import ( AnthropicModelInfo, is_anthropic_oauth_key, @@ -33,11 +34,12 @@ from litellm.router_utils.baseline_request import ( ) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams, LiteLLM_Params -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, ProviderSpecificHeader from litellm.utils import supports_thinking_cache_preservation _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _HEADERS: Final = TypeAdapter(dict[str, str]) +_PROVIDER_HEADERS: Final = TypeAdapter(ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None) _MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) _SYSTEM: Final = TypeAdapter(str | list[dict[str, JsonValue]] | None) _counter: Final = AnthropicCountTokensHandler() @@ -637,6 +639,17 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M return None source: Final = {**parameters, "messages": request.get("messages"), "stream": request.get("stream", False)} try: + headers: Final = { + **_HEADERS.validate_python(request.get("headers") or {}), + **_HEADERS.validate_python(request.get("extra_headers") or {}), + **_HEADERS.validate_python( + ProviderSpecificHeaderUtils.get_provider_specific_headers( + _PROVIDER_HEADERS.validate_python(request.get("provider_specific_header")), "anthropic" + ) + ), + } + if not supported_prediction_headers(headers): + return None owned: Final = _JSON_OBJECT.validate_python(source) context: Final = {**{k: v for k, v in request.items() if k not in ("metadata", "litellm_metadata")}, **owned} resolved_model: Final = litellm.get_llm_provider(model=model, custom_llm_provider="anthropic")[0] diff --git a/litellm/llms/prompt_cache_estimation.py b/litellm/llms/prompt_cache_estimation.py index baab58d5feb..82192ac55ea 100644 --- a/litellm/llms/prompt_cache_estimation.py +++ b/litellm/llms/prompt_cache_estimation.py @@ -198,6 +198,34 @@ def _ttl(control: object, default: int) -> int: return _TTLS.get(value, default) if isinstance(value, str) else default +def _cache_lifetime(request: Mapping[str, JsonValue], provider: str, prices: ModelInfo | None) -> int: + retention: Final = request.get("prompt_cache_retention") + requested: Final = _TTLS.get(retention) if provider in ("openai", "azure") and isinstance(retention, str) else None + duration_pricing: Final = provider == "anthropic" or bool( + prices and prices.get("cache_creation_input_token_cost_above_1hr") is not None + ) + return _ttl( + request.get("prompt_cache_options"), + requested or (300 if duration_pricing else 600 if retention == "in_memory" else 1800), + ) + + +def _explicit_ttls(value: JsonValue) -> Iterator[int]: + if isinstance(value, dict): + yield _ttl(value.get("cache_control"), 0) + for child in value.values(): + yield from _explicit_ttls(child) + elif isinstance(value, list): + for child in value: + yield from _explicit_ttls(child) + + +def cache_ttl_upper_bound(request: Mapping[str, JsonValue] | None, provider: str, prices: ModelInfo | None) -> int: + if request is None: + return 3600 if provider == "anthropic" else max(_TTLS.values()) + return max((_cache_lifetime(request, provider, prices), *_explicit_ttls(dict(request)))) + + def estimate_cache_plan( request: dict[str, JsonValue], model: str, @@ -273,8 +301,7 @@ def estimate_cache_plan( duration_pricing: Final = provider == "anthropic" or bool( prices and prices.get("cache_creation_input_token_cost_above_1hr") is not None ) - default_ttl: Final = requested_ttl or (300 if duration_pricing else 600 if retention == "in_memory" else 1800) - lifetime: Final = _ttl(options, default_ttl) + lifetime: Final = _cache_lifetime(request, provider, prices) implicit: Final = not duration_pricing and options.get("mode") != "explicit" initial: Final = next( (index - 1 for index, part in enumerate(parts) if part.role not in ("system", "developer")), len(parts) - 1 diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index 8f8112c6a70..7e4686fb926 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -35,6 +35,7 @@ from litellm.llms.anthropic.prompt_cache_prediction import ( ) from litellm.llms.prompt_cache_estimation import ( EstimatedCachePlan, + cache_ttl_upper_bound, count_prefix_tokens, estimate_cache_plan, normalize_cache_usage, @@ -242,6 +243,7 @@ class AutoRouterBaselineCache(CustomLogger): identity.provider, _METADATA.validate_python(deployment.model_info.model_dump()) if deployment else {}, ) + native_plan: Final = native and isinstance(target, NativePredictionTarget) projected: Final = ( baseline_request( kwargs, @@ -254,19 +256,28 @@ class AutoRouterBaselineCache(CustomLogger): ) estimated_request: Final = ( prepare_cache_request(projected, identity.model, identity.provider, native=native) - if estimated and projected is not None + if estimated and not native_plan and projected is not None + else None + ) + baseline_body: Final = ( + prepare_native_baseline_body(projected, target.model) + if native_plan + and projected is not None + and isinstance(target, NativePredictionTarget) + and not (estimated and compaction_applied(kwargs)) else None ) scope: Final = "autorouter-baseline:v3:" + _digest( ( - "baseline_request_v3" if estimated else "baseline_request_v4", + "baseline_request_v5", request.user_api_key_hash, session, request.route.router_name, request.route.baseline_deployment_id, params, prices, - *((identity, "estimated_prefixes_v5") if estimated else ()), + identity, + "native_prefixes_v2" if native_plan else "estimated_prefixes_v5", ) ) started: Final = logging_obj.start_time.timestamp() @@ -288,6 +299,9 @@ class AutoRouterBaselineCache(CustomLogger): reason="incomplete_response", cache_policy="estimated" if estimated else "anthropic", assumptions=(), + cache_ttl_seconds=cache_ttl_upper_bound( + baseline_body or estimated_request, identity.provider, prices + ), cache_write_pricing="standard" if estimated and (prices is None or prices.get("cache_creation_input_token_cost_above_1hr") is None) else "duration", @@ -307,9 +321,7 @@ class AutoRouterBaselineCache(CustomLogger): request.route.baseline_deployment_id, estimated_request, estimated, - prepare_native_baseline_body(projected, target.model) - if not estimated and projected is not None and isinstance(target, NativePredictionTarget) - else None, + baseline_body, _native_body_digest(selected_body) if selected_body is not None and not compaction_applied(kwargs) else None, @@ -317,9 +329,22 @@ class AutoRouterBaselineCache(CustomLogger): except Exception: # noqa: BLE001 # optional observation cannot fail inference verbose_proxy_logger.warning("Auto-router baseline observation could not be initialized") - async def estimate(self, context: BaselineCacheContext, usage: Usage) -> EstimatedCachePlan | None: + async def estimate( + self, context: BaselineCacheContext, usage: Usage + ) -> tuple[EstimatedCachePlan | None, str | None]: + if isinstance(context.target, NativePredictionTarget) and context.baseline_body is not None: + plan, reason = await self.plan(context.target, context.baseline_body) + return ( + EstimatedCachePlan(plan, ("cold_cache_at_session_start", "same_output_tokens")) + if plan is not None + else None, + reason, + ) if self.estimate_slots.locked() or context.estimated_request is None: - return None + return ( + None, + "unsupported_cache_request" if context.estimated_request is None else "estimation_capacity_exhausted", + ) await self.estimate_slots.acquire() async def count() -> EstimatedCachePlan | None: @@ -336,7 +361,7 @@ class AutoRouterBaselineCache(CustomLogger): task: Final = asyncio.create_task(count()) task.add_done_callback(self._release_estimate_slot) - return await asyncio.shield(task) + return await asyncio.shield(task), None def _release_estimate_slot(self, task: asyncio.Task[EstimatedCachePlan | None]) -> None: self.estimate_slots.release() @@ -362,16 +387,20 @@ class AutoRouterBaselineCache(CustomLogger): return tokens async def plan( - self, target: NativePredictionTarget, wire: httpx.Request, body: Mapping[str, JsonValue], usage: Usage | None + self, + target: NativePredictionTarget, + body: Mapping[str, JsonValue], + wire: httpx.Request | None = None, + usage: Usage | None = None, ) -> tuple[CountedPromptCachePlan | None, str | None]: deadline: Final = asyncio.get_running_loop().time() + optional_callback_budget(_COUNT_TIMEOUT, fraction=0.75) - if not supported_prediction_headers(wire.headers): + if wire is not None and not supported_prediction_headers(wire.headers): return None, "unsupported_request_headers" plan: Final = parse_cache_plan(body) if isinstance(plan, UnsupportedCachePlan): return None, plan.reason details: Final = usage.prompt_tokens_details if usage is not None else None - selected: Final = parse_cache_plan(_JSON_BODY.validate_json(wire.content)) + selected: Final = parse_cache_plan(_JSON_BODY.validate_json(wire.content)) if wire is not None else plan if ( (isinstance(selected, UnsupportedCachePlan) or not selected.breakpoints) and details is not None @@ -423,7 +452,7 @@ async def finalize_baseline_cache(logging_obj: Logging, response_obj: object) -> try: capture: Final = ( await asyncio.wait_for(asyncio.shield(task), timeout=budget) - if context.estimated + if context.estimated and context.baseline_body is None else await asyncio.shield(task) ) if logging_obj.baseline_cache_context is active: @@ -491,7 +520,7 @@ async def _capture_native( same: Final = logging_obj.get_router_model_id() == context.baseline_deployment_id and _native_body_digest( projected ) == _native_body_digest(body) - plan, reason = await context.collector.plan(target, wire, projected, usage) + plan, reason = await context.collector.plan(target, projected, wire, usage) return capture.with_observation( BaselineObservation( request_id=original.request_id, @@ -505,6 +534,7 @@ async def _capture_native( plan=plan, reason=reason, minimum_cache_tokens=get_prompt_cache_min_tokens(target.model), + cache_ttl_seconds=original.cache_ttl_seconds, ) ) @@ -553,10 +583,10 @@ async def _capture_estimated( if raw_usage is not None else None ) - estimated: Final = ( + estimated, reason = ( await context.collector.estimate(context, usage) - if context.estimated_request is not None and usage is not None and not context.invalidated - else None + if usage is not None and not context.invalidated + else (None, None) ) plan: Final = estimated.plan if estimated is not None else None provider: Final = _METADATA.validate_python(logging_obj.model_call_details).get("custom_llm_provider") @@ -576,13 +606,16 @@ async def _capture_estimated( update={ "available_at": available, "outcome": "uncertain" if context.invalidated or usage is None else "complete", - "usage": prepare_baseline_usage(usage, context.capture.provider, context.estimated_request), + "usage": prepare_baseline_usage( + usage, context.capture.provider, context.baseline_body or context.estimated_request + ), "plan": plan, "minimum_cache_tokens": get_prompt_cache_min_tokens( f"{context.capture.provider}/{context.capture.model}" ), "assumptions": estimated.assumptions if estimated else (), "reason": context.invalidated + or reason or ("missing_usage" if usage is None else "unsupported_cache_request" if plan is None else None), } ), diff --git a/litellm/proxy/spend_tracking/baseline_accounting.py b/litellm/proxy/spend_tracking/baseline_accounting.py index 9deccb30dcf..17ed248459a 100644 --- a/litellm/proxy/spend_tracking/baseline_accounting.py +++ b/litellm/proxy/spend_tracking/baseline_accounting.py @@ -7,7 +7,7 @@ owners. Replaying these values in event order is independent of callback order. from __future__ import annotations from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from itertools import groupby from math import isfinite from types import MappingProxyType @@ -41,6 +41,7 @@ class BaselineObservation(LiteLLMBaseModel): cache_policy: Literal["anthropic", "estimated"] = "anthropic" assumptions: tuple[str, ...] = () cache_write_pricing: Literal["duration", "standard"] = "duration" + cache_ttl_seconds: int | None = Field(default=None, ge=0, le=MAX_ESTIMATED_CACHE_TTL) @dataclass(frozen=True, slots=True) @@ -70,6 +71,9 @@ class BaselineHistory: uncertain_before: float = 0.0 entries: tuple[CacheEntry, ...] = () blocked_until: float = 0.0 + uncertain_until: float = 0.0 + uncertain_ttl_seconds: int = 0 + version: Literal[1, 2] = 1 def _complete_usage(usage: Usage | None, policy: str = "anthropic") -> bool: @@ -217,6 +221,8 @@ def _estimate(history: BaselineHistory, observation: BaselineObservation, equiva if not _valid_plan(observation) or plan is None: return BaselineEstimate(observation.request_id, observation.reason or "unsupported_cache_plan") markers: Final = _markers(observation) + if markers and observation.started_at < history.uncertain_until: + return BaselineEstimate(observation.request_id, "history_unavailable") if any(_ambiguous(entry, markers, observation.started_at) for entry in history.entries): return BaselineEstimate(observation.request_id, "cache_ttl_changed") read: Final = max( @@ -268,6 +274,7 @@ def _writes(history: BaselineHistory, observation: BaselineObservation) -> tuple if ( observation.outcome != "complete" or observation.started_at < history.blocked_until + or observation.started_at < history.uncertain_until or not _complete_usage(observation.usage, observation.cache_policy) or not _valid_plan(observation) ): @@ -308,7 +315,7 @@ def _writes(history: BaselineHistory, observation: BaselineObservation) -> tuple marker.fingerprint, marker.content_fingerprint, marker.prefix_tokens, - marker.ttl_seconds, + max((marker.ttl_seconds, *(entry.ttl_seconds for entry in ambiguous))), observation.available_at, observation.started_at + max((marker.ttl_seconds, *(entry.ttl_seconds for entry in ambiguous))), uncertain=bool(ambiguous), @@ -372,9 +379,34 @@ def advance_baseline_history( if item.outcome != "response_cache" and item.cache_policy == "estimated" and item.plan is None ) relevant: Final = tuple(item for item in simultaneous if item.outcome != "response_cache" and item not in skipped) - equivalent: Final = history.equivalent and all(item.baseline_equivalent for item in relevant) - before: Final = BaselineHistory( - first, history.last_at, equivalent, uncertain, history.entries, history.blocked_until + equivalent: Final = history.equivalent and all( + item.baseline_equivalent for item in simultaneous if item.outcome != "response_cache" + ) + legacy_lifetime: Final = ( + MAX_ESTIMATED_CACHE_TTL + if history.version == 1 and (history.uncertain_before > 0 or any(entry.uncertain for entry in history.entries)) + else 0 + ) + legacy_until: Final = ( + max( + ( + history.uncertain_before, + history.last_at or 0, + history.blocked_until, + *(entry.available_at for entry in history.entries), + ) + ) + + legacy_lifetime + if legacy_lifetime + else 0 + ) + before: Final = replace( + history, + first_at=first, + equivalent=equivalent, + uncertain_before=uncertain, + uncertain_until=max(history.uncertain_until, legacy_until), + uncertain_ttl_seconds=max(history.uncertain_ttl_seconds, legacy_lifetime), ) estimates: Final = tuple(_estimate(before, item, equivalent) for item in simultaneous) invalidated: Final = any( @@ -387,13 +419,35 @@ def advance_baseline_history( started, ) overflow: Final = len(entries) > MAX_CACHE_ENTRIES + active: Final = started < before.uncertain_until + unknown: Final = tuple( + item + for item in simultaneous + if item.outcome != "response_cache" + and (invalidated or overflow or item in skipped or active and bool(_markers(item))) + ) + lifetime: Final = max( + ( + before.uncertain_ttl_seconds if active else 0, + *(entry.ttl_seconds for entry in history.entries if unknown and entry.expires_at > started), + *( + item.cache_ttl_seconds + if item.cache_ttl_seconds is not None + else MAX_ESTIMATED_CACHE_TTL + if item.cache_policy == "estimated" + else MAX_CACHE_TTL + for item in unknown + ), + ) + ) return BaselineHistory( first_at=first, last_at=started, equivalent=equivalent, - uncertain_before=max(started, blocked) - if invalidated or overflow - else max((uncertain, *(item.available_at for item in skipped))), + uncertain_before=uncertain, entries=() if overflow else entries, blocked_until=blocked, + uncertain_until=max((before.uncertain_until, *(item.available_at + lifetime for item in unknown))), + uncertain_ttl_seconds=lifetime, + version=2, ), estimates diff --git a/tests/proxy_behavior/spend/test_baseline_accounting.py b/tests/proxy_behavior/spend/test_baseline_accounting.py index 4298b85cab4..fc322f82417 100644 --- a/tests/proxy_behavior/spend/test_baseline_accounting.py +++ b/tests/proxy_behavior/spend/test_baseline_accounting.py @@ -108,6 +108,35 @@ async def _user_sessions(db: Prisma, record: BaselineAccountingRecord) -> dict[s return {str(row["user_id"]): row for row in rows} +async def test_skipped_lifetime_survives_checkpoints_without_changing_spend( + db: Prisma, record: Callable[..., BaselineAccountingRecord], +) -> None: + store: Final = _store(db) + marker: Final = CountedBreakpoint("prefix:300", 300, 6000, ("prefix:300",), "content", ("content",)) + for label, started, skipped, expected in ( + ("cold", 10000.0, False, "cache_prefix_cold"), + ("skipped", 10002.0, True, "unsupported_cache_plan"), + ("shorter", 10400.0, False, "history_unavailable"), + ("expired", 14001.0, False, "cache_prefix_expired"), + ): + original: Final = record(label, started, identical=False) + item: Final = original.model_copy(update={"observation": original.observation.model_copy(update={ + "cache_policy": "estimated", "cache_ttl_seconds": 3600 if skipped else 300, + "plan": None if skipped else CountedPromptCachePlan(6200, (marker,)), + })}) + await _log(db, item) + assert await store.append(item) == "recorded" + assert await store.project(item.scope) == "published" + rows: Final = await db.query_raw( + 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', item.observation.request_id, + ) + assert rows[0]["spend"] == 0.17 + assert rows[0]["metadata"]["autorouter_savings_estimate"]["reason"] == expected + session: Final = await _session(db, item) + assert session["turns"] == 4 and session["spend"] == 0.68 + assert session["savings_estimated_turns"] == 2 + + async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: store: Final = _store(db) late: Final = record("late", 10001.0, user_id="late-user") diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 9bb26b66aa7..dfaf54dfdb8 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -25,6 +25,7 @@ from litellm.llms.anthropic.prompt_cache_prediction import ( parse_cache_plan, parse_observed_cache, parse_prompt, + prepare_native_baseline_body, resolve_baseline_prediction_target, resolve_prediction_target, supported_prediction_headers, @@ -55,6 +56,24 @@ def _body() -> dict[str, JsonValue]: } +@pytest.mark.parametrize("source", ("headers", "extra_headers", "provider_specific_header")) +@pytest.mark.parametrize("name,value,supported", ( + ("anthropic-beta", "unmodeled-cache-policy", False), + ("anthropic-version", "2099-01-01", False), + ("anthropic-version", DEFAULT_ANTHROPIC_API_VERSION, True), +)) +def test_native_baseline_preparation_preserves_header_eligibility( + source: str, name: str, value: str, supported: bool, +) -> None: + headers: Final = {name: value} + request: Final = { + **_body(), "max_tokens": 16, + source: {"custom_llm_provider": "anthropic", "extra_headers": headers} + if source == "provider_specific_header" else headers, + } + assert (prepare_native_baseline_body(request, _MODEL) is not None) == supported + + @pytest.mark.parametrize("version", [None, "2099-01-01", DEFAULT_ANTHROPIC_API_VERSION]) @pytest.mark.asyncio async def test_observer_records_only_version_supported_by_token_counter(version: str | None) -> None: diff --git a/tests/unit/llms/test_prompt_cache_estimation.py b/tests/unit/llms/test_prompt_cache_estimation.py index 45b48bf0f44..771b80e6f71 100644 --- a/tests/unit/llms/test_prompt_cache_estimation.py +++ b/tests/unit/llms/test_prompt_cache_estimation.py @@ -7,7 +7,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.prompt_cache_estimation import estimate_cache_plan, normalize_cache_usage, prepare_cache_request +from litellm.llms.prompt_cache_estimation import cache_ttl_upper_bound, estimate_cache_plan, normalize_cache_usage, prepare_cache_request from litellm.proxy.spend_tracking.baseline_accounting import ( BaselineHistory, BaselineObservation, @@ -56,6 +56,14 @@ def test_plain_responses_input_remains_estimatable() -> None: assert captured.plan.total_tokens == _usage().prompt_tokens +@pytest.mark.parametrize("ttl,seconds", (("5m", 300), ("1h", 3600), ("24h", 86400))) +def test_skipped_request_lifetime_is_captured_without_a_token_plan(ttl: str, seconds: int) -> None: + request: Final = prepare_cache_request(_request(cache_control={"type": "ephemeral", "ttl": ttl})) + assert request is not None + assert cache_ttl_upper_bound(request, "anthropic", None) == seconds + assert cache_ttl_upper_bound(None, "openai", None) >= seconds + + def _observation( request: dict[str, object], started: float = 10000.0, @@ -79,6 +87,7 @@ def _observation( cache_policy="estimated", cache_write_pricing="standard", assumptions=captured.assumptions, + cache_ttl_seconds=cache_ttl_upper_bound(prepared, provider, _PRICES), ) @@ -503,7 +512,7 @@ def test_skipped_estimation_preserves_paid_cache_without_refreshing_it(reason: s _, followup = advance_baseline_history( retained, (first.model_copy(update={"request_id": "followup", "started_at": 10004.0, "available_at": 10005.0}),) ) - assert followup[0].usage is not None and followup[0].usage.prompt_tokens_details.cached_tokens == 8000 + assert followup[0].usage is None and followup[0].reason == "history_unavailable" _, refreshed = advance_baseline_history( retained, (first.model_copy(update={"request_id": "refreshed", "started_at": 11800.0, "available_at": 11801.0}),), diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index 4b605673182..55b61ff627d 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -599,11 +599,13 @@ async def test_plain_fallback_invalidates_existing_autorouter_capture(monkeypatc assert logging.baseline_observation.observation.reason == "retried_request" +@pytest.mark.parametrize("adapted", (False, True)) async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and_spend( - monkeypatch: pytest.MonkeyPatch, + monkeypatch: pytest.MonkeyPatch, adapted: bool, ) -> None: from litellm.litellm_core_utils import logging_worker from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.types.utils import ModelResponse, Usage release: Final = asyncio.Event() @@ -615,9 +617,22 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and worker: Final = LoggingWorker(timeout=8.0) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) - rig: Final = _Rig(monkeypatch, count=count) + models: Final = _MESSAGES.validate_python([ + _MODELS[0], + {**_MODELS[1], "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "test-selected"}} + if adapted else _MODELS[1], + _MODELS[2], + ]) + rig: Final = _Rig(monkeypatch, count=count, models=models) + monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", True) try: - with _transport(_upstream): + with respx.mock(assert_all_called=False) as transport: + transport.post("https://api.anthropic.com/v1/messages").mock(side_effect=_upstream) + transport.post("https://api.openai.com/v1/chat/completions").respond(json=ModelResponse( + model="gpt-6-astra", + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "OK"}}], + usage=Usage(prompt_tokens=6000, completion_tokens=10, total_tokens=6010), + ).model_dump()) await _call(rig.router, rig.logging()) payload: Final = await rig.capture.payload() observed: Final = _observation(payload).observation @@ -1279,9 +1294,13 @@ async def test_messages_estimate_keeps_adapted_baseline_extra_body_retention( @pytest.mark.parametrize("chat_adapter", (False, True)) +@pytest.mark.parametrize("native_first", (False, True)) +@pytest.mark.parametrize("count_available", (False, True)) async def test_adapted_messages_tier_estimates_native_anthropic_baseline( - monkeypatch: pytest.MonkeyPatch, chat_adapter: bool + monkeypatch: pytest.MonkeyPatch, chat_adapter: bool, native_first: bool, count_available: bool ) -> None: + from litellm.proxy.spend_tracking.baseline_accounting import BaselineHistory, advance_baseline_history + models: Final = _MESSAGES.validate_python( [ _MODELS[0], @@ -1297,7 +1316,10 @@ async def test_adapted_messages_tier_estimates_native_anthropic_baseline( _MODELS[2], ] ) - rig: Final = _Rig(monkeypatch, models=models) + async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: + return await _count(model, api_key, body) if count_available else None + + rig: Final = _Rig(monkeypatch, models=models, count=count) monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", chat_adapter) def openai_response(request: httpx.Request) -> httpx.Response: @@ -1311,7 +1333,7 @@ async def test_adapted_messages_tier_estimates_native_anthropic_baseline( "created": 1, "model": body["model"], "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "OK"}}], - "usage": {"prompt_tokens": 6000, "completion_tokens": 10, "total_tokens": 6010}, + "usage": {"prompt_tokens": 9000, "completion_tokens": 10, "total_tokens": 9010}, } if chat_adapter else { @@ -1321,22 +1343,49 @@ async def test_adapted_messages_tier_estimates_native_anthropic_baseline( "status": "completed", "model": body["model"], "output": [], - "usage": {"input_tokens": 6000, "output_tokens": 10, "total_tokens": 6010}, + "usage": {"input_tokens": 9000, "output_tokens": 10, "total_tokens": 9010}, }, ) with respx.mock(assert_all_called=False) as transport: - anthropic: Final = transport.post("https://api.anthropic.com/v1/messages").mock(side_effect=_upstream) + transport.post("https://api.anthropic.com/v1/messages").mock(side_effect=_upstream) transport.post("https://api.openai.com/v1/" + ("chat/completions" if chat_adapter else "responses")).mock( side_effect=openai_response ) - log: Final = rig.logging() - await _call(rig.router, log) - captured: Final = _observation(await rig.capture.payload()) - assert not anthropic.calls - assert log.baseline_cache_context is not None and log.baseline_cache_context.estimated - assert captured.observation.cache_policy == "estimated" and captured.observation.outcome == "complete" - assert captured.observation.plan is not None and captured.observation.reason is None + async def capture(native: bool, started: float) -> CapturedBaselineObservation: + messages: Final = _MESSAGES_JSON.replace("1h", "5m").replace( + "question", "question USE_OPUS" if native else "question" + ) + await _call(rig.router, rig.logging(), messages=messages) + captured: Final = _observation(await rig.capture.payload()) + assert captured.observation.usage is not None + assert captured.observation.usage.prompt_tokens == (6000 if native else 9000) + return captured.with_observation( + captured.observation.model_copy( + update={"request_id": str(started), "started_at": started, "available_at": started + 1} + ) + ) + + captures: Final = ( + await capture(native_first, 1.0), + await capture(not native_first, 101.0), + await capture(native_first, 351.0), + ) + assert len({captured.scope for captured in captures}) == 1 + assert all(captured.observation.cache_ttl_seconds == 300 for captured in captures) + if not count_available: + assert all(captured.observation.plan is None for captured in captures) + assert all(captured.observation.reason == "token_count_unavailable" for captured in captures) + return + for captured in captures: + assert captured.observation.plan is not None and captured.observation.plan.total_tokens == 6000 + first, _ = advance_baseline_history(BaselineHistory(), (captures[0].observation,)) + second, _ = advance_baseline_history(first, (captures[1].observation,)) + _, last = advance_baseline_history(second, (captures[2].observation,)) + assert last[0].reason == "cache_prefix_available" + assert last[0].usage is not None + assert last[0].usage.prompt_tokens_details.cached_tokens == 5000 + assert last[0].usage.prompt_tokens_details.cache_creation_tokens == 0 @pytest.mark.parametrize("passthrough", (False, True)) diff --git a/tests/unit/proxy/spend_tracking/test_baseline_accounting.py b/tests/unit/proxy/spend_tracking/test_baseline_accounting.py index 499fd16aa40..a71b0dfda19 100644 --- a/tests/unit/proxy/spend_tracking/test_baseline_accounting.py +++ b/tests/unit/proxy/spend_tracking/test_baseline_accounting.py @@ -3,6 +3,7 @@ from itertools import groupby from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.llms.anthropic.cost_calculation import cost_per_token @@ -84,6 +85,68 @@ def _replay(*observations: BaselineObservation) -> tuple[BaselineEstimate, ...]: return tuple(results) +@pytest.mark.parametrize("ttl", (3600, 86400, None)) +@pytest.mark.parametrize("outcome", ("complete", "uncertain")) +def test_skipped_long_cache_lifetime_survives_shorter_followups(ttl: int | None, outcome: str) -> None: + bound: Final = ttl or 86400 + first: Final = _observation( + "known", 1.0, cache_policy="estimated", cache_ttl_seconds=300, + plan=CountedPromptCachePlan(6200, (_marker(ttl=300),)), + ) + skipped: Final = first.model_copy( + update={ + "request_id": "skipped", "started_at": 10.0, "available_at": 11.0, + "plan": None, "cache_ttl_seconds": ttl, "outcome": outcome, + } + ) + history, _ = advance_baseline_history(BaselineHistory(), (first,)) + missing, _ = advance_baseline_history(history, (skipped,)) + assert missing.entries == history.entries + + def followup(started: float) -> BaselineObservation: + return first.model_copy(update={"request_id": str(started), "started_at": started, "available_at": started + 1}) + + known, still_unknown = advance_baseline_history(missing, (followup(12.0),)) + assert still_unknown[0].reason == "history_unavailable" + assert known.entries == history.entries + refreshed, later = advance_baseline_history(known, (followup(float(bound)),)) + assert later[0].reason == "history_unavailable" + _, prior_deadline = advance_baseline_history(refreshed, (followup(bound + 15.0),)) + assert prior_deadline[0].reason == "history_unavailable" + recovered, cold = advance_baseline_history(refreshed, (followup(2 * bound + 2.0),)) + assert cold[0].usage is not None and cold[0].usage.prompt_tokens_details.cached_tokens == 0 + _, warm = advance_baseline_history(recovered, (followup(2 * bound + 4.0),)) + assert warm[0].usage is not None and warm[0].usage.prompt_tokens_details.cached_tokens == 6000 + + +@pytest.mark.parametrize("ambiguous_entry", (False, True)) +def test_legacy_checkpoint_accounts_for_refreshes_after_a_skipped_request(ambiguous_entry: bool) -> None: + restored: Final = TypeAdapter(BaselineHistory).validate_python({ + "first_at": 1.0, "last_at": 80000.0, "equivalent": False, + "uncertain_before": 0.0 if ambiguous_entry else 1000.0, + "entries": (CacheEntry("p:300", "p", 6000, 300, 80000.0, 83600.0, uncertain=True),) + if ambiguous_entry else (), + }) + updated, pending = advance_baseline_history(restored, (_observation("pending", 90000.0),)) + assert pending[0].reason == "history_unavailable" and pending[0].usage is None + checkpoint: Final = TypeAdapter(BaselineHistory).validate_json(TypeAdapter(BaselineHistory).dump_json(updated)) + _, expired = advance_baseline_history(checkpoint, (_observation("expired", 180001.0),)) + assert expired[0].usage is not None and expired[0].usage.prompt_tokens_details.cached_tokens == 0 + + +def test_skipped_turn_preserves_lifetime_of_an_earlier_ambiguous_write() -> None: + first: Final = _observation("long", 1.0, cache_policy="estimated", cache_ttl_seconds=3600) + short: Final = _observation( + "short", 3000.0, cache_policy="estimated", cache_ttl_seconds=300, + plan=CountedPromptCachePlan(6200, (_marker(ttl=300),)), + ) + skipped: Final = short.model_copy(update={"request_id": "skip", "started_at": 4000.0, "available_at": 4001.0, "plan": None}) + later: Final = short.model_copy(update={"request_id": "later", "started_at": 7000.0, "available_at": 7001.0}) + estimates: Final = _replay(first, short, skipped, later) + assert estimates[1].reason == "cache_ttl_changed" + assert estimates[3].reason == "history_unavailable" and estimates[3].usage is None + + @pytest.mark.parametrize("write_rate", (None, 0.0, 0.017)) @pytest.mark.parametrize("prefix", (4000, 8000)) def test_multimodal_history_prices_absent_zero_and_premium_writes(write_rate: float | None, prefix: int) -> None: @@ -192,6 +255,7 @@ def test_growth_lookback_and_mixed_ttl_keep_distinct_read_write_buckets(warm_tai second: Final = _replay(first, _observation("second", 10001.0, plan=grown))[-1] assert second.reason == "history_unavailable" history: Final = BaselineHistory( + version=2, first_at=1.0, last_at=10000.0, equivalent=False, @@ -252,11 +316,15 @@ def test_overlapping_uncertain_request_cannot_be_warmed_by_a_later_callback() -> overlap: Final = _observation("overlap", 10001.0) during: Final = _observation("during", 10002.0) after: Final = _observation("after", 10011.0) - warmed: Final = _observation("warmed", 10012.0) - estimates: Final = _replay(uncertain, overlap, during, after, warmed) + during_lifetime: Final = _observation("during_lifetime", 10012.0) + expired: Final = _observation("expired", 13613.0) + warmed: Final = _observation("warmed", 13614.0) + estimates: Final = _replay(uncertain, overlap, during, after, during_lifetime, expired, warmed) assert estimates[1].reason == estimates[2].reason == "concurrent_uncertainty" assert estimates[3].usage is None - assert estimates[4].usage is not None and estimates[4].usage.prompt_tokens_details.cached_tokens == 6000 + assert estimates[4].reason == "history_unavailable" + assert estimates[5].usage is not None and estimates[5].usage.prompt_tokens_details.cached_tokens == 0 + assert estimates[6].usage is not None and estimates[6].usage.prompt_tokens_details.cached_tokens == 6000 def test_modeled_read_cannot_recharge_the_original_private_write_count() -> None: @@ -277,6 +345,7 @@ def test_modeled_read_cannot_recharge_the_original_private_write_count() -> None def test_mixed_lifetime_lookback_preserves_a_compatible_native_hit() -> None: marker: Final = _marker("prefix", 3600, 6000) history: Final = BaselineHistory( + version=2, first_at=1.0, last_at=10000.0, equivalent=False,