mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(cost_tracking): bill cache writes at the request's service-tier rate
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
8177230a29
commit
1af27283d5
2 changed files with 47 additions and 2 deletions
|
|
@ -199,6 +199,7 @@ from litellm.types.utils import (
|
|||
SandboxProviders,
|
||||
SearchProviders,
|
||||
SelectTokenizerResponse,
|
||||
ServiceTier,
|
||||
StreamingChoices,
|
||||
TextChoices,
|
||||
TextCompletionResponse,
|
||||
|
|
@ -5219,7 +5220,13 @@ def _is_potential_model_name_in_model_cost(
|
|||
)
|
||||
|
||||
|
||||
_ABOVE_THRESHOLD_COST_KEY = re.compile(r"_above_\d+k?_tokens$")
|
||||
# Cost keys the cost calculator composes at request time (context-length
|
||||
# thresholds and service tiers) instead of reading from a declared ModelInfo
|
||||
# field, e.g. "cache_creation_input_token_cost_priority" or
|
||||
# "input_cost_per_token_above_272k_tokens". They are copied onto ModelInfo
|
||||
# verbatim so a new pricing dimension in the cost map does not need a matching
|
||||
# field before it can be billed.
|
||||
_DYNAMIC_COST_KEY = re.compile(r"cost.*(?:_above_\d+k?_tokens|_(?:" + "|".join(st.value for st in ServiceTier) + r"))$")
|
||||
|
||||
|
||||
def _get_model_info_helper(
|
||||
|
|
@ -5535,7 +5542,7 @@ def _get_model_info_helper(
|
|||
supports_image_size=_model_info.get("supports_image_size", None),
|
||||
)
|
||||
for cost_key, cost_value in _model_info.items():
|
||||
if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None:
|
||||
if cost_key not in returned_model_info and _DYNAMIC_COST_KEY.search(cost_key) is not None:
|
||||
returned_model_info[cost_key] = cost_value # type: ignore[literal-required]
|
||||
return returned_model_info
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -2399,3 +2399,41 @@ def test_generic_cost_per_token_gemini_35_flash_lite():
|
|||
)
|
||||
assert prompt_cost == pytest.approx(0.0003)
|
||||
assert completion_cost == pytest.approx(0.00125)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"service_tier, expected_cache_write_rate, expected_input_rate, expected_output_rate",
|
||||
[
|
||||
(None, 6.25e-6, 5e-6, 3e-5),
|
||||
("priority", 1.25e-5, 1e-5, 6e-5),
|
||||
("flex", 3.125e-6, 2.5e-6, 1.5e-5),
|
||||
],
|
||||
)
|
||||
def test_service_tier_cache_write_pricing(
|
||||
service_tier,
|
||||
expected_cache_write_rate,
|
||||
expected_input_rate,
|
||||
expected_output_rate,
|
||||
_local_model_cost_map,
|
||||
):
|
||||
"""Regression: cache-write tokens on a service-tier request must bill at that tier's
|
||||
cache_creation_input_token_cost_<tier> rate. gpt-5.6-sol publishes $6.25/M standard,
|
||||
$12.50/M priority and $3.125/M flex; billing every tier at the standard rate
|
||||
under-charges priority by 50% and over-charges flex by 100%."""
|
||||
usage = Usage(
|
||||
prompt_tokens=4_020,
|
||||
completion_tokens=4,
|
||||
total_tokens=4_024,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_write_tokens=4_017),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="gpt-5.6-sol",
|
||||
usage=usage,
|
||||
custom_llm_provider="openai",
|
||||
service_tier=service_tier,
|
||||
)
|
||||
|
||||
expected_prompt = 4_017 * expected_cache_write_rate + 3 * expected_input_rate
|
||||
assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
|
||||
assert completion_cost == pytest.approx(4 * expected_output_rate, rel=1e-9)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue