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:
shivam 2026-07-24 22:34:11 +00:00
parent 8177230a29
commit 1af27283d5
2 changed files with 47 additions and 2 deletions

View file

@ -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:

View file

@ -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)