fix: apply cached token rate for custom pricing

This commit is contained in:
Genmin 2026-04-30 09:56:36 -07:00
parent 3583ac1159
commit 92e8e98750
2 changed files with 86 additions and 1 deletions

View file

@ -172,6 +172,8 @@ _MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
completion_tokens: float = 0,
cache_read_input_tokens: Optional[int] = 0,
usage_object: Optional[Usage] = None,
response_time_ms: Optional[float] = 0.0,
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
@ -182,7 +184,23 @@ def _cost_per_token_custom_pricing_helper(
return None
if custom_cost_per_token is not None:
input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens
cache_read_input_token_cost = cast(
Optional[float],
custom_cost_per_token.get("cache_read_input_token_cost"),
)
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_tokens = max(0, cache_read_input_tokens or 0)
uncached_prompt_tokens = max(0, prompt_tokens - cache_read_tokens)
input_cost = (
custom_cost_per_token["input_cost_per_token"] * uncached_prompt_tokens
+ cache_read_input_token_cost * cache_read_tokens
)
else:
input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens
output_cost = custom_cost_per_token["output_cost_per_token"] * completion_tokens
return input_cost, output_cost
elif custom_cost_per_second is not None:
@ -326,6 +344,8 @@ def cost_per_token( # noqa: PLR0915
response_cost = _cost_per_token_custom_pricing_helper(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
cache_read_input_tokens=cache_read_input_tokens,
usage_object=usage_block,
response_time_ms=response_time_ms,
custom_cost_per_second=custom_cost_per_second,
custom_cost_per_token=custom_cost_per_token,

View file

@ -48,6 +48,71 @@ def test_completion_cost_uses_response_model_for_dynamic_routing():
assert cost > 0, "Cost should be calculated using response model"
def test_completion_cost_custom_pricing_uses_cache_read_rate():
usage = Usage(
prompt_tokens=6074,
completion_tokens=285,
total_tokens=6359,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=3456,
audio_tokens=0,
),
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/gpt-5.4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = completion_cost(
completion_response=response,
model="openai/gpt-5.4",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
"cache_read_input_token_cost": 0.00000025,
},
)
assert cost == pytest.approx(0.011684)
def test_completion_cost_custom_pricing_without_cache_read_rate_preserves_input_rate():
usage = Usage(
prompt_tokens=6074,
completion_tokens=285,
total_tokens=6359,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=3456,
audio_tokens=0,
),
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/gpt-5.4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = completion_cost(
completion_response=response,
model="openai/gpt-5.4",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
},
)
assert cost == pytest.approx(0.01946)
def test_cost_calculator_with_response_cost_in_additional_headers():
class MockResponse(BaseModel):
_hidden_params = {