diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d0bc9b78941..649d0a46495 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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, diff --git a/tests/test_litellm/test_custom_pricing_cache_cost.py b/tests/test_litellm/test_custom_pricing_cache_cost.py index 36c7e161f6a..b4b7beb54c1 100644 --- a/tests/test_litellm/test_custom_pricing_cache_cost.py +++ b/tests/test_litellm/test_custom_pricing_cache_cost.py @@ -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))