test(cost): cover custom cache pricing fields

This commit is contained in:
Kcstring 2026-04-30 20:56:57 +08:00
parent a3f24d0f6d
commit be73bf1472
2 changed files with 65 additions and 3 deletions

View file

@ -40,7 +40,7 @@ from pydantic import (
field_validator,
model_validator,
)
from typing_extensions import Required, TypedDict
from typing_extensions import NotRequired, Required, TypedDict
from litellm._uuid import uuid
from litellm.types.llms.base import (
@ -110,6 +110,8 @@ SupportedCacheControls = ["ttl", "s-maxage", "no-cache", "no-store"]
class CostPerToken(TypedDict):
input_cost_per_token: float
output_cost_per_token: float
cache_creation_input_token_cost: NotRequired[float]
cache_read_input_token_cost: NotRequired[float]
class ProviderField(TypedDict):
@ -1264,7 +1266,7 @@ class Delta(SafeAttributeModel, OpenAIObject):
reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None
provider_specific_fields: Optional[Dict[str, Any]] = Field(default=None)
def __init__(
def __init__( # noqa: PLR0915
self,
content=None,
role=None,

View file

@ -31,7 +31,67 @@ def test_custom_cost_per_token_uses_cache_read_pricing():
"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)
def test_custom_cost_per_token_uses_cache_read_usage_field():
usage = Usage(
prompt_tokens=100,
completion_tokens=10,
total_tokens=110,
cache_read_input_tokens=40,
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/custom",
object="chat.completion",
choices=[],
usage=usage,
)
cost = completion_cost(
completion_response=response,
model="openai/custom",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_read_input_token_cost": 0.001,
},
)
assert cost == pytest.approx((60 * 0.01) + (40 * 0.001) + (10 * 0.02))
def test_custom_cost_per_token_uses_cache_creation_pricing():
usage = Usage(
prompt_tokens=100,
completion_tokens=10,
total_tokens=110,
cache_creation_input_tokens=30,
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/custom",
object="chat.completion",
choices=[],
usage=usage,
)
cost = completion_cost(
completion_response=response,
model="openai/custom",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_creation_input_token_cost": 0.015,
},
)
assert cost == pytest.approx((70 * 0.01) + (30 * 0.015) + (10 * 0.02))