mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix: apply cached token rate for custom pricing
This commit is contained in:
parent
3583ac1159
commit
92e8e98750
2 changed files with 86 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue