mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
Merge pull request #40877 from BerriAI/litellm_lit7658_cache_cost_v0_fresh
feat(proxy): predict prompt-cache costs across deployments
This commit is contained in:
tin-berri
GitHub
commit
c626ff098b
20 changed files with 3106 additions and 89 deletions
388
litellm/llms/anthropic/prompt_cache_prediction.py
Normal file
388
litellm/llms/anthropic/prompt_cache_prediction.py
Normal file
|
|
@ -0,0 +1,388 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_HEADERS: Final = TypeAdapter(dict[str, str])
|
||||
_counter: Final = AnthropicCountTokensHandler()
|
||||
|
||||
|
||||
_NATIVE_HEADERS: Final = frozenset(
|
||||
(
|
||||
"host",
|
||||
"accept",
|
||||
"accept-encoding",
|
||||
"connection",
|
||||
"user-agent",
|
||||
"content-length",
|
||||
"content-type",
|
||||
"x-api-key",
|
||||
"anthropic-version",
|
||||
)
|
||||
)
|
||||
|
||||
_DEPLOYMENT_OPTIONS: Final = frozenset(
|
||||
{
|
||||
"model",
|
||||
"api_key",
|
||||
"api_base",
|
||||
"custom_llm_provider",
|
||||
"rpm",
|
||||
"tpm",
|
||||
"timeout",
|
||||
"stream_timeout",
|
||||
"max_retries",
|
||||
"num_retries",
|
||||
"max_parallel_requests",
|
||||
"input_cost_per_token",
|
||||
"output_cost_per_token",
|
||||
"cache_read_input_token_cost",
|
||||
"cache_creation_input_token_cost",
|
||||
"cache_creation_input_token_cost_above_1hr",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _StrictModel(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
|
||||
|
||||
|
||||
class _CacheControl(_StrictModel):
|
||||
type: Literal["ephemeral"]
|
||||
ttl: Literal["5m", "1h"] = "5m"
|
||||
|
||||
|
||||
class _Text(_StrictModel):
|
||||
type: Literal["text"]
|
||||
text: str = Field(min_length=1, pattern=r"\S")
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
class _ToolUse(_StrictModel):
|
||||
type: Literal["tool_use"]
|
||||
id: str = Field(min_length=1)
|
||||
name: str = Field(min_length=1)
|
||||
input: Mapping[str, JsonValue]
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
class _ResultText(_StrictModel):
|
||||
type: Literal["text"]
|
||||
text: str
|
||||
|
||||
|
||||
class _ToolResult(_StrictModel):
|
||||
type: Literal["tool_result"]
|
||||
tool_use_id: str = Field(min_length=1)
|
||||
content: str | Annotated[tuple[_ResultText, ...], Field(strict=False)]
|
||||
is_error: bool | None = None
|
||||
cache_control: _CacheControl | None = None
|
||||
|
||||
|
||||
_Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminator="type")]
|
||||
|
||||
|
||||
class _Message(_StrictModel):
|
||||
role: Literal["user", "assistant"]
|
||||
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
|
||||
|
||||
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
|
||||
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
|
||||
|
||||
|
||||
class _Tool(_StrictModel):
|
||||
name: str = Field(min_length=1)
|
||||
description: str | None = None
|
||||
input_schema: Mapping[str, JsonValue]
|
||||
type: Literal["custom"] | None = None
|
||||
|
||||
|
||||
class _Request(_StrictModel):
|
||||
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
|
||||
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
|
||||
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
|
||||
model: str | None = None
|
||||
max_tokens: int | None = None
|
||||
stream: bool | None = None
|
||||
temperature: float | int | None = None
|
||||
top_p: float | int | None = None
|
||||
top_k: int | None = None
|
||||
stop_sequences: Annotated[tuple[str, ...], Field(strict=False)] | None = None
|
||||
metadata: Mapping[str, JsonValue] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptPrefix:
|
||||
prefix_body: Mapping[str, JsonValue]
|
||||
fingerprint: str
|
||||
fingerprints: tuple[str, ...]
|
||||
ttl_seconds: int
|
||||
|
||||
|
||||
def _digest(value: object) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
|
||||
return _digest((previous, boundary))
|
||||
|
||||
|
||||
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
|
||||
try:
|
||||
request: Final = _Request.model_validate(body)
|
||||
blocks: Final = tuple(message.blocks() for message in request.messages)
|
||||
except ValidationError:
|
||||
return None
|
||||
markers: Final = tuple(
|
||||
(message_index, block_index, block.cache_control)
|
||||
for message_index, message_blocks in enumerate(blocks)
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if block.cache_control is not None
|
||||
)
|
||||
if len(markers) != 1:
|
||||
return None
|
||||
message_end, block_end, marker = markers[0]
|
||||
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
|
||||
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
|
||||
boundaries: Final = tuple(
|
||||
(
|
||||
message_index,
|
||||
request.messages[message_index].role,
|
||||
_JSON_OBJECT.validate_python(
|
||||
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
|
||||
),
|
||||
)
|
||||
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if message_index < message_end or block_index <= block_end
|
||||
)
|
||||
hashes: Final = tuple(
|
||||
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
|
||||
)[1:]
|
||||
prefix_messages: Final = tuple(
|
||||
_Message(
|
||||
role=request.messages[message_index].role,
|
||||
content=tuple(
|
||||
block
|
||||
for block_index, block in enumerate(message_blocks)
|
||||
if message_index < message_end or block_index <= block_end
|
||||
),
|
||||
)
|
||||
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
|
||||
)
|
||||
return PromptPrefix(
|
||||
prefix_body=MappingProxyType(
|
||||
_JSON_OBJECT.validate_python(
|
||||
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
|
||||
mode="json", exclude_none=True
|
||||
)
|
||||
)
|
||||
),
|
||||
fingerprint=hashes[-1],
|
||||
fingerprints=tuple(reversed(hashes[-20:])),
|
||||
ttl_seconds=3600 if marker.ttl == "1h" else 300,
|
||||
)
|
||||
|
||||
|
||||
def cache_scope(
|
||||
caller_key_hash: str,
|
||||
deployment_id: str,
|
||||
provider_key: str,
|
||||
model: str,
|
||||
anthropic_version: str = DEFAULT_ANTHROPIC_API_VERSION,
|
||||
) -> str:
|
||||
return _digest((caller_key_hash, deployment_id, provider_key, model, anthropic_version))
|
||||
|
||||
|
||||
class _TTLUsage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
ephemeral_5m_input_tokens: int = Field(default=0, ge=0)
|
||||
ephemeral_1h_input_tokens: int = Field(default=0, ge=0)
|
||||
|
||||
|
||||
class _CacheUsage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
cached_tokens: int = Field(default=0, ge=0)
|
||||
cache_creation_tokens: int = Field(default=0, ge=0)
|
||||
cache_creation_token_details: _TTLUsage | None = None
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
prompt_tokens: int = Field(ge=0)
|
||||
prompt_tokens_details: _CacheUsage
|
||||
|
||||
|
||||
class _Choice(BaseModel):
|
||||
finish_reason: str = Field(min_length=1)
|
||||
|
||||
|
||||
class _Response(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
model: str
|
||||
usage: _Usage
|
||||
choices: tuple[_Choice, ...] = Field(min_length=1, strict=False)
|
||||
|
||||
|
||||
class _CountBody(BaseModel):
|
||||
messages: Sequence[Mapping[str, JsonValue]]
|
||||
tools: Sequence[Mapping[str, JsonValue]] | None = None
|
||||
system: str | Sequence[Mapping[str, JsonValue]] | None = None
|
||||
|
||||
|
||||
class _CountResult(BaseModel):
|
||||
input_tokens: Annotated[StrictInt, Field(ge=0)]
|
||||
|
||||
|
||||
class TokenCounter(Protocol):
|
||||
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: ...
|
||||
|
||||
|
||||
def _count_objects(
|
||||
values: Sequence[Mapping[str, JsonValue]],
|
||||
) -> list[dict[str, JsonValue]]: # mutable-ok: the existing provider count API requires JSON lists/dicts
|
||||
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
|
||||
|
||||
|
||||
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
native: Final = _CountBody.model_validate(body)
|
||||
try:
|
||||
result: Final = _CountResult.model_validate(
|
||||
await _counter.handle_count_tokens_request(
|
||||
model=model,
|
||||
messages=_count_objects(native.messages),
|
||||
tools=_count_objects(native.tools) if native.tools is not None else None,
|
||||
system=native.system,
|
||||
api_key=api_key,
|
||||
timeout=15.0,
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # provider/count validation failures are unavailable estimates, not zero tokens
|
||||
return None
|
||||
return result.input_tokens
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativePredictionTarget:
|
||||
model: str
|
||||
api_key: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnsupportedPredictionTarget:
|
||||
reason: Literal[
|
||||
"unsupported_deployment_configuration",
|
||||
"unsupported_provider_endpoint",
|
||||
"unsupported_provider",
|
||||
"unsupported_provider_credentials",
|
||||
]
|
||||
|
||||
|
||||
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
|
||||
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
|
||||
if configured_options - _DEPLOYMENT_OPTIONS:
|
||||
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
|
||||
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
|
||||
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
|
||||
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
|
||||
try:
|
||||
model, provider, _, _ = litellm.get_llm_provider(
|
||||
model=params.model, custom_llm_provider=params.custom_llm_provider
|
||||
)
|
||||
except Exception: # noqa: BLE001 # the shared provider resolver raises for unknown deployments
|
||||
return UnsupportedPredictionTarget("unsupported_provider")
|
||||
if provider != "anthropic":
|
||||
return UnsupportedPredictionTarget("unsupported_provider")
|
||||
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
|
||||
if api_key is None or not _supported_provider_key(api_key):
|
||||
return UnsupportedPredictionTarget("unsupported_provider_credentials")
|
||||
return NativePredictionTarget(model=model, api_key=api_key)
|
||||
|
||||
|
||||
def _supported_provider_key(api_key: str) -> bool:
|
||||
return bool(api_key) and not is_anthropic_oauth_key(api_key)
|
||||
|
||||
|
||||
def supported_prediction_headers(headers: Mapping[str, str]) -> bool:
|
||||
return all(
|
||||
name.lower() != "anthropic-beta"
|
||||
and (name.lower() != "anthropic-version" or value == DEFAULT_ANTHROPIC_API_VERSION)
|
||||
for name, value in headers.items()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ObservedCachePrefix:
|
||||
prefix: PromptPrefix
|
||||
scope: str
|
||||
cached_tokens: int
|
||||
cache_creation_tokens: int
|
||||
|
||||
|
||||
def parse_observed_cache(
|
||||
wire: httpx.Request, response_obj: ModelResponse, caller_key_hash: str, deployment_id: str
|
||||
) -> ObservedCachePrefix | None:
|
||||
try:
|
||||
response: Final = _Response.model_validate(response_obj, from_attributes=True)
|
||||
body: Final = _JSON_OBJECT.validate_json(wire.content)
|
||||
headers: Final = _HEADERS.validate_python(wire.headers)
|
||||
except (ValidationError, RuntimeError, httpx.RequestNotRead):
|
||||
return None
|
||||
if (
|
||||
wire.url.scheme != "https"
|
||||
or wire.url.host != "api.anthropic.com"
|
||||
or wire.url.path != "/v1/messages"
|
||||
or wire.url.query
|
||||
or wire.url.port not in (None, 443)
|
||||
):
|
||||
return None
|
||||
if (
|
||||
frozenset(headers) - _NATIVE_HEADERS
|
||||
or not supported_prediction_headers(headers)
|
||||
or headers.get("anthropic-version") != DEFAULT_ANTHROPIC_API_VERSION
|
||||
):
|
||||
return None
|
||||
provider_key: Final = headers.get("x-api-key", "")
|
||||
model: Final = body.get("model")
|
||||
if not _supported_provider_key(provider_key) or not isinstance(model, str) or model != response.model:
|
||||
return None
|
||||
prefix: Final = parse_prompt(body)
|
||||
if prefix is None:
|
||||
return None
|
||||
usage: Final = response.usage.prompt_tokens_details
|
||||
cache_tokens: Final = usage.cached_tokens + usage.cache_creation_tokens
|
||||
if cache_tokens <= 0 or cache_tokens > response.usage.prompt_tokens:
|
||||
return None
|
||||
split: Final = usage.cache_creation_token_details
|
||||
if usage.cache_creation_tokens and split is None:
|
||||
return None
|
||||
if split is not None and (
|
||||
split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens != usage.cache_creation_tokens
|
||||
or (prefix.ttl_seconds == 300 and split.ephemeral_1h_input_tokens > 0)
|
||||
or (prefix.ttl_seconds == 3600 and split.ephemeral_5m_input_tokens > 0)
|
||||
):
|
||||
return None
|
||||
return ObservedCachePrefix(
|
||||
prefix=prefix,
|
||||
scope=cache_scope(caller_key_hash, deployment_id, provider_key, model),
|
||||
cached_tokens=cache_tokens,
|
||||
cache_creation_tokens=usage.cache_creation_tokens,
|
||||
)
|
||||
|
|
@ -887,6 +887,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/auto_router/validate_complexity_router_config",
|
||||
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
|
||||
"/auto_router/session",
|
||||
"/cost/predict-cache",
|
||||
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
|
||||
# inside agent_endpoints/endpoints.py
|
||||
*agent_management_routes,
|
||||
|
|
|
|||
|
|
@ -895,6 +895,7 @@ async def common_checks(
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
llm_router=llm_router,
|
||||
request=request,
|
||||
team_id=valid_token.team_id if valid_token is not None else None,
|
||||
)
|
||||
|
||||
skip_all_budget_checks: Final = skip_budget_checks or (
|
||||
|
|
@ -4471,7 +4472,7 @@ async def stamp_matched_model_access_groups(
|
|||
|
||||
async def can_key_call_model(
|
||||
model: str | list[str],
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> Literal[True]:
|
||||
|
|
@ -4518,7 +4519,7 @@ async def can_key_call_model(
|
|||
|
||||
async def can_key_call_resolved_model(
|
||||
model: str,
|
||||
llm_model_list: list | None,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me
|
|||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
|
||||
)
|
||||
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
|
||||
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams
|
||||
|
||||
|
||||
|
|
@ -1736,7 +1736,7 @@ def _append_model_candidates(candidates: list[str], value: Any) -> None:
|
|||
candidates.extend(model for model in model_names if model)
|
||||
|
||||
|
||||
def _dedupe_model_candidates(candidates: list[str]) -> list[str]:
|
||||
def _dedupe_model_candidates(candidates: Collection[str]) -> list[str]:
|
||||
deduped: Final[list[str]] = []
|
||||
for model in candidates:
|
||||
if model not in deduped:
|
||||
|
|
@ -1845,13 +1845,42 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non
|
|||
return model_id
|
||||
|
||||
|
||||
def get_cache_prediction_deployments(
|
||||
*, current_deployment_id: str, candidate_deployment_id: str, llm_router: Router, team_id: str | None
|
||||
) -> tuple[Deployment, Deployment] | None:
|
||||
current: Final = llm_router.get_deployment(current_deployment_id)
|
||||
candidate: Final = llm_router.get_deployment(candidate_deployment_id)
|
||||
if current is None or candidate is None:
|
||||
return None
|
||||
if any(deployment.model_info.team_id not in (None, team_id) for deployment in (current, candidate)):
|
||||
return None
|
||||
return current, candidate
|
||||
|
||||
|
||||
def _cache_prediction_model_candidates(
|
||||
request_data: Mapping[str, object], llm_router: Router | None, team_id: str | None
|
||||
) -> tuple[str, ...]:
|
||||
current_id: Final = request_data.get("current_deployment_id")
|
||||
candidate_id: Final = request_data.get("candidate_deployment_id")
|
||||
if llm_router is None or not isinstance(current_id, str) or not isinstance(candidate_id, str):
|
||||
return ()
|
||||
deployments: Final = get_cache_prediction_deployments(
|
||||
current_deployment_id=current_id, candidate_deployment_id=candidate_id, llm_router=llm_router, team_id=team_id
|
||||
)
|
||||
return tuple(deployment.model_name for deployment in deployments) if deployments is not None else ()
|
||||
|
||||
|
||||
def _extract_model_candidates_from_request(
|
||||
request_data: dict,
|
||||
route: str,
|
||||
request_headers: Mapping[str, object] | None = None,
|
||||
request_query_params: Mapping[str, object] | None = None,
|
||||
llm_router: Router | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> list[str]:
|
||||
if route == "/cost/predict-cache":
|
||||
prediction_models: Final = _cache_prediction_model_candidates(request_data, llm_router, team_id) # pyright: ignore[reportUnknownArgumentType] # the typed reader validates each deployment ID from this legacy payload
|
||||
return _dedupe_model_candidates(prediction_models)
|
||||
candidates: Final[list[str]] = []
|
||||
uses_model_routing_sources: Final = _route_uses_model_routing_sources(route=route)
|
||||
uses_header_or_query_model_sources: Final = _route_matches_any_marker(
|
||||
|
|
@ -1945,6 +1974,7 @@ def get_model_from_request(
|
|||
request_query_params: Mapping[str, object] | None = None,
|
||||
llm_router: Router | None = None,
|
||||
request: Request | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> str | list[str] | None:
|
||||
"""Resolve the model(s) a request targets, for model-access and budget checks.
|
||||
|
||||
|
|
@ -1967,6 +1997,7 @@ def get_model_from_request(
|
|||
request_headers=request_headers,
|
||||
request_query_params=request_query_params,
|
||||
llm_router=llm_router,
|
||||
team_id=team_id,
|
||||
)
|
||||
model = _format_model_candidates(candidates)
|
||||
|
||||
|
|
|
|||
|
|
@ -191,6 +191,7 @@ def _get_model_from_request_context(
|
|||
route: str,
|
||||
request: Request | None,
|
||||
llm_router: Any | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> str | list[str] | None:
|
||||
return get_model_from_request(
|
||||
request_data=request_data,
|
||||
|
|
@ -199,6 +200,7 @@ def _get_model_from_request_context(
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
llm_router=llm_router,
|
||||
request=request,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -217,7 +219,7 @@ async def _normalize_claude_model(
|
|||
return
|
||||
if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True:
|
||||
return
|
||||
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router)
|
||||
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router, valid_token.team_id)
|
||||
if not isinstance(requested, str) or requested != request_data.get("model"):
|
||||
return
|
||||
if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"):
|
||||
|
|
@ -1652,6 +1654,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
skip_budget_checks = False
|
||||
if model is not None and llm_router is not None:
|
||||
|
|
@ -1692,6 +1695,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
|
@ -2091,6 +2095,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
skip_budget_checks = False
|
||||
if model is not None and llm_router is not None:
|
||||
|
|
@ -2209,6 +2214,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
|
|
@ -2239,6 +2245,7 @@ async def _user_api_key_auth_builder(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
|
|
@ -2734,6 +2741,7 @@ async def _run_centralized_common_checks(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=user_api_key_auth_obj.team_id,
|
||||
)
|
||||
|
||||
# Pin the metadata variable name (litellm_metadata vs metadata) before
|
||||
|
|
@ -2850,12 +2858,14 @@ def _should_skip_budget_checks(
|
|||
route: str,
|
||||
request: Request | None,
|
||||
llm_router: Any | None,
|
||||
team_id: str | None = None,
|
||||
) -> bool:
|
||||
model: Final = _get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=team_id,
|
||||
)
|
||||
if model is not None and llm_router is not None:
|
||||
return _is_model_cost_zero(model=model, llm_router=llm_router)
|
||||
|
|
@ -3301,6 +3311,7 @@ async def _enforce_key_and_fallback_model_access(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
|
||||
if model is not None:
|
||||
|
|
@ -3408,6 +3419,7 @@ async def _run_post_custom_auth_checks(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
|
|
@ -3449,6 +3461,7 @@ async def _run_post_custom_auth_checks(
|
|||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
|
|
|
|||
91
litellm/proxy/common_utils/prompt_cache_pricing.py
Normal file
91
litellm/proxy/common_utils/prompt_cache_pricing.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
from collections.abc import Mapping
|
||||
from math import isfinite
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.cost_calculator import (
|
||||
_select_model_name_for_cost_calc, # pyright: ignore[reportPrivateUsage] # shares completion_cost's deployment tariff selection
|
||||
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
|
||||
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _valid_price(value: object) -> bool:
|
||||
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
|
||||
|
||||
|
||||
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets) -> bool:
|
||||
required: Final = (
|
||||
("input_cost_per_token", True),
|
||||
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
|
||||
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
|
||||
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
|
||||
)
|
||||
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
|
||||
return False
|
||||
return all(
|
||||
_valid_price(value)
|
||||
for key, value in prices.items()
|
||||
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
|
||||
)
|
||||
|
||||
|
||||
def price_cache_tokens(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> float | None:
|
||||
try:
|
||||
selected_model: Final = _select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=None,
|
||||
custom_pricing=True,
|
||||
custom_llm_provider="anthropic",
|
||||
router_model_id=deployment_id,
|
||||
)
|
||||
if selected_model is None:
|
||||
return None
|
||||
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
|
||||
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
|
||||
price_entry: Final = registry.get(model_info["key"])
|
||||
if price_entry is None:
|
||||
return None
|
||||
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
|
||||
if not _has_required_prices(prices, tokens):
|
||||
return None
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=tokens.total_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=tokens.total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=tokens.cache_read_input_tokens,
|
||||
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
|
||||
cache_creation_token_details=CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
|
||||
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
|
||||
),
|
||||
),
|
||||
)
|
||||
logging_obj: Final = Logging(
|
||||
model=model,
|
||||
messages=[], # mutable-ok: Logging requires a list
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=None,
|
||||
litellm_call_id="prompt-cache-prediction",
|
||||
function_id="prompt-cache-prediction",
|
||||
)
|
||||
completion_cost(
|
||||
completion_response=ModelResponse(model=model, usage=usage),
|
||||
model=model,
|
||||
custom_llm_provider="anthropic",
|
||||
custom_pricing=True,
|
||||
router_model_id=deployment_id,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
cost: Final = logging_obj.cost_breakdown.get("input_cost") if logging_obj.cost_breakdown is not None else None
|
||||
return cost if cost is not None and _valid_price(cost) else None
|
||||
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
|
||||
return None
|
||||
|
|
@ -9,6 +9,7 @@ from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
|
|||
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
|
||||
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
|
||||
from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
|
||||
from .prompt_cache_prediction import PromptCacheObserver
|
||||
from .responses_id_security import ResponsesIDSecurity
|
||||
from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
|
||||
|
||||
|
|
@ -25,6 +26,7 @@ PROXY_HOOKS: Final = {
|
|||
"max_iterations_limiter": _PROXY_MaxIterationsHandler,
|
||||
"max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler,
|
||||
"sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler,
|
||||
"prompt_cache_prediction": PromptCacheObserver,
|
||||
}
|
||||
|
||||
## FEATURE FLAG HOOKS ##
|
||||
|
|
|
|||
|
|
@ -9,10 +9,12 @@ import binascii
|
|||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence, Set
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
|
||||
from contextlib import asynccontextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
|
|
@ -23,6 +25,7 @@ from typing import (
|
|||
TypedDict,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly
|
||||
|
||||
from litellm import DualCache
|
||||
|
|
@ -84,6 +87,9 @@ else:
|
|||
InternalUsageCache = Any
|
||||
|
||||
|
||||
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
BATCH_RATE_LIMITER_SCRIPT: Final = """
|
||||
local results = {}
|
||||
local now = tonumber(ARGV[1])
|
||||
|
|
@ -2673,12 +2679,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
Returns list of descriptors for API key, user, team, team member, end user,
|
||||
model-specific, agent, and agent-session limits.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
get_team_model_rpm_limit,
|
||||
get_team_model_tpm_limit,
|
||||
)
|
||||
|
||||
descriptors: Final = []
|
||||
descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place
|
||||
|
||||
# API Key rate limits
|
||||
if user_api_key_dict.api_key and (
|
||||
|
|
@ -2803,34 +2804,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
if (
|
||||
get_team_model_rpm_limit(user_api_key_dict) is not None
|
||||
or get_team_model_tpm_limit(user_api_key_dict) is not None
|
||||
):
|
||||
_tpm_limit_for_team_model: Final = get_team_model_tpm_limit(user_api_key_dict) or {}
|
||||
_rpm_limit_for_team_model: Final = get_team_model_rpm_limit(user_api_key_dict) or {}
|
||||
should_check_rate_limit = False
|
||||
if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model:
|
||||
should_check_rate_limit = True
|
||||
|
||||
if should_check_rate_limit:
|
||||
model_specific_tpm_limit = None
|
||||
model_specific_rpm_limit = None
|
||||
if requested_model in _tpm_limit_for_team_model:
|
||||
model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model]
|
||||
if requested_model in _rpm_limit_for_team_model:
|
||||
model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model]
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
self._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model if isinstance(requested_model, str) else None,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Agent-level and session-level rate limits
|
||||
resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data)
|
||||
|
|
@ -3416,6 +3394,108 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
requested_model,
|
||||
)
|
||||
|
||||
async def _build_request_rate_limit_descriptors(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: Mapping[str, object],
|
||||
call_type: str | None,
|
||||
) -> list[RateLimitDescriptor]: # mutable-ok: the shared generation reservation helpers require a list
|
||||
metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(
|
||||
user_api_key_dict.metadata or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # validates the legacy auth metadata boundary
|
||||
)
|
||||
rpm_value: Final = metadata.get("rpm_limit_type")
|
||||
tpm_value: Final = metadata.get("tpm_limit_type")
|
||||
rpm_limit_type: Final = rpm_value if isinstance(rpm_value, str) else None
|
||||
tpm_limit_type: Final = tpm_value if isinstance(tpm_value, str) else None
|
||||
model_value: Final = data.get("model")
|
||||
requested_model: Final = model_value if isinstance(model_value, str) else None
|
||||
model_has_failures: Final = (
|
||||
await self._check_model_has_recent_failures(
|
||||
model=requested_model,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
if requested_model and self._is_dynamic_rate_limiting_enabled(rpm_limit_type, tpm_limit_type)
|
||||
else False
|
||||
)
|
||||
descriptors: Final = self._create_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # legacy helper reads a dictionary with validated keys
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=dict(data), # mutable-ok: legacy descriptor helpers accept a request dictionary
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
call_type=call_type,
|
||||
)
|
||||
self._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
self.add_project_io_token_rate_limit_descriptors_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
return [ # mutable-ok: the shared generation reservation helpers require a list
|
||||
*descriptors,
|
||||
*self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model),
|
||||
]
|
||||
|
||||
async def _release_request_capacity_when_admitted(
|
||||
self,
|
||||
admission: asyncio.Task[RateLimitResponse],
|
||||
acquisition: ParallelSlotAcquisition,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
response: Final = await admission
|
||||
if response["overall_code"] == "OK":
|
||||
await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span)
|
||||
|
||||
@asynccontextmanager
|
||||
async def request_capacity(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model: str,
|
||||
*,
|
||||
request_data: Mapping[str, object] | None = None,
|
||||
) -> AsyncGenerator[None, None]:
|
||||
"""Charge one non-generation provider request to RPM and hold its concurrency slot."""
|
||||
data: Final = MappingProxyType({**(request_data or MappingProxyType({})), "model": model})
|
||||
descriptors: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None)
|
||||
acquisition: Final = ParallelSlotAcquisition(
|
||||
slot_id=uuid.uuid4().hex,
|
||||
counter_keys=[ # mutable-ok: the shared slot-release contract requires a list
|
||||
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
|
||||
for d in descriptors
|
||||
if d["rate_limit"] is not None and d["rate_limit"].get("max_parallel_requests") is not None
|
||||
],
|
||||
)
|
||||
admission: Final = asyncio.create_task(
|
||||
self.should_rate_limit(
|
||||
descriptors=descriptors,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
skip_tpm_check=True,
|
||||
parallel_slot_id=acquisition["slot_id"],
|
||||
)
|
||||
)
|
||||
try:
|
||||
response: Final = await asyncio.shield(admission)
|
||||
if response["overall_code"] == "OVER_LIMIT":
|
||||
self._handle_rate_limit_error(response, descriptors, model)
|
||||
yield
|
||||
finally:
|
||||
cleanup: Final = asyncio.create_task(
|
||||
self._release_request_capacity_when_admitted(admission, acquisition, user_api_key_dict)
|
||||
)
|
||||
cancellation: asyncio.CancelledError | None = None # rebind-ok: retain cancellation until cleanup finishes
|
||||
while not cleanup.done():
|
||||
try:
|
||||
await asyncio.shield(cleanup)
|
||||
except asyncio.CancelledError as exc:
|
||||
cancellation = exc # rebind-ok: retain the latest cancellation without interrupting slot release
|
||||
cleanup.result()
|
||||
if cancellation is not None:
|
||||
raise cancellation
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -3444,59 +3524,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
call_type=call_type,
|
||||
)
|
||||
|
||||
# Get rate limit types from metadata
|
||||
metadata: Final = user_api_key_dict.metadata or {}
|
||||
rpm_limit_type: Final = metadata.get("rpm_limit_type")
|
||||
tpm_limit_type: Final = metadata.get("tpm_limit_type")
|
||||
|
||||
# For dynamic mode, check if the model has recent failures
|
||||
model_has_failures = False
|
||||
requested_model: Final = data.get("model", None)
|
||||
|
||||
if (
|
||||
self._is_dynamic_rate_limiting_enabled(
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
)
|
||||
and requested_model
|
||||
):
|
||||
model_has_failures = await self._check_model_has_recent_failures(
|
||||
model=requested_model,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
|
||||
# Create rate limit descriptors
|
||||
descriptors: Final = self._create_rate_limit_descriptors(
|
||||
request_data: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
|
||||
model_value: Final = request_data.get("model")
|
||||
requested_model: Final = model_value if isinstance(model_value, str) else None
|
||||
descriptors: Final = await self._build_request_rate_limit_descriptors(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
rpm_limit_type=rpm_limit_type,
|
||||
tpm_limit_type=tpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
# Add team model rate limits from team_metadata
|
||||
self._add_team_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Project Level Rate Limits
|
||||
self._add_project_model_rate_limit_descriptor_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
self.add_project_io_token_rate_limit_descriptors_from_metadata(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
requested_model=requested_model,
|
||||
descriptors=descriptors,
|
||||
)
|
||||
|
||||
# Org Level Rate Limits
|
||||
descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model))
|
||||
|
||||
# Only check rate limits if we have descriptors with actual limits
|
||||
if descriptors:
|
||||
# First pass: RPM and max_parallel_requests sliding-window check.
|
||||
|
|
|
|||
142
litellm/proxy/hooks/prompt_cache_prediction.py
Normal file
142
litellm/proxy/hooks/prompt_cache_prediction.py
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, parse_observed_cache
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
|
||||
_RETENTION_SECONDS: Final = 86_400
|
||||
|
||||
|
||||
class CacheObservation(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
|
||||
|
||||
fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
|
||||
cached_tokens: int = Field(gt=0)
|
||||
observed_at: float = Field(ge=0, allow_inf_nan=False)
|
||||
expires_at: float = Field(ge=0, allow_inf_nan=False)
|
||||
|
||||
|
||||
_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
|
||||
|
||||
|
||||
def _cache_key(scope: str, fingerprint: str) -> str:
|
||||
return f"prompt-cache-observation:{scope}:{fingerprint}"
|
||||
|
||||
|
||||
async def lookup(
|
||||
cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
|
||||
) -> CacheObservation | None:
|
||||
checked_at: Final = time.time() if now is None else now
|
||||
exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
|
||||
if exact is not None and exact.expires_at > checked_at:
|
||||
return exact
|
||||
older: Final = await asyncio.gather(
|
||||
*(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
|
||||
)
|
||||
observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
|
||||
return next(
|
||||
(observation for observation in observations if observation.expires_at > checked_at),
|
||||
next(iter(observations), None),
|
||||
)
|
||||
|
||||
|
||||
async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
|
||||
try:
|
||||
value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validate the legacy cache's untyped result at the I/O boundary
|
||||
if value is None:
|
||||
return None
|
||||
observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
|
||||
except ValidationError:
|
||||
return None
|
||||
return observation if observation.fingerprint == fingerprint else None
|
||||
|
||||
|
||||
class _Metadata(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
user_api_key_hash: str = Field(min_length=1)
|
||||
|
||||
|
||||
class _Logged(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
status: Literal["success"]
|
||||
model_id: str = Field(min_length=1)
|
||||
metadata: _Metadata
|
||||
|
||||
|
||||
class _Event(BaseModel):
|
||||
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
|
||||
call_type: Literal["anthropic_messages"]
|
||||
custom_llm_provider: Literal["anthropic"]
|
||||
cache_hit: bool | None = None
|
||||
httpx_response: httpx.Response
|
||||
first_api_call_start_time: datetime
|
||||
standard_logging_object: _Logged
|
||||
stream: bool = False
|
||||
prompt_cache_response_complete: bool = False
|
||||
|
||||
|
||||
class PromptCacheObserver(CustomLogger):
|
||||
def __init__(self, internal_usage_cache: InternalUsageCache, clock: Callable[[], float] = time.time) -> None:
|
||||
super().__init__() # pyright: ignore[reportUnknownMemberType] # base callback constructor accepts untyped kwargs
|
||||
self.cache = internal_usage_cache.dual_cache
|
||||
self.clock = clock
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
if not isinstance(response_obj, ModelResponse):
|
||||
return
|
||||
try:
|
||||
event: Final = _Event.model_validate(kwargs)
|
||||
wire: Final = event.httpx_response.request
|
||||
except (ValidationError, RuntimeError, httpx.RequestNotRead):
|
||||
return
|
||||
if (
|
||||
event.cache_hit
|
||||
or event.httpx_response.status_code != 200
|
||||
or (event.stream and not event.prompt_cache_response_complete)
|
||||
):
|
||||
return
|
||||
observed: Final = parse_observed_cache(
|
||||
wire,
|
||||
response_obj,
|
||||
event.standard_logging_object.metadata.user_api_key_hash,
|
||||
event.standard_logging_object.model_id,
|
||||
)
|
||||
if observed is None:
|
||||
return
|
||||
prefix: Final = observed.prefix
|
||||
scope: Final = observed.scope
|
||||
cache_tokens: Final = observed.cached_tokens
|
||||
now: Final = self.clock()
|
||||
started: Final = event.first_api_call_start_time.timestamp()
|
||||
if started > now:
|
||||
return
|
||||
if observed.cache_creation_tokens == 0:
|
||||
previous: Final = await _read_exact(self.cache, scope, prefix.fingerprint)
|
||||
if previous is None or previous.fingerprint != prefix.fingerprint or previous.cached_tokens != cache_tokens:
|
||||
return
|
||||
observation: Final = CacheObservation(
|
||||
fingerprint=prefix.fingerprint,
|
||||
cached_tokens=cache_tokens,
|
||||
observed_at=now,
|
||||
expires_at=started + prefix.ttl_seconds,
|
||||
)
|
||||
key: Final = _cache_key(scope, prefix.fingerprint)
|
||||
payload: Final = observation.model_dump_json()
|
||||
await self.cache.async_set_cache(key, payload, ttl=_RETENTION_SECONDS) # pyright: ignore[reportUnknownMemberType] # legacy cache accepts a serialized validated observation
|
||||
if self.cache.redis_cache is not None:
|
||||
await self.cache.async_set_cache(key, payload, local_only=True, ttl=1) # pyright: ignore[reportUnknownMemberType] # keep the local copy short-lived while Redis retains stale evidence
|
||||
|
|
@ -28,6 +28,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.prompt_cache_prediction import router as prompt_cache_prediction_router
|
||||
from litellm.types.utils import (
|
||||
CostBreakdown,
|
||||
CostPerToken,
|
||||
|
|
@ -39,6 +40,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
router: Final = APIRouter()
|
||||
router.include_router(prompt_cache_prediction_router)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
278
litellm/proxy/management_endpoints/prompt_cache_prediction.py
Normal file
278
litellm/proxy/management_endpoints/prompt_cache_prediction.py
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
import time
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import current_billing_time, pinned_billing_time
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import (
|
||||
PromptPrefix,
|
||||
TokenCounter,
|
||||
UnsupportedPredictionTarget,
|
||||
cache_scope,
|
||||
count_prompt_tokens,
|
||||
parse_prompt,
|
||||
resolve_prediction_target,
|
||||
supported_prediction_headers,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
from litellm.proxy.auth.auth_utils import get_cache_prediction_deployments
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # canonical parsed-body owner; validate its legacy result at the endpoint boundary
|
||||
)
|
||||
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
_PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # use the configured proxy limiter's shared capacity owner
|
||||
)
|
||||
from litellm.proxy.hooks.prompt_cache_prediction import lookup
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import (
|
||||
CacheCostScenario,
|
||||
CacheEvidence,
|
||||
CachePredictionArm,
|
||||
CachePredictionRequest,
|
||||
CachePredictionResponse,
|
||||
CacheTokenBuckets,
|
||||
)
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.utils import get_prompt_cache_min_tokens
|
||||
|
||||
router: Final = APIRouter()
|
||||
_REQUEST_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class _CallerSettings(BaseModel):
|
||||
config: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def has_request_transforms() -> bool:
|
||||
from litellm.proxy.hooks import PROXY_HOOKS
|
||||
|
||||
builtins: Final = frozenset(PROXY_HOOKS.values())
|
||||
hooks: Final = ("async_pre_call_hook", "async_pre_request_hook", "async_pre_call_deployment_hook")
|
||||
callbacks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomLogger)
|
||||
return any(
|
||||
type(callback) not in builtins
|
||||
and any(getattr(type(callback), hook) is not getattr(CustomLogger, hook) for hook in hooks)
|
||||
for callback in callbacks
|
||||
)
|
||||
|
||||
|
||||
def _buckets(prefix_tokens: int, suffix_tokens: int, read_tokens: int, ttl_seconds: int) -> CacheTokenBuckets:
|
||||
return CacheTokenBuckets(
|
||||
uncached_input_tokens=suffix_tokens,
|
||||
cache_read_input_tokens=read_tokens,
|
||||
cache_creation_5m_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 300 else 0,
|
||||
cache_creation_1h_input_tokens=prefix_tokens - read_tokens if ttl_seconds == 3600 else 0,
|
||||
)
|
||||
|
||||
|
||||
def _scenario(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> CacheCostScenario | None:
|
||||
cost: Final = price_cache_tokens(model=model, deployment_id=deployment_id, tokens=tokens)
|
||||
return CacheCostScenario(tokens=tokens, input_cost=cost) if cost is not None else None
|
||||
|
||||
|
||||
def _capacity_counter(
|
||||
limiter: _PROXY_MaxParallelRequestsHandler_v3,
|
||||
caller: UserAPIKeyAuth,
|
||||
model_name: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> TokenCounter:
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
async with limiter.request_capacity(caller, model_name, request_data=request_data):
|
||||
return await count_prompt_tokens(model, api_key, body)
|
||||
|
||||
return count
|
||||
|
||||
|
||||
def _capacity_request_data(
|
||||
http_request: Request, caller: UserAPIKeyAuth, request_data: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
# The parsed-body cache retains only original top-level keys. Replay the
|
||||
# shared idempotent tag merges on limiter-only data when auth added metadata.
|
||||
data: Final = dict(request_data) # mutable-ok: the existing tag merge owners accept a dictionary out-param
|
||||
LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(http_request, data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner takes the validated capacity dictionary
|
||||
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(data, caller) # pyright: ignore[reportUnknownMemberType] # legacy tag owner merges trusted key tags into capacity metadata
|
||||
return MappingProxyType(data)
|
||||
|
||||
|
||||
async def predict_arm(
|
||||
deployment: Deployment,
|
||||
body: Mapping[str, JsonValue],
|
||||
prefix: PromptPrefix,
|
||||
caller_key_hash: str,
|
||||
cache: DualCache,
|
||||
token_counter: TokenCounter,
|
||||
) -> CachePredictionArm:
|
||||
deployment_id: Final = deployment.model_info.id or ""
|
||||
params: Final = deployment.litellm_params
|
||||
unknown: Final = CachePredictionArm(deployment_id=deployment_id, model=params.model)
|
||||
if deployment.model_info.blocked:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "unsupported_deployment_configuration"}))
|
||||
target: Final = resolve_prediction_target(params)
|
||||
if isinstance(target, UnsupportedPredictionTarget):
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": target.reason}))
|
||||
model: Final = target.model
|
||||
api_key: Final = target.api_key
|
||||
total_count: Final = await token_counter(model, api_key, body)
|
||||
prefix_count: Final = await token_counter(model, api_key, prefix.prefix_body)
|
||||
if total_count is None or prefix_count is None or total_count < prefix_count:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "token_count_unavailable"}))
|
||||
scope: Final = cache_scope(caller_key_hash, deployment_id, api_key, model)
|
||||
observation: Final = await lookup(cache, scope, prefix)
|
||||
exact: Final = observation is not None and observation.fingerprint == prefix.fingerprint
|
||||
cacheable: Final = observation.cached_tokens if exact and observation is not None else prefix_count
|
||||
if cacheable > total_count or (observation is not None and observation.cached_tokens > cacheable):
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "inconsistent_prefix_token_count"}))
|
||||
suffix: Final = total_count - cacheable
|
||||
evidence: Final = (
|
||||
CacheEvidence(observed_at=observation.observed_at, expires_at=observation.expires_at)
|
||||
if observation is not None
|
||||
else None
|
||||
)
|
||||
if cacheable < get_prompt_cache_min_tokens(params.model):
|
||||
disabled: Final = _scenario(model, deployment_id, CacheTokenBuckets(uncached_input_tokens=total_count))
|
||||
if disabled is None:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
|
||||
return CachePredictionArm(
|
||||
deployment_id=deployment_id,
|
||||
model=model,
|
||||
cache_state="disabled",
|
||||
reason="below_cache_minimum",
|
||||
estimate=disabled,
|
||||
cold=disabled,
|
||||
warm=disabled,
|
||||
token_count_source="anthropic_count_tokens",
|
||||
)
|
||||
fresh: Final = observation is not None and observation.expires_at > time.time()
|
||||
read: Final = observation.cached_tokens if fresh and observation is not None else 0
|
||||
with pinned_billing_time(current_billing_time()):
|
||||
cold: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, 0, prefix.ttl_seconds))
|
||||
warm: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, cacheable, prefix.ttl_seconds))
|
||||
estimate: Final = _scenario(model, deployment_id, _buckets(cacheable, suffix, read, prefix.ttl_seconds))
|
||||
if cold is None or warm is None or estimate is None:
|
||||
return unknown.model_copy(update=MappingProxyType({"reason": "pricing_unavailable"}))
|
||||
return CachePredictionArm(
|
||||
deployment_id=deployment_id,
|
||||
model=model,
|
||||
cache_state="warm" if fresh and exact else "partial" if fresh else "stale" if observation else "unknown",
|
||||
reason=None if fresh else "observation_expired" if observation else "no_compatible_observation",
|
||||
estimate=estimate,
|
||||
cold=cold,
|
||||
warm=warm,
|
||||
evidence=evidence,
|
||||
token_count_source="anthropic_count_tokens",
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/cost/predict-cache",
|
||||
tags=["Cost Tracking"], # mutable-ok: FastAPI requires a list for OpenAPI tags
|
||||
response_model=CachePredictionResponse,
|
||||
)
|
||||
async def predict_cache_cost(
|
||||
request: CachePredictionRequest,
|
||||
http_request: Request,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> CachePredictionResponse:
|
||||
"""Compare the next native Anthropic request on two configured deployment IDs.
|
||||
|
||||
Estimates use provider token counting and recent successful cache telemetry for this key.
|
||||
Unknown cache state uses the cold scenario when prices/counts are available. Cache observations
|
||||
do not guarantee retention. v0 supports one message-content breakpoint, text and client tools;
|
||||
system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and
|
||||
request transforms are unknown.
|
||||
Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses
|
||||
up to four counts. The legacy rate limiter returns unknown without contacting the provider.
|
||||
This endpoint does not generate tokens, prewarm caches, choose a model or alter routing.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
|
||||
|
||||
if llm_router is None:
|
||||
raise HTTPException(status_code=503, detail="Model router is unavailable")
|
||||
deployments: Final = get_cache_prediction_deployments(
|
||||
current_deployment_id=request.current_deployment_id,
|
||||
candidate_deployment_id=request.candidate_deployment_id,
|
||||
llm_router=llm_router,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
if deployments is None:
|
||||
raise HTTPException(status_code=404, detail="Deployment not found")
|
||||
current, candidate = deployments
|
||||
for deployment in (current, candidate):
|
||||
await can_key_call_resolved_model(
|
||||
model=deployment.model_name,
|
||||
llm_model_list=llm_router.get_model_list(),
|
||||
valid_token=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
prefix: Final = parse_prompt(request.request)
|
||||
caller: Final = user_api_key_dict.api_key
|
||||
caller_settings: Final = _CallerSettings.model_validate(user_api_key_dict, from_attributes=True)
|
||||
unsupported_transform: Final = bool(caller_settings.config) or has_request_transforms()
|
||||
unsupported_headers: Final = not supported_prediction_headers(http_request.headers)
|
||||
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
|
||||
if (
|
||||
prefix is None
|
||||
or not caller
|
||||
or unsupported_transform
|
||||
or unsupported_headers
|
||||
or not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3)
|
||||
):
|
||||
reason: Final = (
|
||||
"unsupported_provider_headers"
|
||||
if unsupported_headers
|
||||
else "unsupported_request_transform"
|
||||
if unsupported_transform
|
||||
else "unsupported_prompt_shape"
|
||||
if prefix is None
|
||||
else "caller_identity_unavailable"
|
||||
if not caller
|
||||
else "limiter_unavailable"
|
||||
)
|
||||
return CachePredictionResponse(
|
||||
stay=CachePredictionArm(deployment_id=request.current_deployment_id, reason=reason),
|
||||
switch=CachePredictionArm(deployment_id=request.candidate_deployment_id, reason=reason),
|
||||
switch_delta=None,
|
||||
cache_rebuild_penalty=None,
|
||||
)
|
||||
request_data: Final = _capacity_request_data(
|
||||
http_request, user_api_key_dict, _REQUEST_DATA.validate_python(await _read_request_body(http_request))
|
||||
)
|
||||
stay: Final = await predict_arm(
|
||||
current,
|
||||
request.request,
|
||||
prefix,
|
||||
caller,
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache,
|
||||
_capacity_counter(limiter, user_api_key_dict, current.model_name, request_data),
|
||||
)
|
||||
switch: Final = (
|
||||
stay
|
||||
if current.model_info.id == candidate.model_info.id
|
||||
else await predict_arm(
|
||||
candidate,
|
||||
request.request,
|
||||
prefix,
|
||||
caller,
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache,
|
||||
_capacity_counter(limiter, user_api_key_dict, candidate.model_name, request_data),
|
||||
)
|
||||
)
|
||||
return CachePredictionResponse(
|
||||
stay=stay,
|
||||
switch=switch,
|
||||
switch_delta=(switch.estimate.input_cost - stay.estimate.input_cost)
|
||||
if switch.estimate is not None and stay.estimate is not None
|
||||
else None,
|
||||
cache_rebuild_penalty=(switch.estimate.input_cost - switch.warm.input_cost)
|
||||
if switch.estimate is not None and switch.warm is not None
|
||||
else None,
|
||||
)
|
||||
|
|
@ -270,6 +270,24 @@ class PassThroughStreamingHandler:
|
|||
- Vertex AI
|
||||
- OpenAI
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
_is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection
|
||||
_is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence
|
||||
)
|
||||
|
||||
# Transport reads can split event names and JSON payloads. Recognize terminal
|
||||
# events only after the shared SSE framer has reassembled the collected bytes.
|
||||
complete_frames, incomplete_tail = split_complete_sse_frames(
|
||||
b"".join(raw_bytes) if endpoint_type == EndpointType.ANTHROPIC else b""
|
||||
)
|
||||
litellm_logging_obj.model_call_details[ # rebind-ok: stamp evidence on the per-request state read by callbacks
|
||||
"prompt_cache_response_complete"
|
||||
] = (
|
||||
endpoint_type == EndpointType.ANTHROPIC
|
||||
and not incomplete_tail.strip()
|
||||
and _is_message_stop_chunk(complete_frames)
|
||||
and not _is_provider_error_chunk(complete_frames)
|
||||
)
|
||||
try:
|
||||
(
|
||||
standard_logging_response_object,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,67 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Annotated, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt
|
||||
|
||||
TokenCount: TypeAlias = Annotated[StrictInt, Field(ge=0)]
|
||||
|
||||
|
||||
class CacheTokenBuckets(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
uncached_input_tokens: TokenCount = 0
|
||||
cache_read_input_tokens: TokenCount = 0
|
||||
cache_creation_5m_input_tokens: TokenCount = 0
|
||||
cache_creation_1h_input_tokens: TokenCount = 0
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
return (
|
||||
self.uncached_input_tokens
|
||||
+ self.cache_read_input_tokens
|
||||
+ self.cache_creation_5m_input_tokens
|
||||
+ self.cache_creation_1h_input_tokens
|
||||
)
|
||||
|
||||
|
||||
class CacheEvidence(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
observed_at: float
|
||||
expires_at: float
|
||||
source: Literal["provider_usage"] = "provider_usage"
|
||||
confidence: Literal["observed"] = "observed"
|
||||
|
||||
|
||||
class CacheCostScenario(BaseModel):
|
||||
tokens: CacheTokenBuckets
|
||||
input_cost: float
|
||||
|
||||
|
||||
class CachePredictionArm(BaseModel):
|
||||
deployment_id: str
|
||||
model: str | None = None
|
||||
cache_state: Literal["warm", "partial", "stale", "unknown", "disabled"] = "unknown"
|
||||
reason: str | None = None
|
||||
estimate: CacheCostScenario | None = None
|
||||
cold: CacheCostScenario | None = None
|
||||
warm: CacheCostScenario | None = None
|
||||
evidence: CacheEvidence | None = None
|
||||
token_count_source: Literal["anthropic_count_tokens"] | None = None
|
||||
|
||||
|
||||
class CachePredictionRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
current_deployment_id: str = Field(min_length=1, max_length=256)
|
||||
candidate_deployment_id: str = Field(min_length=1, max_length=256)
|
||||
request: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
class CachePredictionResponse(BaseModel):
|
||||
stay: CachePredictionArm
|
||||
switch: CachePredictionArm
|
||||
switch_delta: float | None
|
||||
cache_rebuild_penalty: float | None
|
||||
pricing_basis: Literal["input_before_discounts_and_margins"] = "input_before_discounts_and_margins"
|
||||
cache_guarantee: Literal[False] = False
|
||||
|
|
@ -0,0 +1,209 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.llms.anthropic.count_tokens import handler as count_handler
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import (
|
||||
NativePredictionTarget,
|
||||
cache_scope,
|
||||
count_prompt_tokens,
|
||||
parse_observed_cache,
|
||||
parse_prompt,
|
||||
resolve_prediction_target,
|
||||
supported_prediction_headers,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.models.credentials import CredentialItem
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver, lookup
|
||||
from litellm.proxy.management_endpoints.prompt_cache_prediction import predict_arm
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
_MODEL: Final = "claude-sonnet-5"
|
||||
_KEY: Final = "test-provider-key"
|
||||
_CALLER: Final = "test-caller-hash"
|
||||
_DEPLOYMENT: Final = "test-native-deployment"
|
||||
|
||||
|
||||
def _body() -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": _MODEL,
|
||||
"system": "Keep this context",
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}],
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "A cacheable prefix", "cache_control": {"type": "ephemeral"}}
|
||||
]}],
|
||||
}
|
||||
|
||||
|
||||
@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:
|
||||
cache: Final = DualCache()
|
||||
observer: Final = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: 1010.0)
|
||||
body: Final = _body()
|
||||
prefix: Final = parse_prompt(body)
|
||||
assert prefix is not None
|
||||
headers: Final = {"x-api-key": _KEY, **({"anthropic-version": version} if version is not None else {})}
|
||||
wire: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=body)
|
||||
response: Final = ModelResponse(
|
||||
model=_MODEL,
|
||||
usage=Usage(
|
||||
prompt_tokens=311,
|
||||
completion_tokens=2,
|
||||
total_tokens=313,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=100,
|
||||
cache_creation_tokens=200,
|
||||
cache_creation_token_details=CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=200, ephemeral_1h_input_tokens=0
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
await observer.async_log_success_event(
|
||||
{
|
||||
"call_type": "anthropic_messages",
|
||||
"custom_llm_provider": "anthropic",
|
||||
"httpx_response": httpx.Response(200, request=wire),
|
||||
"first_api_call_start_time": datetime.fromtimestamp(1000.0),
|
||||
"standard_logging_object": {
|
||||
"status": "success", "model_id": _DEPLOYMENT,
|
||||
"metadata": {"user_api_key_hash": _CALLER},
|
||||
},
|
||||
},
|
||||
response,
|
||||
datetime.fromtimestamp(1010.0),
|
||||
datetime.fromtimestamp(1010.0),
|
||||
)
|
||||
default_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL)
|
||||
found: Final = await lookup(cache, default_scope, prefix, now=1010.0)
|
||||
assert (found is not None) == (version == DEFAULT_ANTHROPIC_API_VERSION)
|
||||
if version != DEFAULT_ANTHROPIC_API_VERSION:
|
||||
other_scope: Final = cache_scope(_CALLER, _DEPLOYMENT, _KEY, _MODEL, version or "")
|
||||
assert await lookup(cache, other_scope, prefix, now=1010.0) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers, supported", [
|
||||
({}, True),
|
||||
({"Anthropic-Version": DEFAULT_ANTHROPIC_API_VERSION}, True),
|
||||
({"anthropic-version": "2099-01-01"}, False),
|
||||
({"Anthropic-Beta": ""}, False),
|
||||
({"anthropic-beta": "future-feature"}, False),
|
||||
])
|
||||
def test_prediction_header_eligibility(headers: Mapping[str, str], supported: bool) -> None:
|
||||
assert supported_prediction_headers(headers) is supported
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_count_uses_same_version_and_preserves_native_input(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
body: Final = _body()
|
||||
requests: Final[list[httpx.Request]] = []
|
||||
|
||||
def provider(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(200, json={"input_tokens": 311})
|
||||
|
||||
client: Final = AsyncHTTPHandler()
|
||||
await client.client.aclose()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
|
||||
monkeypatch.setattr(count_handler, "get_async_httpx_client", lambda **kwargs: client)
|
||||
try:
|
||||
assert await count_prompt_tokens(_MODEL, _KEY, body) == 311
|
||||
finally:
|
||||
await client.client.aclose()
|
||||
assert len(requests) == 1
|
||||
assert requests[0].headers["anthropic-version"] == DEFAULT_ANTHROPIC_API_VERSION
|
||||
assert requests[0].url == "https://api.anthropic.com/v1/messages/count_tokens"
|
||||
assert json.loads(requests[0].content) == body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["static", "database"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_environment_credential_matches_native_count_and_observed_scope(
|
||||
source: str, monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("LIT7658_PROVIDER_KEY", _KEY)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
|
||||
params: Final = {
|
||||
"model": f"anthropic/{_MODEL}", "api_key": "os.environ/LIT7658_PROVIDER_KEY",
|
||||
"api_base": "https://api.anthropic.com",
|
||||
}
|
||||
router: Final = litellm.Router(model_list=[{
|
||||
"model_name": "test-native", "litellm_params": dict(params), "model_info": {"id": _DEPLOYMENT},
|
||||
}] if source == "static" else [], num_retries=0)
|
||||
if source == "database":
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
assert proxy_server.ProxyConfig()._add_deployment([SimpleNamespace(
|
||||
model_id=_DEPLOYMENT, model_name="test-native", model_info={}, litellm_params=dict(params),
|
||||
)]) == 1
|
||||
deployment: Final = router.get_deployment(_DEPLOYMENT)
|
||||
assert deployment is not None
|
||||
target: Final = resolve_prediction_target(deployment.litellm_params)
|
||||
assert isinstance(target, NativePredictionTarget)
|
||||
body: Final = _body()
|
||||
with respx.mock() as upstream:
|
||||
native: Final = upstream.post("https://api.anthropic.com/v1/messages").respond(200, json={
|
||||
"id": "msg_test", "type": "message", "role": "assistant", "model": _MODEL,
|
||||
"content": [{"type": "text", "text": "Hello"}], "stop_reason": "end_turn", "stop_sequence": None,
|
||||
"usage": {"input_tokens": 11, "output_tokens": 1, "cache_read_input_tokens": 300},
|
||||
})
|
||||
counter: Final = upstream.post("https://api.anthropic.com/v1/messages/count_tokens").respond(
|
||||
200, json={"input_tokens": 311},
|
||||
)
|
||||
await router.aanthropic_messages(
|
||||
model="test-native", max_tokens=1, **{key: value for key, value in body.items() if key != "model"},
|
||||
)
|
||||
assert await count_prompt_tokens(target.model, target.api_key, body) == 311
|
||||
assert native.call_count == counter.call_count == 1
|
||||
assert native.calls.last.request.headers["x-api-key"] == counter.calls.last.request.headers["x-api-key"] == _KEY
|
||||
observed: Final = parse_observed_cache(native.calls.last.request, ModelResponse(
|
||||
model=_MODEL, usage=Usage(
|
||||
prompt_tokens=311, completion_tokens=1, total_tokens=312,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300),
|
||||
),
|
||||
), _CALLER, _DEPLOYMENT)
|
||||
assert observed is not None
|
||||
assert observed.scope == cache_scope(_CALLER, _DEPLOYMENT, target.api_key, target.model)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("inline_key", [None, _KEY])
|
||||
@pytest.mark.asyncio
|
||||
async def test_named_credential_is_explicitly_unsupported_before_count(
|
||||
inline_key: str | None, monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [CredentialItem(
|
||||
credential_name="test-named", credential_info={}, credential_values={"api_key": "test-named-provider-key"},
|
||||
)])
|
||||
deployment: Final = Deployment(
|
||||
model_name="test-native",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model=f"anthropic/{_MODEL}", api_key=inline_key, litellm_credential_name="test-named",
|
||||
),
|
||||
model_info=ModelInfo(id=_DEPLOYMENT),
|
||||
)
|
||||
body: Final = _body()
|
||||
prefix: Final = parse_prompt(body)
|
||||
assert prefix is not None
|
||||
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
pytest.fail("Unsupported named credentials must not reach provider counting")
|
||||
|
||||
arm: Final = await predict_arm(deployment, body, prefix, _CALLER, DualCache(), count)
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "unsupported_deployment_configuration"
|
||||
assert arm.estimate is None and arm.cold is None and arm.warm is None
|
||||
|
|
@ -433,6 +433,232 @@ def test_get_model_from_request_no_request_extracts_model():
|
|||
)
|
||||
|
||||
|
||||
def _cache_prediction_router():
|
||||
from litellm.router import Router
|
||||
|
||||
return Router(model_list=[
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"},
|
||||
"model_info": {"id": deployment_id, "team_id": team_id},
|
||||
}
|
||||
for group, deployment_id, team_id in (
|
||||
("current-group", "current-id", None), ("candidate-group", "candidate-id", None),
|
||||
("own-group", "own-id", "prediction-team"), ("foreign-group", "foreign-id", "foreign-team"),
|
||||
)
|
||||
])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("candidate,team_id,expected", [
|
||||
("candidate-id", None, ["current-group", "candidate-group"]),
|
||||
("current-id", None, "current-group"),
|
||||
("missing-id", None, None),
|
||||
("candidate-group", None, None),
|
||||
("own-id", None, None),
|
||||
("own-id", "prediction-team", ["current-group", "own-group"]),
|
||||
("foreign-id", "prediction-team", None),
|
||||
])
|
||||
def test_cache_prediction_auth_resolves_only_exact_deployment_ids(candidate, team_id, expected):
|
||||
assert get_model_from_request(
|
||||
request_data={
|
||||
"current_deployment_id": "current-id", "candidate_deployment_id": candidate,
|
||||
"request": {"model": "caller-controlled-provider-model"},
|
||||
},
|
||||
route="/cost/predict-cache",
|
||||
llm_router=_cache_prediction_router(),
|
||||
team_id=team_id,
|
||||
) == expected
|
||||
|
||||
|
||||
def _cache_prediction_auth_app(
|
||||
monkeypatch, allowed_routes, user_models, metadata=None, *, team_id=None, key_models=None, team_models=None
|
||||
):
|
||||
import importlib
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, ProxyException
|
||||
from litellm.proxy.auth import auth_checks
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint
|
||||
from litellm.proxy.utils import InternalUsageCache, ProxyLogging
|
||||
|
||||
auth = importlib.import_module("litellm.proxy.auth.user_api_key_auth")
|
||||
router = _cache_prediction_router()
|
||||
allowed_models = ["current-group", "candidate-group", "own-group"]
|
||||
token = UserAPIKeyAuth(
|
||||
api_key="test-proxy-key-hash", user_id="prediction-user", user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=allowed_models if key_models is None else key_models, team_id=team_id,
|
||||
team_models=allowed_models if team_models is None else team_models,
|
||||
allowed_routes=allowed_routes, metadata=metadata or {},
|
||||
)
|
||||
user = LiteLLM_UserTable(
|
||||
user_id=token.user_id, user_role=LitellmUserRoles.INTERNAL_USER.value, models=user_models,
|
||||
)
|
||||
async def authenticate(request, request_data, **_headers):
|
||||
await auth._enforce_key_and_fallback_model_access(
|
||||
valid_token=token, request_data=request_data, route=request.url.path, request=request,
|
||||
llm_model_list=router.get_model_list(), llm_router=router,
|
||||
)
|
||||
return token
|
||||
|
||||
monkeypatch.setattr(auth, "_user_api_key_auth_builder", authenticate)
|
||||
monkeypatch.setattr(auth, "get_user_object", AsyncMock(return_value=user))
|
||||
team = LiteLLM_TeamTableCachedObj(team_id=team_id, models=token.team_models) if team_id else None
|
||||
monkeypatch.setattr(auth, "get_team_object", AsyncMock(return_value=team))
|
||||
monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team))
|
||||
monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None))
|
||||
monkeypatch.setattr(auth, "get_global_proxy_spend", AsyncMock(return_value=0))
|
||||
monkeypatch.setattr(proxy_server, "master_key", "test-master-key")
|
||||
monkeypatch.setattr(proxy_server, "user_custom_auth", None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list())
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
|
||||
logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
logging.proxy_hook_mapping["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v3(
|
||||
InternalUsageCache(dual_cache=DualCache())
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging)
|
||||
counts = AsyncMock(return_value=6_000)
|
||||
monkeypatch.setattr(endpoint, "count_prompt_tokens", counts)
|
||||
app = FastAPI()
|
||||
app.include_router(endpoint.router)
|
||||
app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
|
||||
return app, counts
|
||||
|
||||
|
||||
def _cache_prediction_payload(candidate="candidate-id", current="current-id"):
|
||||
return {
|
||||
"current_deployment_id": current, "candidate_deployment_id": candidate,
|
||||
"request": {"messages": [{"role": "user", "content": [{
|
||||
"type": "text", "text": "Stable cached context",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}]}]},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("allowed_routes,user_models,candidate,status_code", [
|
||||
(["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403),
|
||||
(["/cost/predict-cache"], ["current-group"], "candidate-id", 403),
|
||||
(["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200),
|
||||
(["/cost/predict-cache"], ["current-group"], "current-id", 200),
|
||||
(["/cost/predict-cache"], ["current-group"], "missing-id", 404),
|
||||
])
|
||||
async def test_cache_prediction_authorizes_route_and_personal_models_before_provider_counts(
|
||||
monkeypatch, allowed_routes, user_models, candidate, status_code
|
||||
):
|
||||
import httpx
|
||||
|
||||
app, counts = _cache_prediction_auth_app(monkeypatch, allowed_routes, user_models)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(candidate))
|
||||
|
||||
assert response.status_code == status_code, response.text
|
||||
if status_code == 200:
|
||||
assert counts.await_count == (2 if candidate == "current-id" else 4)
|
||||
else:
|
||||
assert counts.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
|
||||
@pytest.mark.parametrize("team_id,key_models,user_models,team_models", [
|
||||
(None, ["*"], ["*"], None),
|
||||
(None, ["current-group", "candidate-group"], ["*"], None),
|
||||
(None, ["*"], ["current-group", "candidate-group"], None),
|
||||
("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]),
|
||||
])
|
||||
async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_authorization(
|
||||
monkeypatch, arm, team_id, key_models, user_models, team_models
|
||||
):
|
||||
import httpx
|
||||
|
||||
app, counts = _cache_prediction_auth_app(
|
||||
monkeypatch, ["/cost/predict-cache"], user_models,
|
||||
team_id=team_id, key_models=key_models, team_models=team_models,
|
||||
)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
missing = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "missing-id"})
|
||||
foreign = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "foreign-id"})
|
||||
|
||||
assert missing.status_code == foreign.status_code == 404, foreign.text
|
||||
assert missing.json() == foreign.json() == {"detail": "Deployment not found"}
|
||||
assert counts.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
|
||||
@pytest.mark.parametrize("key_models,team_models,status_code", [
|
||||
(["*"], ["*"], 200),
|
||||
(["current-group", "candidate-group"], ["*"], 403),
|
||||
(["*"], ["current-group", "candidate-group"], 403),
|
||||
])
|
||||
async def test_cache_prediction_checks_each_visible_team_deployment_model(
|
||||
monkeypatch, arm, key_models, team_models, status_code
|
||||
):
|
||||
import httpx
|
||||
|
||||
app, counts = _cache_prediction_auth_app(
|
||||
monkeypatch, ["/cost/predict-cache"], ["*"],
|
||||
team_id="prediction-team", key_models=key_models, team_models=team_models,
|
||||
)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "own-id"})
|
||||
|
||||
assert response.status_code == status_code, response.text
|
||||
assert counts.await_count == (4 if status_code == 200 else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"])
|
||||
async def test_cache_prediction_checks_each_visible_personal_deployment_model(monkeypatch, arm):
|
||||
import httpx
|
||||
|
||||
app, counts = _cache_prediction_auth_app(monkeypatch, ["/cost/predict-cache"], ["current-group"])
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
"/cost/predict-cache", json={**_cache_prediction_payload(candidate="current-id"), arm: "candidate-id"}
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json()["error"]["type"] == "user_model_access_denied"
|
||||
assert counts.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("header_tag,key_tags,limit,status_code,provider_calls", [
|
||||
("limited", [], 1, 429, 1),
|
||||
(None, ["limited"], 1, 429, 1),
|
||||
("limited", ["limited"], 4, 200, 4),
|
||||
("unlimited", [], 1, 200, 4),
|
||||
])
|
||||
async def test_cache_prediction_preserves_authenticated_header_and_key_tag_rpm(
|
||||
monkeypatch, header_tag, key_tags, limit, status_code, provider_calls
|
||||
):
|
||||
import httpx
|
||||
|
||||
app, counts = _cache_prediction_auth_app(
|
||||
monkeypatch, ["/cost/predict-cache"], ["current-group", "candidate-group"],
|
||||
metadata={"tag_rpm_limit": {"limited": limit}, "tags": key_tags},
|
||||
)
|
||||
headers = {"x-litellm-tags": header_tag} if header_tag else {}
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers)
|
||||
assert response.status_code == status_code, response.text
|
||||
assert counts.await_count == provider_calls
|
||||
if limit == 4:
|
||||
exhausted = await client.post("/cost/predict-cache", json=_cache_prediction_payload(), headers=headers)
|
||||
assert exhausted.status_code == 429, exhausted.text
|
||||
assert counts.await_count == 4
|
||||
assert all("metadata" not in call.args[2] for call in counts.await_args_list)
|
||||
|
||||
|
||||
def test_get_model_from_request_supports_google_model_names_with_slashes():
|
||||
assert (
|
||||
get_model_from_request(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,105 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected"),
|
||||
[("anthropic/claude-sonnet-4-5", 1.26), ("anthropic/claude-sonnet-4-6", 0.63)],
|
||||
)
|
||||
def test_prices_all_cache_buckets_at_total_context_tier(model: str, expected: float) -> None:
|
||||
tokens: Final = CacheTokenBuckets(
|
||||
uncached_input_tokens=100_000,
|
||||
cache_read_input_tokens=50_000,
|
||||
cache_creation_5m_input_tokens=20_000,
|
||||
cache_creation_1h_input_tokens=40_000,
|
||||
)
|
||||
assert price_cache_tokens(model, "unconfigured-deployment", tokens) == pytest.approx(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("total", "expected"), [(200_000, 0.387), (200_001, 0.774006)])
|
||||
def test_long_context_tier_starts_above_threshold(total: int, expected: float) -> None:
|
||||
tokens: Final = CacheTokenBuckets(
|
||||
uncached_input_tokens=total - 100_000,
|
||||
cache_creation_1h_input_tokens=10_000,
|
||||
cache_read_input_tokens=90_000,
|
||||
)
|
||||
actual: Final = price_cache_tokens("anthropic/claude-sonnet-4-5", "unconfigured-deployment", tokens)
|
||||
assert actual == pytest.approx(expected)
|
||||
|
||||
|
||||
def test_deployment_tariff_wins_without_proxy_discounts_or_margins(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.model_cost.copy())
|
||||
litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "cache-pricing-test",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-6",
|
||||
"api_key": "test-only",
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"cache_read_input_token_cost": 0.000001,
|
||||
"cache_creation_input_token_cost": 0.0000125,
|
||||
"cache_creation_input_token_cost_above_1hr": 0.00002,
|
||||
},
|
||||
"model_info": {"id": "cache-pricing-test-a"},
|
||||
}
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(litellm, "cost_discount_config", {"anthropic": 0.5})
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"percentage": 0.3, "fixed_amount": 1.0}})
|
||||
tokens: Final = CacheTokenBuckets(
|
||||
uncached_input_tokens=3_000,
|
||||
cache_read_input_tokens=4_000,
|
||||
cache_creation_5m_input_tokens=1_000,
|
||||
cache_creation_1h_input_tokens=2_000,
|
||||
)
|
||||
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-test-a", tokens) == pytest.approx(0.0865)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rate", [None, -1.0, float("nan"), float("inf"), "0.00001", True])
|
||||
def test_unknown_for_absent_or_invalid_active_cache_rate(monkeypatch: pytest.MonkeyPatch, rate: object) -> None:
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"cache-pricing-invalid",
|
||||
{
|
||||
"litellm_provider": "anthropic",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"cache_creation_input_token_cost_above_1hr": rate,
|
||||
},
|
||||
)
|
||||
tokens: Final = CacheTokenBuckets(cache_creation_1h_input_tokens=4_000)
|
||||
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-invalid", tokens) is None
|
||||
|
||||
|
||||
def test_missing_input_price_is_unknown_even_when_get_model_info_defaults_to_zero(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.model_cost, "cache-pricing-missing", {"litellm_provider": "anthropic", "mode": "chat"})
|
||||
tokens: Final = CacheTokenBuckets(uncached_input_tokens=4_000)
|
||||
assert price_cache_tokens("cache-pricing-missing", "unconfigured-deployment", tokens) is None
|
||||
|
||||
|
||||
def test_explicit_free_pricing_is_not_unknown(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"cache-pricing-free",
|
||||
{
|
||||
"litellm_provider": "anthropic",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"cache_read_input_token_cost": 0.0,
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_creation_input_token_cost_above_1hr": 0.0,
|
||||
},
|
||||
)
|
||||
tokens: Final = CacheTokenBuckets(uncached_input_tokens=100, cache_read_input_tokens=5_000)
|
||||
assert price_cache_tokens("anthropic/claude-sonnet-4-6", "cache-pricing-free", tokens) == 0.0
|
||||
|
|
@ -6284,3 +6284,248 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_
|
|||
assert isinstance(values, list)
|
||||
assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == []
|
||||
assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limits, request_data, counter_scope",
|
||||
[
|
||||
({"rpm_limit": 1}, {}, "api_key"),
|
||||
({"user_id": "u", "user_rpm_limit": 1}, {}, "user"),
|
||||
({"team_id": "t", "team_rpm_limit": 1}, {}, "team"),
|
||||
(
|
||||
{"team_id": "t", "user_id": "u", "team_member_rpm_limit": 1},
|
||||
{},
|
||||
"team_member",
|
||||
),
|
||||
({"end_user_id": "e", "end_user_rpm_limit": 1}, {}, "end_user"),
|
||||
(
|
||||
{"metadata": {"model_rpm_limit": {"test-model": 1}}},
|
||||
{},
|
||||
"model_per_key",
|
||||
),
|
||||
(
|
||||
{"metadata": {"tag_rpm_limit": {"test-tag": 1}}},
|
||||
{"metadata": {"tags": ["test-tag"]}},
|
||||
"tag_per_key",
|
||||
),
|
||||
(
|
||||
{
|
||||
"team_id": "t",
|
||||
"metadata": {"model_rpm_limit": {"test-model": 100}},
|
||||
"team_metadata": {"model_rpm_limit": {"test-model": 1}},
|
||||
},
|
||||
{},
|
||||
"model_per_team",
|
||||
),
|
||||
(
|
||||
{"project_id": "p", "project_metadata": {"model_rpm_limit": {"test-model": 1}}},
|
||||
{},
|
||||
"model_per_project",
|
||||
),
|
||||
({"org_id": "o", "organization_rpm_limit": 1}, {}, "organization"),
|
||||
(
|
||||
{"org_id": "o", "organization_metadata": {"model_rpm_limit": {"test-model": 1}}},
|
||||
{},
|
||||
"model_per_organization",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("request_kind", ["count", "generation"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_enforces_shared_rpm_scopes(
|
||||
limits, request_data, counter_scope, request_kind
|
||||
):
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-rpm"), **limits)
|
||||
async def request():
|
||||
if request_kind == "generation":
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=cache,
|
||||
data={**request_data, "model": "test-model"},
|
||||
call_type="acompletion",
|
||||
)
|
||||
return
|
||||
async with handler.request_capacity(auth, "test-model", request_data=request_data):
|
||||
pass
|
||||
|
||||
await request()
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await request()
|
||||
assert exc.value.status_code == 429
|
||||
assert counter_scope in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
router = Router(model_list=[{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "openai/gpt-test", "api_key": "test-key"},
|
||||
"model_info": {"id": "test-deployment"},
|
||||
}])
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-count-dynamic"),
|
||||
rpm_limit=1,
|
||||
metadata={"rpm_limit_type": "dynamic"},
|
||||
)
|
||||
for _ in range(2):
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pass
|
||||
router.cache.set_cache("test-deployment:fails", 100, ttl=60, local_only=True)
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pass
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pytest.fail("dynamic RPM must enforce after deployment failures")
|
||||
assert exc.value.status_code == 429
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_skips_tokens_and_preserves_parent_stash():
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-count-tpm"),
|
||||
rpm_limit=5,
|
||||
tpm_limit=1,
|
||||
max_parallel_requests=1,
|
||||
project_id="p",
|
||||
project_metadata={
|
||||
"model_tpm_limit": {"test-model": 1},
|
||||
"model_itpm_limit": {"test-model": 1},
|
||||
"model_otpm_limit": {"test-model": 1},
|
||||
},
|
||||
)
|
||||
token_scopes = (
|
||||
("api_key", auth.api_key),
|
||||
("model_per_project", "p:test-model"),
|
||||
("model_per_project_itpm", "p:test-model"),
|
||||
("model_per_project_otpm", "p:test-model"),
|
||||
)
|
||||
for scope, value in token_scopes:
|
||||
token_key = handler.create_rate_limit_keys(scope, value, "tokens")
|
||||
await cache.async_set_cache(token_key, 100, ttl=60)
|
||||
await cache.async_set_cache(f"{{{scope}:{value}}}:window", int(time.time()), ttl=60)
|
||||
parent = get_or_create_request_stash()
|
||||
parent.reserved_tokens = 123
|
||||
parent.parallel_slot = ParallelSlotAcquisition(slot_id="parent", counter_keys=["parent-gauge"])
|
||||
for _ in range(2):
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
assert get_request_stash() is parent
|
||||
assert parent.parallel_slot["slot_id"] == "parent"
|
||||
assert parent.reserved_tokens == 123
|
||||
for scope, value in token_scopes:
|
||||
assert await cache.async_get_cache(handler.create_rate_limit_keys(scope, value, "tokens")) == 100
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exit_mode", ["success", "failure", "cancel"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_releases_exact_parallel_slot(exit_mode):
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-parallel"), max_parallel_requests=1)
|
||||
entered = asyncio.Event()
|
||||
finish = asyncio.Event()
|
||||
|
||||
async def provider():
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
entered.set()
|
||||
await finish.wait()
|
||||
if exit_mode == "failure":
|
||||
raise RuntimeError("provider failed")
|
||||
|
||||
task = asyncio.create_task(provider())
|
||||
await asyncio.wait_for(entered.wait(), timeout=2)
|
||||
try:
|
||||
for _ in range(2):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pytest.fail("rejected request freed the occupied slot")
|
||||
assert exc.value.status_code == 429
|
||||
finally:
|
||||
if exit_mode == "cancel":
|
||||
task.cancel()
|
||||
else:
|
||||
finish.set()
|
||||
if exit_mode == "success":
|
||||
await task
|
||||
else:
|
||||
with pytest.raises(asyncio.CancelledError if exit_mode == "cancel" else RuntimeError):
|
||||
await task
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pass
|
||||
|
||||
|
||||
class _DelayedCapacityUsageCache:
|
||||
def __init__(self):
|
||||
self.delegate = InternalUsageCache(DualCache())
|
||||
self.dual_cache = self.delegate.dual_cache
|
||||
self.acquired = asyncio.Event()
|
||||
self.finish_admission = asyncio.Event()
|
||||
self.releasing = asyncio.Event()
|
||||
self.finish_release = asyncio.Event()
|
||||
|
||||
async def async_get_cache(self, *args, **kwargs):
|
||||
return await self.delegate.async_get_cache(*args, **kwargs)
|
||||
|
||||
async def async_batch_get_cache(self, *args, **kwargs):
|
||||
return await self.delegate.async_batch_get_cache(*args, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key, value, **kwargs):
|
||||
await self.delegate.async_set_cache(key=key, value=value, **kwargs)
|
||||
if not key.endswith(":max_parallel_requests"):
|
||||
return
|
||||
if value:
|
||||
self.acquired.set()
|
||||
await self.finish_admission.wait()
|
||||
else:
|
||||
self.releasing.set()
|
||||
await self.finish_release.wait()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_finishes_admission_and_release_despite_repeated_cancel():
|
||||
cache = _DelayedCapacityUsageCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-cancel-admission"), max_parallel_requests=1)
|
||||
|
||||
async def provider():
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pytest.fail("cancelled admission entered provider body")
|
||||
|
||||
task = asyncio.create_task(provider())
|
||||
await asyncio.wait_for(cache.acquired.wait(), timeout=2)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
cache.finish_admission.set()
|
||||
await asyncio.wait_for(cache.releasing.wait(), timeout=2)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
cache.finish_release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=2)
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_capacity_rejection_keeps_existing_redis_mirror():
|
||||
cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-count-mirror"), max_parallel_requests=1)
|
||||
counter_key = handler.create_rate_limit_keys("api_key", auth.api_key, "max_parallel_requests")
|
||||
await cache.async_set_cache(counter_key, 1, ttl=60, local_only=True)
|
||||
for _ in range(2):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
async with handler.request_capacity(auth, "test-model"):
|
||||
pytest.fail("rejection released another request's mirrored slot")
|
||||
assert exc.value.status_code == 429
|
||||
assert await cache.async_get_cache(counter_key, local_only=True) == 1
|
||||
|
|
|
|||
300
tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py
Normal file
300
tests/test_litellm/proxy/hooks/test_prompt_cache_observer.py
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import cache_scope, parse_prompt
|
||||
from litellm.proxy.hooks.prompt_cache_prediction import (
|
||||
PromptCacheObserver,
|
||||
lookup,
|
||||
)
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
MODEL = "claude-sonnet-5"
|
||||
CALLER = "a" * 64
|
||||
DEPLOYMENT = "native-deployment"
|
||||
KEY = "test-provider-key"
|
||||
|
||||
|
||||
def body(ttl="5m", texts=("private cache prefix",)):
|
||||
return {
|
||||
"model": MODEL,
|
||||
"max_tokens": 2,
|
||||
"system": "private system instructions",
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}],
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": text, **(
|
||||
{"cache_control": {"type": "ephemeral", "ttl": ttl}}
|
||||
if index == len(texts) - 1 else {}
|
||||
)}
|
||||
for index, text in enumerate(texts)
|
||||
]}],
|
||||
}
|
||||
|
||||
|
||||
def usage(ttl="5m", read=100, write=200):
|
||||
return {
|
||||
"input_tokens": 11,
|
||||
"output_tokens": 2,
|
||||
"cache_read_input_tokens": read,
|
||||
"cache_creation_input_tokens": write,
|
||||
"cache_creation": {
|
||||
"ephemeral_5m_input_tokens": write if ttl == "5m" else 0,
|
||||
"ephemeral_1h_input_tokens": write if ttl == "1h" else 0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def event(request_body, started=1000.0, headers=None, **overrides):
|
||||
request = httpx.Request(
|
||||
"POST", "https://api.anthropic.com/v1/messages", json=request_body,
|
||||
headers={"x-api-key": KEY, "anthropic-version": "2023-06-01", **(headers or {})},
|
||||
)
|
||||
return {
|
||||
"call_type": "anthropic_messages",
|
||||
"custom_llm_provider": "anthropic",
|
||||
"cache_hit": False,
|
||||
"httpx_response": httpx.Response(200, request=request),
|
||||
"first_api_call_start_time": datetime.fromtimestamp(started),
|
||||
"standard_logging_object": {
|
||||
"status": "success", "model_id": DEPLOYMENT,
|
||||
"metadata": {"user_api_key_hash": CALLER},
|
||||
},
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
async def observe(cache, request_body=None, native_usage=None, now=1010.0, **overrides):
|
||||
observer = PromptCacheObserver(InternalUsageCache(dual_cache=cache), clock=lambda: now)
|
||||
response = ModelResponse(
|
||||
model=MODEL,
|
||||
usage=AnthropicConfig().calculate_usage(native_usage or usage(), reasoning_content=None),
|
||||
)
|
||||
await observer.async_log_success_event(
|
||||
event(request_body or body(), **overrides), response,
|
||||
datetime.fromtimestamp(now), datetime.fromtimestamp(now),
|
||||
)
|
||||
|
||||
|
||||
def scope(**overrides):
|
||||
return cache_scope(**{
|
||||
"caller_key_hash": CALLER, "deployment_id": DEPLOYMENT,
|
||||
"provider_key": KEY, "model": MODEL, **overrides,
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ttl,expires", [("5m", 1300), ("1h", 4600)])
|
||||
@pytest.mark.asyncio
|
||||
async def test_observed_cache_count_and_request_start_expiry_survive_as_stale(ttl, expires):
|
||||
cache = DualCache()
|
||||
request_body = body(ttl=ttl)
|
||||
await observe(cache, request_body, usage(ttl=ttl))
|
||||
prefix = parse_prompt(request_body)
|
||||
observed = await lookup(cache, scope(), prefix, now=1200)
|
||||
assert observed.cached_tokens == 300
|
||||
assert observed.observed_at == 1010
|
||||
assert observed.expires_at == expires
|
||||
assert await lookup(cache, scope(), prefix, now=expires) == observed
|
||||
saved = json.dumps(cache.in_memory_cache.cache_dict)
|
||||
assert "private cache prefix" not in saved
|
||||
assert "private system instructions" not in saved
|
||||
assert KEY not in saved
|
||||
assert CALLER not in saved
|
||||
|
||||
|
||||
@pytest.mark.parametrize("changed", [
|
||||
{"caller_key_hash": "b" * 64}, {"deployment_id": "other"},
|
||||
{"provider_key": "rotated"}, {"model": "claude-opus-5"},
|
||||
{"anthropic_version": "different"},
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_evidence_is_isolated_by_every_scope_dimension(changed):
|
||||
cache = DualCache()
|
||||
await observe(cache)
|
||||
assert await lookup(cache, scope(**changed), parse_prompt(body()), now=1010) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_only_prefix_finds_prior_evidence_but_edit_or_context_change_does_not():
|
||||
cache = DualCache()
|
||||
await observe(cache)
|
||||
extended = parse_prompt(body(texts=("private cache prefix", "new turn")))
|
||||
prior = await lookup(cache, scope(), extended, now=1010)
|
||||
assert prior.cached_tokens == 300
|
||||
assert prior.fingerprint != extended.fingerprint
|
||||
for changed in (
|
||||
body(texts=("edited prefix", "new turn")),
|
||||
{**body(), "system": "different system"},
|
||||
{**body(), "tools": [{"name": "other", "input_schema": {"type": "object"}}]},
|
||||
body(ttl="1h"),
|
||||
):
|
||||
assert await lookup(cache, scope(), parse_prompt(changed), now=1010) is None
|
||||
outside_lookback = parse_prompt(body(texts=("private cache prefix", *[str(i) for i in range(20)])))
|
||||
assert await lookup(cache, scope(), outside_lookback, now=1010) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", [
|
||||
{"thinking": {"type": "enabled", "budget_tokens": 1024}},
|
||||
{"tool_choice": {"type": "auto"}},
|
||||
{"cache_control": {"type": "ephemeral"}},
|
||||
{"tools": [{"type": "web_search_20250305", "name": "web_search"}]},
|
||||
{"system": [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral"}}]},
|
||||
{"messages": [{"role": "user", "content": [{"type": "image", "source": {}}]}]},
|
||||
{"messages": [{"role": "user", "content": "no breakpoint"}]},
|
||||
])
|
||||
def test_unsupported_or_ambiguous_shapes_have_no_cache_identity(change):
|
||||
assert parse_prompt({**body(), **change}) is None
|
||||
duplicate = body()
|
||||
duplicate["messages"][0]["content"].append(duplicate["messages"][0]["content"][0])
|
||||
assert parse_prompt(duplicate) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("overrides", [
|
||||
{"cache_hit": True}, {"call_type": "completion"},
|
||||
{"custom_llm_provider": "bedrock"}, {"stream": True},
|
||||
{"headers": {"anthropic-beta": "unverified-feature"}},
|
||||
{"headers": {"x-custom-header": "unverified"}},
|
||||
{"standard_logging_object": {"status": "success", "model_id": DEPLOYMENT, "metadata": {}}},
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_unverified_source_never_creates_observations(overrides):
|
||||
cache = DualCache()
|
||||
await observe(cache, **overrides)
|
||||
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("native_usage", [
|
||||
usage(write=0),
|
||||
{**usage(), "cache_creation": None},
|
||||
{**usage(), "cache_creation": {"ephemeral_5m_input_tokens": 199, "ephemeral_1h_input_tokens": 0}},
|
||||
usage(ttl="1h"),
|
||||
{**usage(), "cache_creation_input_tokens": -200},
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_or_contradictory_telemetry_cannot_create_observations(native_usage):
|
||||
cache = DualCache()
|
||||
await observe(cache, native_usage=native_usage)
|
||||
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pure_read_refresh_requires_prior_matching_evidence():
|
||||
cache = DualCache()
|
||||
await observe(cache, native_usage=usage(read=300, write=0))
|
||||
assert await lookup(cache, scope(), parse_prompt(body()), now=1010) is None
|
||||
await observe(cache)
|
||||
await observe(cache, native_usage=usage(read=300, write=0), started=1100, now=1110)
|
||||
assert (await lookup(cache, scope(), parse_prompt(body()), now=1110)).expires_at == 1400
|
||||
|
||||
|
||||
class RecordingObserver(PromptCacheObserver):
|
||||
def __init__(self, cache):
|
||||
super().__init__(InternalUsageCache(dual_cache=cache))
|
||||
self.finished = asyncio.Event()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
self.finished.set()
|
||||
|
||||
|
||||
def native_response():
|
||||
return {
|
||||
"id": "msg_prediction", "type": "message", "role": "assistant", "model": MODEL,
|
||||
"content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn",
|
||||
"stop_sequence": None, "usage": usage(ttl="1h"),
|
||||
}
|
||||
|
||||
|
||||
def stream_response(completed, provider_error=False):
|
||||
response = native_response()
|
||||
events = [
|
||||
{"type": "message_start", "message": {**response, "content": [], "stop_reason": None}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
|
||||
]
|
||||
if completed:
|
||||
events.append({"type": "message_stop"})
|
||||
if provider_error:
|
||||
events.append({"type": "error", "error": {"type": "overloaded_error", "message": "temporary failure"}})
|
||||
return "".join(f"event: {item['type']}\ndata: {json.dumps(item)}\n\n" for item in events)
|
||||
|
||||
|
||||
class TransportChunks(httpx.AsyncByteStream):
|
||||
def __init__(self, payload, chunk_size, fragment_error_only=False):
|
||||
self.payload = payload.encode()
|
||||
self.chunk_size = chunk_size or len(self.payload)
|
||||
self.prefix_length = self.payload.index(b"event: error") if fragment_error_only else 0
|
||||
|
||||
async def __aiter__(self):
|
||||
if self.prefix_length:
|
||||
yield self.payload[:self.prefix_length]
|
||||
for offset in range(self.prefix_length, len(self.payload), self.chunk_size):
|
||||
yield self.payload[offset:offset + self.chunk_size]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream,completed,provider_error,transport", [
|
||||
(False, True, False, "whole"),
|
||||
(True, True, False, "whole"),
|
||||
(True, False, False, "whole"),
|
||||
(True, True, True, "whole"),
|
||||
(True, True, False, "fragmented"),
|
||||
(True, False, False, "fragmented"),
|
||||
(True, True, True, "fragmented"),
|
||||
(True, True, True, "fragmented_error"),
|
||||
(True, True, False, "unterminated"),
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_production_callback_records_only_completed_wire_requests(stream, completed, provider_error, transport):
|
||||
cache = DualCache()
|
||||
observer = RecordingObserver(cache)
|
||||
litellm.logging_callback_manager.add_litellm_callback(observer)
|
||||
|
||||
def provider(request):
|
||||
if stream:
|
||||
payload = stream_response(completed, provider_error)
|
||||
if transport == "unterminated":
|
||||
payload = payload.removesuffix("\n\n")
|
||||
return httpx.Response(
|
||||
200, request=request, headers={"content-type": "text/event-stream"},
|
||||
stream=TransportChunks(
|
||||
payload, 1 if transport.startswith("fragmented") else None,
|
||||
fragment_error_only=transport == "fragmented_error",
|
||||
),
|
||||
)
|
||||
return httpx.Response(200, request=request, json=native_response())
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
await client.client.aclose()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
|
||||
try:
|
||||
request_body = body(ttl="1h")
|
||||
before = time.time()
|
||||
result = await litellm.anthropic_messages(
|
||||
**{**request_body, "model": f"anthropic/{MODEL}"},
|
||||
api_key=KEY, client=client, stream=stream, model_info={"id": DEPLOYMENT},
|
||||
litellm_metadata={"user_api_key_hash": CALLER, "model_info": {"id": DEPLOYMENT}},
|
||||
)
|
||||
if stream:
|
||||
async for _ in result:
|
||||
pass
|
||||
await asyncio.wait_for(observer.finished.wait(), timeout=5)
|
||||
found = await lookup(cache, scope(), parse_prompt(request_body))
|
||||
if completed and not provider_error and transport != "unterminated":
|
||||
assert found is not None
|
||||
assert found.cached_tokens == 300
|
||||
assert before + 3600 <= found.expires_at <= time.time() + 3600
|
||||
else:
|
||||
assert found is None
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(observer)
|
||||
await client.client.aclose()
|
||||
|
|
@ -0,0 +1,698 @@
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body, _safe_set_request_parsed_body
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, cache_scope, parse_prompt
|
||||
from litellm.proxy.hooks.prompt_cache_prediction import (
|
||||
CacheObservation,
|
||||
_cache_key,
|
||||
)
|
||||
from litellm.proxy.management_endpoints import prompt_cache_prediction as endpoint
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.types.management_endpoints.prompt_cache_prediction import CachePredictionResponse
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
|
||||
_PROVIDER_KEY: Final = "cache-prediction-test-provider-key"
|
||||
_CALLER: Final = "cache-prediction-test-caller-hash"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def anthropic_endpoint_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
|
||||
|
||||
|
||||
def _body(ttl: str = "5m", *, extended: bool = False) -> dict[str, JsonValue]:
|
||||
blocks: Final[list[JsonValue]] = [
|
||||
{"type": "text", "text": "Stable context"},
|
||||
*([{"type": "text", "text": "Appended context"}] if extended else []),
|
||||
]
|
||||
return {
|
||||
"max_tokens": 10,
|
||||
"system": "Follow the project conventions",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*blocks[:-1],
|
||||
{**blocks[-1], "cache_control": {"type": "ephemeral", "ttl": ttl}},
|
||||
{"type": "text", "text": "Follow-up question"},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _prefix(body: Mapping[str, JsonValue]) -> PromptPrefix:
|
||||
prefix: Final = parse_prompt(body)
|
||||
assert prefix is not None
|
||||
return prefix
|
||||
|
||||
|
||||
def _deployment(
|
||||
deployment_id: str = "sonnet",
|
||||
model: str = "claude-sonnet-5",
|
||||
*,
|
||||
team_id: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> Deployment:
|
||||
return Deployment(
|
||||
model_name=deployment_id,
|
||||
litellm_params=LiteLLM_Params(model=f"anthropic/{model}", api_key=_PROVIDER_KEY, api_base=api_base),
|
||||
model_info=ModelInfo(id=deployment_id, team_id=team_id),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Counts:
|
||||
total: int | None = 6_000
|
||||
prefix: int | None = 5_000
|
||||
|
||||
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
assert api_key == _PROVIDER_KEY
|
||||
assert model.startswith("claude-")
|
||||
return self.total if "max_tokens" in body else self.prefix
|
||||
|
||||
|
||||
async def _observe(
|
||||
cache: DualCache,
|
||||
body: Mapping[str, JsonValue],
|
||||
*,
|
||||
deployment_id: str = "sonnet",
|
||||
model: str = "claude-sonnet-5",
|
||||
cached_tokens: int = 5_000,
|
||||
expired: bool = False,
|
||||
caller: str = _CALLER,
|
||||
) -> None:
|
||||
prefix: Final = _prefix(body)
|
||||
now: Final = time.time()
|
||||
observation: Final = CacheObservation(
|
||||
fingerprint=prefix.fingerprint,
|
||||
cached_tokens=cached_tokens,
|
||||
observed_at=now - 400 if expired else now - 10,
|
||||
expires_at=now - 100 if expired else now + 290,
|
||||
)
|
||||
scope: Final = cache_scope(caller, deployment_id, _PROVIDER_KEY, model)
|
||||
await cache.async_set_cache(_cache_key(scope, prefix.fingerprint), observation.model_dump_json(), ttl=3_600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("ttl", "cold_cost"), [("5m", 0.0145), ("1h", 0.022)])
|
||||
async def test_unobserved_cache_prices_cold_and_warm_bounds(ttl: str, cold_cost: float) -> None:
|
||||
body: Final = _body(ttl)
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts())
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "no_compatible_observation"
|
||||
assert arm.evidence is None
|
||||
assert arm.estimate is not None and arm.cold is not None and arm.warm is not None
|
||||
assert arm.estimate.input_cost == pytest.approx(cold_cost)
|
||||
assert arm.cold.input_cost == pytest.approx(cold_cost)
|
||||
assert arm.warm.input_cost == pytest.approx(0.003)
|
||||
assert arm.cold.tokens.uncached_input_tokens == 1_000
|
||||
assert arm.cold.tokens.cache_read_input_tokens == 0
|
||||
assert arm.cold.tokens.cache_creation_5m_input_tokens == (5_000 if ttl == "5m" else 0)
|
||||
assert arm.cold.tokens.cache_creation_1h_input_tokens == (5_000 if ttl == "1h" else 0)
|
||||
assert arm.warm.tokens.cache_read_input_tokens == 5_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("cached_tokens", "warm_cost", "cold_cost"), [(5_400, 0.00228, 0.0147), (4_600, 0.00372, 0.0143)]
|
||||
)
|
||||
@pytest.mark.parametrize("expired", [False, True])
|
||||
async def test_exact_prefix_conserves_total_with_observed_count_in_all_scenarios(
|
||||
cached_tokens: int, warm_cost: float, cold_cost: float, expired: bool
|
||||
) -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body, cached_tokens=cached_tokens, expired=expired)
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
|
||||
|
||||
assert arm.cache_state == ("stale" if expired else "warm")
|
||||
assert arm.evidence is not None
|
||||
assert arm.estimate is not None and arm.warm is not None and arm.cold is not None
|
||||
assert arm.warm.tokens.cache_read_input_tokens == cached_tokens
|
||||
assert arm.warm.tokens.cache_creation_5m_input_tokens == 0
|
||||
assert arm.cold.tokens.cache_creation_5m_input_tokens == cached_tokens
|
||||
assert arm.cold.tokens.cache_read_input_tokens == 0
|
||||
for scenario in (arm.estimate, arm.cold, arm.warm):
|
||||
assert scenario.tokens.total_tokens == 6_000
|
||||
assert scenario.tokens.uncached_input_tokens == 6_000 - cached_tokens
|
||||
assert arm.warm.input_cost == pytest.approx(warm_cost)
|
||||
assert arm.cold.input_cost == pytest.approx(cold_cost)
|
||||
assert arm.estimate.input_cost == pytest.approx(cold_cost if expired else warm_cost)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observed_prefix_larger_than_full_request_returns_unknown() -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body, cached_tokens=6_001)
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "inconsistent_prefix_token_count"
|
||||
assert arm.estimate is None and arm.cold is None and arm.warm is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("ttl", "expected"), [("5m", 0.0053), ("1h", 0.0068)])
|
||||
async def test_append_only_prefix_reads_old_tokens_and_writes_extension(ttl: str, expected: float) -> None:
|
||||
cache: Final = DualCache()
|
||||
await _observe(cache, _body(ttl), cached_tokens=4_000)
|
||||
body: Final = _body(ttl, extended=True)
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
|
||||
|
||||
assert arm.cache_state == "partial"
|
||||
assert arm.estimate is not None
|
||||
assert arm.estimate.tokens.cache_read_input_tokens == 4_000
|
||||
assert arm.estimate.tokens.cache_creation_5m_input_tokens == (1_000 if ttl == "5m" else 0)
|
||||
assert arm.estimate.tokens.cache_creation_1h_input_tokens == (1_000 if ttl == "1h" else 0)
|
||||
assert arm.estimate.input_cost == pytest.approx(expected)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_observation_estimates_a_cold_rebuild() -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body, expired=True)
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
|
||||
|
||||
assert arm.cache_state == "stale"
|
||||
assert arm.reason == "observation_expired"
|
||||
assert arm.evidence is not None and arm.evidence.expires_at < time.time()
|
||||
assert arm.estimate is not None and arm.cold is not None
|
||||
assert arm.estimate.tokens.cache_read_input_tokens == 0
|
||||
assert arm.estimate.tokens.cache_creation_5m_input_tokens == 5_000
|
||||
assert arm.estimate.input_cost == arm.cold.input_cost
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_below_model_minimum_prices_all_input_as_uncached() -> None:
|
||||
body: Final = _body()
|
||||
arm: Final = await endpoint.predict_arm(
|
||||
_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts(total=1_500, prefix=1_000)
|
||||
)
|
||||
|
||||
assert arm.cache_state == "disabled"
|
||||
assert arm.reason == "below_cache_minimum"
|
||||
assert arm.estimate is not None
|
||||
assert arm.estimate.tokens.uncached_input_tokens == 1_500
|
||||
assert arm.estimate.tokens.cache_read_input_tokens == 0
|
||||
assert arm.estimate.tokens.cache_creation_5m_input_tokens == 0
|
||||
assert arm.estimate.input_cost == pytest.approx(0.003)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("counts", [Counts(total=None), Counts(prefix=None), Counts(total=4_000)])
|
||||
async def test_unavailable_or_inconsistent_token_counts_return_null_estimates(counts: Counts) -> None:
|
||||
body: Final = _body()
|
||||
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), counts)
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "token_count_unavailable"
|
||||
assert arm.estimate is None and arm.cold is None and arm.warm is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("counts", [Counts(), Counts(total=1_500, prefix=1_000)])
|
||||
async def test_missing_prices_return_unknown_and_null_estimates(
|
||||
monkeypatch: pytest.MonkeyPatch, counts: Counts
|
||||
) -> None:
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"claude-cache-unpriced-5",
|
||||
{"litellm_provider": "anthropic", "mode": "chat"},
|
||||
)
|
||||
body: Final = _body()
|
||||
arm: Final = await endpoint.predict_arm(
|
||||
_deployment("cache-prediction-unpriced", "claude-cache-unpriced-5"),
|
||||
body,
|
||||
_prefix(body),
|
||||
_CALLER,
|
||||
DualCache(),
|
||||
counts,
|
||||
)
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "pricing_unavailable"
|
||||
assert arm.estimate is None and arm.cold is None and arm.warm is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_api_base_from_environment_returns_unknown_before_counting(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid")
|
||||
body: Final = _body()
|
||||
arm: Final = await endpoint.predict_arm(
|
||||
_deployment(), body, _prefix(body), _CALLER, DualCache(), _unexpected_count
|
||||
)
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "unsupported_provider_endpoint"
|
||||
assert arm.estimate is None and arm.cold is None and arm.warm is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_official_api_base_overrides_custom_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid")
|
||||
body: Final = _body()
|
||||
arm: Final = await endpoint.predict_arm(
|
||||
_deployment(api_base="https://api.anthropic.com"), body, _prefix(body), _CALLER, DualCache(), Counts()
|
||||
)
|
||||
|
||||
assert arm.cache_state == "unknown"
|
||||
assert arm.reason == "no_compatible_observation"
|
||||
assert arm.estimate is not None
|
||||
assert arm.estimate.input_cost == pytest.approx(0.0145)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ProxyLogging:
|
||||
internal_usage_cache: InternalUsageCache
|
||||
parallel_limiter: CustomLogger | None
|
||||
|
||||
def get_proxy_hook(self, hook: str) -> CustomLogger | None:
|
||||
return self.parallel_limiter if hook == "parallel_request_limiter" else None
|
||||
|
||||
|
||||
def _app(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
cache: DualCache,
|
||||
*,
|
||||
caller: UserAPIKeyAuth | None = None,
|
||||
current_team: str | None = None,
|
||||
candidate_team: str | None = None,
|
||||
counts: endpoint.TokenCounter = Counts(),
|
||||
limiter: CustomLogger | Literal["default"] | None = "default",
|
||||
) -> FastAPI:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
model_list: Final = [
|
||||
_deployment("opus", "claude-opus-5", team_id=current_team).model_dump(exclude_unset=True),
|
||||
_deployment("sonnet", team_id=candidate_team).model_dump(exclude_unset=True),
|
||||
]
|
||||
router: Final = litellm.Router(model_list=model_list)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", model_list)
|
||||
monkeypatch.setattr(endpoint, "count_prompt_tokens", counts)
|
||||
app: Final = FastAPI()
|
||||
app.include_router(endpoint.router)
|
||||
app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
|
||||
if caller is not None:
|
||||
usage_cache: Final = InternalUsageCache(cache)
|
||||
configured_limiter: Final = (
|
||||
_PROXY_MaxParallelRequestsHandler_v3(usage_cache) if isinstance(limiter, str) else limiter
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", _ProxyLogging(usage_cache, configured_limiter))
|
||||
app.dependency_overrides[endpoint.user_api_key_auth] = lambda: caller
|
||||
return app
|
||||
|
||||
|
||||
async def _post(
|
||||
app: FastAPI,
|
||||
body: Mapping[str, JsonValue],
|
||||
*,
|
||||
current_deployment_id: str = "opus",
|
||||
candidate_deployment_id: str = "sonnet",
|
||||
) -> httpx.Response:
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
return await client.post(
|
||||
"/cost/predict-cache",
|
||||
json={
|
||||
"current_deployment_id": current_deployment_id,
|
||||
"candidate_deployment_id": candidate_deployment_id,
|
||||
"request": body,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("warm_deployment", "warm_model", "expected_delta", "expected_penalty"),
|
||||
[("sonnet", "claude-sonnet-5", -0.03325, 0.0), ("opus", "claude-opus-5", 0.007, 0.0115)],
|
||||
)
|
||||
async def test_switch_delta_accounts_for_each_deployment_cache(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
warm_deployment: str,
|
||||
warm_model: str,
|
||||
expected_delta: float,
|
||||
expected_penalty: float,
|
||||
) -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body, deployment_id=warm_deployment, model=warm_model)
|
||||
app: Final = _app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER))
|
||||
response: Final = await _post(app, body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.switch_delta == pytest.approx(expected_delta)
|
||||
assert result.cache_rebuild_penalty == pytest.approx(expected_penalty)
|
||||
assert result.cache_guarantee is False
|
||||
assert result.pricing_basis == "input_before_discounts_and_margins"
|
||||
if warm_deployment == "sonnet":
|
||||
assert result.switch.cache_state == "warm"
|
||||
assert result.stay.cache_state == "unknown"
|
||||
else:
|
||||
assert result.stay.cache_state == "warm"
|
||||
assert result.switch.cache_state == "unknown"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_caller_identity_cannot_reuse_observations(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body)
|
||||
response: Final = await _post(
|
||||
_app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=None), counts=_unexpected_count), body
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.reason == result.switch.reason == "caller_identity_unavailable"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unauthenticated_request_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "master_key", "cache-prediction-test-master-key")
|
||||
response: Final = await _post(_app(monkeypatch, DualCache()), _body())
|
||||
assert response.status_code == 401, response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("arm", ["current", "candidate"])
|
||||
@pytest.mark.parametrize("caller_team", [None, "own-team"])
|
||||
@pytest.mark.parametrize("restricted", [False, True])
|
||||
async def test_foreign_and_missing_deployments_have_identical_authenticated_responses(
|
||||
monkeypatch: pytest.MonkeyPatch, arm: str, caller_team: str | None, restricted: bool
|
||||
) -> None:
|
||||
allowed: Final = ("sonnet",) if arm == "current" else ("opus",)
|
||||
app: Final = _app(
|
||||
monkeypatch,
|
||||
DualCache(),
|
||||
caller=UserAPIKeyAuth(api_key=_CALLER, team_id=caller_team, models=list(allowed) if restricted else []),
|
||||
current_team="foreign-team" if arm == "current" else None,
|
||||
candidate_team="foreign-team" if arm == "candidate" else None,
|
||||
counts=_unexpected_count,
|
||||
)
|
||||
foreign: Final = await _post(app, _body())
|
||||
missing: Final = await _post(
|
||||
app,
|
||||
_body(),
|
||||
current_deployment_id="missing-deployment" if arm == "current" else "opus",
|
||||
candidate_deployment_id="missing-deployment" if arm == "candidate" else "sonnet",
|
||||
)
|
||||
|
||||
assert foreign.status_code == missing.status_code == 404
|
||||
assert foreign.json() == missing.json() == {"detail": "Deployment not found"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deployment_team", [None, "own-team"])
|
||||
async def test_visible_public_and_own_team_deployments_remain_available(
|
||||
monkeypatch: pytest.MonkeyPatch, deployment_team: str | None
|
||||
) -> None:
|
||||
app: Final = _app(
|
||||
monkeypatch,
|
||||
DualCache(),
|
||||
caller=UserAPIKeyAuth(api_key=_CALLER, team_id="own-team"),
|
||||
current_team=deployment_team,
|
||||
candidate_team=deployment_team,
|
||||
)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.estimate is not None and result.switch.estimate is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("arm", ["current", "candidate"])
|
||||
async def test_visible_deployment_outside_key_model_permissions_is_forbidden(
|
||||
monkeypatch: pytest.MonkeyPatch, arm: str
|
||||
) -> None:
|
||||
allowed: Final = "sonnet" if arm == "current" else "opus"
|
||||
denied: Final = "opus" if arm == "current" else "sonnet"
|
||||
app: Final = _app(monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER, models=[allowed]))
|
||||
response: Final = await _post(app, _body())
|
||||
assert response.status_code == 403, response.text
|
||||
assert denied in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_other_callers_warm_cache_is_not_prediction_evidence(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache: Final = DualCache()
|
||||
body: Final = _body()
|
||||
await _observe(cache, body, caller="other-caller")
|
||||
response: Final = await _post(_app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER)), body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.switch.cache_state == "unknown"
|
||||
assert result.switch.reason == "no_compatible_observation"
|
||||
assert result.switch.evidence is None
|
||||
assert result.switch.estimate is not None
|
||||
assert result.switch.estimate.tokens.cache_read_input_tokens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_failure_nulls_switch_comparison(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
app: Final = _app(
|
||||
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=Counts(total=None)
|
||||
)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.reason == result.switch.reason == "token_count_unavailable"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("limiter", [None, CustomLogger()])
|
||||
async def test_missing_or_unsupported_limiter_returns_unknown_before_counting(
|
||||
monkeypatch: pytest.MonkeyPatch, limiter: CustomLogger | None
|
||||
) -> None:
|
||||
app: Final = _app(
|
||||
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count, limiter=limiter
|
||||
)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.reason == result.switch.reason == "limiter_unavailable"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_occupied_parallel_capacity_rejects_before_provider_count(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache: Final = DualCache()
|
||||
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
|
||||
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
|
||||
app: Final = _app(monkeypatch, cache, caller=caller, counts=_unexpected_count, limiter=limiter)
|
||||
async with limiter.request_capacity(caller, "opus"):
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 429, response.text
|
||||
assert "max_parallel_requests" in response.text
|
||||
recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body())
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_count_consumes_the_deployment_group_rpm_limit(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls: Final = asyncio.Queue[str]()
|
||||
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
calls.put_nowait(model)
|
||||
return await Counts()(model, api_key, body)
|
||||
|
||||
caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"model_rpm_limit": {"sonnet": 1}})
|
||||
app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 429, response.text
|
||||
assert calls.qsize() == 3
|
||||
assert tuple(calls.get_nowait() for _ in range(3)) == (
|
||||
"claude-opus-5", "claude-opus-5", "claude-sonnet-5"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
async def test_each_count_preserves_auth_cached_request_tag_limits(
|
||||
monkeypatch: pytest.MonkeyPatch, metadata_key: str
|
||||
) -> None:
|
||||
calls: Final = asyncio.Queue[str]()
|
||||
caller: Final = UserAPIKeyAuth(api_key=_CALLER, metadata={"tag_rpm_limit": {"cache-cost": 1}})
|
||||
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
calls.put_nowait(model)
|
||||
return await Counts()(model, api_key, body)
|
||||
|
||||
async def authenticated_request(request: Request) -> UserAPIKeyAuth:
|
||||
data: Final = await _read_request_body(request)
|
||||
_safe_set_request_parsed_body(request, {**data, metadata_key: {"tags": ["cache-cost"]}})
|
||||
return caller
|
||||
|
||||
app: Final = _app(monkeypatch, DualCache(), caller=caller, counts=count)
|
||||
app.dependency_overrides[endpoint.user_api_key_auth] = authenticated_request
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 429, response.text
|
||||
assert "tag_per_key" in response.text
|
||||
assert calls.qsize() == 1
|
||||
assert calls.get_nowait() == "claude-opus-5"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_counter_failure_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache: Final = DualCache()
|
||||
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
|
||||
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
|
||||
|
||||
async def fail_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
raise RuntimeError("provider counter failed")
|
||||
|
||||
app: Final = _app(monkeypatch, cache, caller=caller, counts=fail_count, limiter=limiter)
|
||||
with pytest.raises(RuntimeError, match="provider counter failed"):
|
||||
await _post(app, _body())
|
||||
recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body())
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_provider_counter_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache: Final = DualCache()
|
||||
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
|
||||
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def wait_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
started.set()
|
||||
await release.wait()
|
||||
return await Counts()(model, api_key, body)
|
||||
|
||||
app: Final = _app(monkeypatch, cache, caller=caller, counts=wait_count, limiter=limiter)
|
||||
pending: Final = asyncio.create_task(_post(app, _body()))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), timeout=5)
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await pending
|
||||
release.set()
|
||||
recovered: Final = await asyncio.wait_for(_post(app, _body()), timeout=5)
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
|
||||
finally:
|
||||
pending.cancel()
|
||||
release.set()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
|
||||
|
||||
async def _unexpected_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
|
||||
pytest.fail("Unsupported prediction must return before contacting the token counter")
|
||||
|
||||
|
||||
class RequestMutator(CustomLogger):
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
|
||||
) -> dict[str, object]:
|
||||
return {**data, "system": "Injected policy"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def request_mutator() -> Iterator[RequestMutator]:
|
||||
callback: Final = RequestMutator()
|
||||
litellm.logging_callback_manager.add_litellm_callback(callback)
|
||||
try:
|
||||
yield callback
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_transform_callback_returns_unknown_before_token_counting(
|
||||
monkeypatch: pytest.MonkeyPatch, request_mutator: RequestMutator
|
||||
) -> None:
|
||||
app: Final = _app(
|
||||
monkeypatch, DualCache(), caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count
|
||||
)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.cache_state == result.switch.cache_state == "unknown"
|
||||
assert result.stay.reason == result.switch.reason == "unsupported_request_transform"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_config_returns_unknown_before_token_counting(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
app: Final = _app(
|
||||
monkeypatch,
|
||||
DualCache(),
|
||||
caller=UserAPIKeyAuth(api_key=_CALLER, config={"model_list": []}),
|
||||
counts=_unexpected_count,
|
||||
)
|
||||
response: Final = await _post(app, _body())
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.cache_state == result.switch.cache_state == "unknown"
|
||||
assert result.stay.reason == result.switch.reason == "unsupported_request_transform"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers", [
|
||||
{"anthropic-version": "2099-01-01"},
|
||||
{"anthropic-beta": "future-feature"},
|
||||
])
|
||||
@pytest.mark.asyncio
|
||||
async def test_unsupported_provider_headers_cannot_reuse_default_version_evidence(
|
||||
monkeypatch: pytest.MonkeyPatch, headers: dict[str, str]
|
||||
) -> None:
|
||||
cache: Final = DualCache()
|
||||
await _observe(cache, _body(), deployment_id="sonnet")
|
||||
app: Final = _app(
|
||||
monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER), counts=_unexpected_count
|
||||
)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response: Final = await client.post(
|
||||
"/cost/predict-cache",
|
||||
headers=headers,
|
||||
json={"current_deployment_id": "opus", "candidate_deployment_id": "sonnet", "request": _body()},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
result: Final = CachePredictionResponse.model_validate(response.json())
|
||||
assert result.stay.cache_state == result.switch.cache_state == "unknown"
|
||||
assert result.stay.reason == result.switch.reason == "unsupported_provider_headers"
|
||||
assert result.stay.estimate is None and result.switch.estimate is None
|
||||
assert result.switch_delta is None and result.cache_rebuild_penalty is None
|
||||
164
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
164
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -3427,6 +3427,35 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/cost/predict-cache": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
get?: never;
|
||||
put?: never;
|
||||
/**
|
||||
* Predict Cache Cost
|
||||
* @description Compare the next native Anthropic request on two configured deployment IDs.
|
||||
*
|
||||
* Estimates use provider token counting and recent successful cache telemetry for this key.
|
||||
* Unknown cache state uses the cold scenario when prices/counts are available. Cache observations
|
||||
* do not guarantee retention. v0 supports one message-content breakpoint, text and client tools;
|
||||
* system/tool-only breakpoints, thinking, images, nondefault Anthropic versions, beta headers and
|
||||
* request transforms are unknown.
|
||||
* Each provider count consumes one RPM unit and holds concurrency capacity; a comparison uses
|
||||
* up to four counts. The legacy rate limiter returns unknown without contacting the provider.
|
||||
* This endpoint does not generate tokens, prewarm caches, choose a model or alter routing.
|
||||
*/
|
||||
post: operations["predict_cache_cost_cost_predict_cache_post"];
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/credentials": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -24702,6 +24731,31 @@ export interface components {
|
|||
/** Failed Requests */
|
||||
failed_requests: number;
|
||||
};
|
||||
/** CacheCostScenario */
|
||||
CacheCostScenario: {
|
||||
/** Input Cost */
|
||||
input_cost: number;
|
||||
tokens: components["schemas"]["CacheTokenBuckets"];
|
||||
};
|
||||
/** CacheEvidence */
|
||||
CacheEvidence: {
|
||||
/**
|
||||
* Confidence
|
||||
* @default observed
|
||||
* @constant
|
||||
*/
|
||||
confidence: "observed";
|
||||
/** Expires At */
|
||||
expires_at: number;
|
||||
/** Observed At */
|
||||
observed_at: number;
|
||||
/**
|
||||
* Source
|
||||
* @default provider_usage
|
||||
* @constant
|
||||
*/
|
||||
source: "provider_usage";
|
||||
};
|
||||
/** CachePingResponse */
|
||||
CachePingResponse: {
|
||||
/** Cache Type */
|
||||
|
|
@ -24719,6 +24773,59 @@ export interface components {
|
|||
/** Status */
|
||||
status: string;
|
||||
};
|
||||
/** CachePredictionArm */
|
||||
CachePredictionArm: {
|
||||
/**
|
||||
* Cache State
|
||||
* @default unknown
|
||||
* @enum {string}
|
||||
*/
|
||||
cache_state: "warm" | "partial" | "stale" | "unknown" | "disabled";
|
||||
cold?: components["schemas"]["CacheCostScenario"] | null;
|
||||
/** Deployment Id */
|
||||
deployment_id: string;
|
||||
estimate?: components["schemas"]["CacheCostScenario"] | null;
|
||||
evidence?: components["schemas"]["CacheEvidence"] | null;
|
||||
/** Model */
|
||||
model?: string | null;
|
||||
/** Reason */
|
||||
reason?: string | null;
|
||||
/** Token Count Source */
|
||||
token_count_source?: "anthropic_count_tokens" | null;
|
||||
warm?: components["schemas"]["CacheCostScenario"] | null;
|
||||
};
|
||||
/** CachePredictionRequest */
|
||||
CachePredictionRequest: {
|
||||
/** Candidate Deployment Id */
|
||||
candidate_deployment_id: string;
|
||||
/** Current Deployment Id */
|
||||
current_deployment_id: string;
|
||||
/** Request */
|
||||
request: {
|
||||
[key: string]: components["schemas"]["JsonValue"];
|
||||
};
|
||||
};
|
||||
/** CachePredictionResponse */
|
||||
CachePredictionResponse: {
|
||||
/**
|
||||
* Cache Guarantee
|
||||
* @default false
|
||||
* @constant
|
||||
*/
|
||||
cache_guarantee: false;
|
||||
/** Cache Rebuild Penalty */
|
||||
cache_rebuild_penalty: number | null;
|
||||
/**
|
||||
* Pricing Basis
|
||||
* @default input_before_discounts_and_margins
|
||||
* @constant
|
||||
*/
|
||||
pricing_basis: "input_before_discounts_and_margins";
|
||||
stay: components["schemas"]["CachePredictionArm"];
|
||||
switch: components["schemas"]["CachePredictionArm"];
|
||||
/** Switch Delta */
|
||||
switch_delta: number | null;
|
||||
};
|
||||
/** CacheSettingsField */
|
||||
CacheSettingsField: {
|
||||
/** Field Default */
|
||||
|
|
@ -24800,6 +24907,29 @@ export interface components {
|
|||
*/
|
||||
status: string;
|
||||
};
|
||||
/** CacheTokenBuckets */
|
||||
CacheTokenBuckets: {
|
||||
/**
|
||||
* Cache Creation 1H Input Tokens
|
||||
* @default 0
|
||||
*/
|
||||
cache_creation_1h_input_tokens: number;
|
||||
/**
|
||||
* Cache Creation 5M Input Tokens
|
||||
* @default 0
|
||||
*/
|
||||
cache_creation_5m_input_tokens: number;
|
||||
/**
|
||||
* Cache Read Input Tokens
|
||||
* @default 0
|
||||
*/
|
||||
cache_read_input_tokens: number;
|
||||
/**
|
||||
* Uncached Input Tokens
|
||||
* @default 0
|
||||
*/
|
||||
uncached_input_tokens: number;
|
||||
};
|
||||
/**
|
||||
* CallTypes
|
||||
* @enum {string}
|
||||
|
|
@ -28421,6 +28551,7 @@ export interface components {
|
|||
/** Updated By */
|
||||
updated_by?: string | null;
|
||||
};
|
||||
JsonValue: unknown;
|
||||
/** KeyHealthResponse */
|
||||
KeyHealthResponse: {
|
||||
/**
|
||||
|
|
@ -45322,6 +45453,39 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
predict_cache_cost_cost_predict_cache_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody: {
|
||||
content: {
|
||||
"application/json": components["schemas"]["CachePredictionRequest"];
|
||||
};
|
||||
};
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["CachePredictionResponse"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
get_credentials_credentials_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue