mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(router): preserve cache history across switches and skipped turns
This commit is contained in:
Tin
parent
3755867355
commit
a3a886b40d
9 changed files with 351 additions and 49 deletions
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}),),
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue