scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Fix tiered custom pricing for cache tokens

Browse source
This commit is contained in:
its-amann 2026-06-11 19:40:14 +05:30
parent d698106ed7
commit 1f9fe73bd9
4 changed files with 155 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/litellm_core_utils/litellm_logging.py
View file

@ -7219,7 +7219,8 @@ def get_standard_logging_object_payload(
custom_llm_provider=kwargs.get("custom_llm_provider"),
init_response_obj=init_response_obj,
api_base=litellm_params.get("api_base"),
router_model_id=logging_obj.get_router_model_id(),
router_model_id=logging_obj.get_router_model_id()
or (_model_id if isinstance(_model_id, str) and _model_id else None),
)
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(

22
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -288,7 +288,7 @@ def _parse_above_token_threshold(key: str) -> float:
return _parse_token_threshold(key.rsplit("_above_", 1)[1].split("_tokens")[0])
def _prompt_exceeds_threshold(prompt_tokens: int, threshold: float, inclusive: bool) -> bool:
def _prompt_exceeds_threshold(prompt_tokens: float, threshold: float, inclusive: bool) -> bool:
return prompt_tokens > threshold or (inclusive and prompt_tokens == threshold)
@ -689,6 +689,23 @@ def _apply_off_peak_to_base_costs(
)
def _loose_model_extra_token(usage: Usage, field_name: str) -> float:
extra: Final = usage.model_extra
if not isinstance(extra, dict) or field_name not in extra:
return 0.0
token: Final = extra.get(field_name)
if isinstance(token, bool) or not isinstance(token, (int, float)):
return 0.0
return float(token)
def _get_effective_prompt_tokens_for_tiered_pricing(usage: Usage) -> float:
prompt_tokens: Final = float(usage.prompt_tokens or 0)
cache_read_tokens: Final = _loose_model_extra_token(usage, "cache_read_input_tokens")
cache_creation_tokens: Final = _loose_model_extra_token(usage, "cache_creation_input_tokens")
return prompt_tokens + cache_read_tokens + cache_creation_tokens
def _get_token_base_cost(
model_info: ModelInfo,
usage: Usage,
@ -761,6 +778,7 @@ def _get_token_base_cost(
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
threshold: float | None = None
effective_prompt_tokens: Final = _get_effective_prompt_tokens_for_tiered_pricing(usage)
for key in sorted(threshold_keys, key=_parse_above_token_threshold):
value = model_info.get(key)
if value is not None:
@ -768,7 +786,7 @@ def _get_token_base_cost(
# Handle both formats: _above_128k_tokens and _above_128_tokens
threshold_str = key.rsplit("_above_", 1)[1].split("_tokens")[0]
threshold = _parse_above_token_threshold(key)
if _prompt_exceeds_threshold(usage.prompt_tokens, threshold, threshold_is_inclusive):
if _prompt_exceeds_threshold(effective_prompt_tokens, threshold, threshold_is_inclusive):
# Prefer a service_tier-specific above-threshold key when available,
# e.g. input_cost_per_token_priority_above_200k_tokens for Gemini
# ON_DEMAND_PRIORITY. Falls back to the standard key automatically

81
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -5701,6 +5701,87 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload["litellm_call_id"] == call_id
def test_standard_logging_model_map_uses_router_model_info_for_custom_pricing():
"""Spend-log model_map_value keeps the deployment above-200k rates from request metadata."""
from litellm.litellm_core_utils.litellm_logging import (
Logging,
get_standard_logging_object_payload,
)
model_id = "tiered-custom-pricing-logging-test"
model_info = {
"litellm_provider": "vertex_ai",
"mode": "chat",
"input_cost_per_token": 0.000005,
"output_cost_per_token": 0.000025,
"cache_creation_input_token_cost": 0.00000625,
"cache_read_input_token_cost": 0.0000005,
"input_cost_per_token_above_200k_tokens": 0.00001,
"output_cost_per_token_above_200k_tokens": 0.0000375,
"cache_creation_input_token_cost_above_200k_tokens": 0.0000125,
"cache_read_input_token_cost_above_200k_tokens": 0.000001,
}
litellm.register_model(model_cost={model_id: model_info})
try:
logging_obj = Logging(
model="vertex_ai/claude-opus-4-6",
messages=[{"role": "user", "content": "Hello"}],
stream=False,
call_type="completion",
start_time=datetime.datetime.now(),
litellm_call_id="test-tiered-pricing",
function_id="test-function",
)
payload = get_standard_logging_object_payload(
kwargs={
"model": "vertex_ai/claude-opus-4-6",
"messages": [{"role": "user", "content": "Hello"}],
"response_cost": 1.0,
"custom_llm_provider": "vertex_ai",
"litellm_params": {
"api_base": "",
"metadata": {
"model_info": {
"id": model_id,
**model_info,
},
},
},
},
init_response_obj={
"id": "chatcmpl-test",
"model": "vertex_ai/claude-opus-4-6",
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 585659,
"cache_creation_input_tokens": 585659,
},
"choices": [],
},
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
logging_obj=logging_obj,
status="success",
)
assert payload is not None
model_map_value = payload["model_map_information"]["model_map_value"]
assert model_map_value is not None
assert (
model_map_value["cache_creation_input_token_cost_above_200k_tokens"]
== model_info["cache_creation_input_token_cost_above_200k_tokens"]
)
assert (
model_map_value["input_cost_per_token_above_200k_tokens"]
== model_info["input_cost_per_token_above_200k_tokens"]
)
finally:
litellm.model_cost.pop(model_id, None)
litellm.model_cost.pop(f"vertex_ai/{model_id}", None)
@pytest.mark.parametrize(
"hidden_params, passed_cache_hit, expected_cache_hit, expected_cost, expected_unique_id",
[

52
tests/unit/test_cost_calculator.py
View file

@ -3529,6 +3529,58 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens
assert completion_cost == pytest.approx(expected_completion)
def test_custom_pricing_above_200k_cache_creation_tokens_uses_tiered_rate():
"""Cache creation tokens reported outside prompt_tokens still cross the above-200k tier."""
model_id = "tiered-custom-pricing-deploy-test"
model_info = {
"litellm_provider": "vertex_ai",
"mode": "chat",
"input_cost_per_token": 0.000005,
"output_cost_per_token": 0.000025,
"cache_creation_input_token_cost": 0.00000625,
"cache_read_input_token_cost": 0.0000005,
"input_cost_per_token_above_200k_tokens": 0.00001,
"output_cost_per_token_above_200k_tokens": 0.0000375,
"cache_creation_input_token_cost_above_200k_tokens": 0.0000125,
"cache_read_input_token_cost_above_200k_tokens": 0.000001,
}
cache_creation_tokens = 585659
litellm.register_model(model_cost={model_id: model_info})
try:
response = ModelResponse(
id="test-id",
created=1234567890,
model="vertex_ai/claude-opus-4-6",
object="chat.completion",
choices=[],
usage=Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=cache_creation_tokens,
cache_creation_input_tokens=cache_creation_tokens,
cache_read_input_tokens=0,
),
)
cost = completion_cost(
completion_response=response,
model="vertex_ai/claude-opus-4-6",
custom_llm_provider="vertex_ai",
custom_pricing=True,
router_model_id=model_id,
)
assert cost == pytest.approx(
cache_creation_tokens * model_info["cache_creation_input_token_cost_above_200k_tokens"]
)
assert cost != pytest.approx(cache_creation_tokens * model_info["cache_creation_input_token_cost"])
finally:
litellm.model_cost.pop(model_id, None)
litellm.model_cost.pop(f"vertex_ai/{model_id}", None)
# ---------------------------------------------------------------------------
# Bug 2 — db_spend_update_writer cache token extraction helpers.
# ---------------------------------------------------------------------------