perf(cost-calc): precompute service-tier cost-key suffixes (#31431)

_get_token_base_cost rebuilt f"_{st.value}" for every ServiceTier while
scanning every model_info key on each request, and _get_cost_per_unit rebuilt
the same f-strings in its fallback loop. The suffixes are constant, so compute
them once at module level (matching the existing _IMAGE_RESPONSE_CALL_TYPES /
_VALID_DATA_RESIDENCIES pattern) and use str.endswith(tuple) for the threshold
check. Behavior is identical; ~3.5 us/call to ~2.2 us/call on the threshold scan.
This commit is contained in:
Yassin Kortam 2026-06-26 19:38:58 +03:00 committed by GitHub
parent fc644cff3d
commit 0e1a3babf0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 56 additions and 4 deletions

View file

@ -33,6 +33,11 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset(
# Pre-resolved DataResidency enum values for fast membership checks
_VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency)
# Pre-resolved service-tier cost-key suffixes (e.g. "_priority"). Used per
# request in the cost-calc path, so the f-strings are built once here instead
# of being rebuilt for every model_info key on every call.
_SERVICE_TIER_SUFFIXES: tuple[str, ...] = tuple(f"_{st.value}" for st in ServiceTier)
def _get_token_detail_value(details: object, key: str) -> Optional[int]:
if isinstance(details, dict):
@ -252,7 +257,7 @@ def _get_token_base_cost(
k
for k in model_info
if k.startswith("input_cost_per_token_above_")
and not any(k.endswith(f"_{st.value}") for st in ServiceTier)
and not k.endswith(_SERVICE_TIER_SUFFIXES)
]
if not threshold_keys:
return (
@ -418,9 +423,8 @@ def _get_cost_per_unit(
# If the service tier key doesn't exist or is None, try to fall back to the standard key
if cost_per_unit is None:
# Check if any service tier suffix exists in the cost key using ServiceTier enum
for service_tier in ServiceTier:
suffix = f"_{service_tier.value}"
# Check if any service tier suffix exists in the cost key
for suffix in _SERVICE_TIER_SUFFIXES:
if suffix in cost_key:
# Extract the base key by removing the matched suffix
base_key = cost_key.replace(suffix, "")

View file

@ -1720,3 +1720,51 @@ def test_priority_service_tier_above_threshold_falls_back_to_standard_for_cache_
expected_completion = 1_000 * 2.25e-5
assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
assert completion_cost == pytest.approx(expected_completion, rel=1e-9)
def test_service_tier_suffixes_constant_in_sync_with_enum():
from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES
from litellm.types.utils import ServiceTier
assert _SERVICE_TIER_SUFFIXES == tuple(f"_{st.value}" for st in ServiceTier)
def test_get_cost_per_unit_falls_back_from_service_tier_key_to_base():
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit
model_info = {"input_cost_per_token": 2e-6}
# service-tier key is absent -> falls back to the base key
assert _get_cost_per_unit(model_info, "input_cost_per_token_priority") == 2e-6
# service-tier key present -> used directly, no fallback
model_info_direct = {
"input_cost_per_token_priority": 5e-6,
"input_cost_per_token": 2e-6,
}
assert (
_get_cost_per_unit(model_info_direct, "input_cost_per_token_priority") == 5e-6
)
def test_threshold_keys_exclude_service_tier_variants():
from typing import cast
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
from litellm.types.utils import ModelInfo, Usage
# The service-tier-suffixed above-threshold key must be excluded from
# threshold detection. The _priority variant has a higher threshold (300k),
# so if it were not excluded it would sort first and drive a 9e-6 rate for
# this non-tier request. With the exclusion only the standard 200k key
# applies, giving 3e-6.
model_info = cast(
ModelInfo,
{
"input_cost_per_token": 1e-6,
"input_cost_per_token_above_200k_tokens": 3e-6,
"input_cost_per_token_above_300k_tokens_priority": 9e-6,
"output_cost_per_token": 2e-6,
},
)
usage = Usage(prompt_tokens=350_000, completion_tokens=1_000, total_tokens=351_000)
prompt_base, *_ = _get_token_base_cost(model_info=model_info, usage=usage)
assert prompt_base == 3e-6