fix: use top-level cache read tokens in custom pricing

This commit is contained in:
Genmin 2026-05-01 07:31:17 -07:00
parent 92e8e98750
commit 023e05b0e3
2 changed files with 34 additions and 3 deletions

View file

@ -190,9 +190,10 @@ def _cost_per_token_custom_pricing_helper(
)
if cache_read_input_token_cost is not None:
if not cache_read_input_tokens and usage_object is not None:
cache_read_input_tokens = _parse_prompt_tokens_details(usage_object)[
"cache_hit_tokens"
]
cache_read_input_tokens = (
getattr(usage_object, "cache_read_input_tokens", None)
or _parse_prompt_tokens_details(usage_object)["cache_hit_tokens"]
)
cache_read_tokens = max(0, cache_read_input_tokens or 0)
uncached_prompt_tokens = max(0, prompt_tokens - cache_read_tokens)
input_cost = (

View file

@ -81,6 +81,36 @@ def test_completion_cost_custom_pricing_uses_cache_read_rate():
assert cost == pytest.approx(0.011684)
def test_completion_cost_custom_pricing_uses_top_level_cache_read_tokens():
usage = Usage(
prompt_tokens=100,
completion_tokens=10,
total_tokens=110,
cache_read_input_tokens=80,
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="anthropic/claude-sonnet-4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = completion_cost(
completion_response=response,
model="anthropic/claude-sonnet-4",
custom_llm_provider="anthropic",
custom_cost_per_token={
"input_cost_per_token": 1.0,
"output_cost_per_token": 2.0,
"cache_read_input_token_cost": 0.25,
},
)
assert cost == pytest.approx(60.0)
def test_completion_cost_custom_pricing_without_cache_read_rate_preserves_input_rate():
usage = Usage(
prompt_tokens=6074,