mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
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:
parent
fc644cff3d
commit
0e1a3babf0
2 changed files with 56 additions and 4 deletions
|
|
@ -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, "")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue