scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(cost): apply cache pricing for custom token costs

Browse source
This commit is contained in:
Kcstring 2026-04-30 20:44:07 +08:00
parent 2f22a1293e
commit a3f24d0f6d
2 changed files with 69 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/cost_calculator.py
View file

@ -172,6 +172,7 @@ _MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
completion_tokens: float = 0,
usage_object: Optional[Usage] = None,
response_time_ms: Optional[float] = 0.0,
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
@ -182,7 +183,36 @@ 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
input_cost_per_token = custom_cost_per_token["input_cost_per_token"]
custom_cost_map = cast(dict[str, float], custom_cost_per_token)
cache_read_input_token_cost = custom_cost_map.get(
"cache_read_input_token_cost", input_cost_per_token
)
cache_creation_input_token_cost = custom_cost_map.get(
"cache_creation_input_token_cost", input_cost_per_token
)
cache_read_tokens = 0
cache_creation_tokens = 0
if usage_object is not None:
if usage_object.prompt_tokens_details is not None:
prompt_tokens_details = _parse_prompt_tokens_details(usage_object)
cache_read_tokens = prompt_tokens_details["cache_hit_tokens"]
cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"]
cache_read_tokens = cache_read_tokens or (
getattr(usage_object, "cache_read_input_tokens", None) or 0
)
cache_creation_tokens = cache_creation_tokens or (
getattr(usage_object, "cache_creation_input_tokens", None) or 0
)
uncached_prompt_tokens = max(
0, prompt_tokens - cache_read_tokens - cache_creation_tokens
)
input_cost = (
uncached_prompt_tokens * input_cost_per_token
+ cache_read_tokens * cache_read_input_token_cost
+ cache_creation_tokens * cache_creation_input_token_cost
)
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 +356,7 @@ def cost_per_token( # noqa: PLR0915
response_cost = _cost_per_token_custom_pricing_helper(
prompt_tokens=prompt_tokens,
completion_tokens=completion_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,

37
tests/test_litellm/test_custom_pricing_cache_cost.py Normal file
View file

@ -0,0 +1,37 @@
import pytest
from litellm import completion_cost
from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
def test_custom_cost_per_token_uses_cache_read_pricing():
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,
}, # type: ignore[typeddict-unknown-key]
)
assert cost == pytest.approx(0.011684)