mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
test(cost): cover custom cache pricing fields
This commit is contained in:
parent
a3f24d0f6d
commit
be73bf1472
2 changed files with 65 additions and 3 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue