diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 962d4ea4c10..da420ab0578 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response from pydantic import BaseModel -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm import litellm._logging @@ -1034,6 +1034,8 @@ def get_usage_object( ), ) + if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)): + return None if completion_response.usage is None else _decisions_usage(completion_response.usage) if usage_obj is None: return None if isinstance(usage_obj, Usage): @@ -1064,15 +1066,28 @@ def get_usage_object( return None -def _decisions_usage( - *, input_tokens: int, output_tokens: int, cached_tokens: int, cache_write_tokens: int -) -> dict[str, object]: - return { - "prompt_tokens": input_tokens, - "completion_tokens": output_tokens, - "cache_read_input_tokens": cached_tokens, - "cache_creation_input_tokens": cache_write_tokens, - } +def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper: + match usage: + case DecisionsUsage(): + return PromptTokensDetailsWrapper( + cached_tokens=usage.cached_tokens, cache_write_tokens=usage.cache_write_tokens + ) + case OpenAIDecisionUsage(): + return PromptTokensDetailsWrapper( + cached_tokens=usage.input_tokens_details.cached_tokens, + cache_write_tokens=usage.input_tokens_details.cache_write_tokens, + ) + case _: + assert_never(usage) + + +def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage: + return Usage( + prompt_tokens=usage.input_tokens, + completion_tokens=usage.output_tokens, + total_tokens=usage.input_tokens + usage.output_tokens, + prompt_tokens_details=_decisions_prompt_tokens_details(usage), + ) def _is_known_usage_objects(usage_obj): @@ -1493,20 +1508,8 @@ def completion_cost( "usage", litellm.Usage(**_usage_for_dump.model_dump()), ) - if isinstance(usage_obj, DecisionsUsage): - _usage = _decisions_usage( - input_tokens=usage_obj.input_tokens, - output_tokens=usage_obj.output_tokens, - cached_tokens=usage_obj.cached_tokens, - cache_write_tokens=usage_obj.cache_write_tokens, - ) - elif isinstance(usage_obj, OpenAIDecisionUsage): - _usage = _decisions_usage( - input_tokens=usage_obj.input_tokens, - output_tokens=usage_obj.output_tokens, - cached_tokens=usage_obj.input_tokens_details.cached_tokens, - cache_write_tokens=usage_obj.input_tokens_details.cache_write_tokens, - ) + if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)): + _usage = _decisions_usage(usage_obj).model_dump() elif usage_obj is None: _usage = {} elif isinstance(usage_obj, BaseModel): diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index e216f6ec252..a9f3c67bf7c 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -338,7 +338,7 @@ class DecisionsIRPredicateQuestion: name: str | None instructions: DecisionsJSON | None criteria: NoulCriteria | None = None - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -352,7 +352,7 @@ class DecisionsIRChoiceQuestion: name: str | None instructions: DecisionsJSON | None choices: tuple[DecisionsIRChoiceOption, ...] - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -366,7 +366,7 @@ class DecisionsIRScoreQuestion: name: str | None instructions: DecisionsJSON | None levels: tuple[DecisionsIRScoreLevel, ...] - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) DecisionsIRQuestion: TypeAlias = DecisionsIRPredicateQuestion | DecisionsIRChoiceQuestion | DecisionsIRScoreQuestion @@ -382,7 +382,7 @@ class DecisionsIRRequest: @dataclass(frozen=True, slots=True) class DecisionsIRPredicateAnswer: probability: float - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -396,7 +396,7 @@ class DecisionsIRChoiceAnswer: choice: str | bool confidence: float probabilities: tuple[DecisionsIRChoiceProbability, ...] - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -411,7 +411,7 @@ class DecisionsIRScoreAnswer: score: float confidence: float probabilities: tuple[DecisionsIRScoreProbability, ...] - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -431,7 +431,7 @@ class DecisionsIRUsage: cached_tokens: int = 0 cache_write_tokens: int = 0 reasoning_tokens: int = 0 - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) @@ -439,7 +439,7 @@ class DecisionsIRResponse: model: str | None answers: tuple[DecisionsIRAnswer, ...] usage: DecisionsIRUsage - extra: Mapping[str, object] = field(default=_NO_EXTRA) + extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) @dataclass(frozen=True, slots=True) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index f08d2f49248..d915b1b4c01 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -285,6 +285,46 @@ def test_decisions_cost_uses_litellm_token_pricing() -> None: assert cost == pytest.approx(expected_cost) +@pytest.mark.parametrize( + "response", + ( + DecisionsResponse( + model="jev-latest", + answers={}, + usage=DecisionsUsage( + input_tokens=_INPUT_TOKENS, + output_tokens=_OUTPUT_TOKENS, + cached_tokens=_CACHED_TOKENS, + cache_write_tokens=_CACHE_WRITE_TOKENS, + ), + ), + OpenAIDecisionResponse.model_validate(_OPENAI_RESPONSE), + ), + ids=("systemone", "openai"), +) +def test_custom_token_pricing_bills_cached_decisions_input_tokens_once( + response: DecisionsResponse | OpenAIDecisionResponse, +) -> None: + cost: Final = litellm.completion_cost( + completion_response=response, + model="gpt-6-luna", + custom_llm_provider="openai", + custom_cost_per_token={ + "input_cost_per_token": 1.0, + "output_cost_per_token": 2.0, + "cache_read_input_token_cost": 0.1, + "cache_creation_input_token_cost": 1.25, + }, + ) + + assert cost == pytest.approx( + (_INPUT_TOKENS - _CACHED_TOKENS - _CACHE_WRITE_TOKENS) * 1.0 + + _CACHED_TOKENS * 0.1 + + _CACHE_WRITE_TOKENS * 1.25 + + _OUTPUT_TOKENS * 2.0 + ) + + def test_decisions_response_hidden_params_getter_preserves_mutable_identity() -> None: response: Final = DecisionsResponse(model="decider", answers={}, usage=None)