fix(decisions): load on Python 3.10 and 3.11 and bill cached tokens once under custom pricing

Decisions IR dataclasses used a mappingproxy default, which is only hashable from Python 3.12, so import litellm failed on 3.10 and 3.11. Decisions usage now reaches cost_per_token as a Usage with prompt_tokens_details, so custom pricing no longer adds cached and cache-write tokens on top of input_tokens that already include them
This commit is contained in:
mateo-berri 2026-10-07 19:26:05 -07:00
parent 85f9121192
commit 39012f1c22
3 changed files with 75 additions and 32 deletions

View file

@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast
from httpx import Response from httpx import Response
from pydantic import BaseModel from pydantic import BaseModel
from typing_extensions import ReadOnly, TypedDict from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm import litellm
import litellm._logging 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: if usage_obj is None:
return None return None
if isinstance(usage_obj, Usage): if isinstance(usage_obj, Usage):
@ -1064,15 +1066,28 @@ def get_usage_object(
return None return None
def _decisions_usage( def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper:
*, input_tokens: int, output_tokens: int, cached_tokens: int, cache_write_tokens: int match usage:
) -> dict[str, object]: case DecisionsUsage():
return { return PromptTokensDetailsWrapper(
"prompt_tokens": input_tokens, cached_tokens=usage.cached_tokens, cache_write_tokens=usage.cache_write_tokens
"completion_tokens": output_tokens, )
"cache_read_input_tokens": cached_tokens, case OpenAIDecisionUsage():
"cache_creation_input_tokens": cache_write_tokens, 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): def _is_known_usage_objects(usage_obj):
@ -1493,20 +1508,8 @@ def completion_cost(
"usage", "usage",
litellm.Usage(**_usage_for_dump.model_dump()), litellm.Usage(**_usage_for_dump.model_dump()),
) )
if isinstance(usage_obj, DecisionsUsage): if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
_usage = _decisions_usage( _usage = _decisions_usage(usage_obj).model_dump()
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,
)
elif usage_obj is None: elif usage_obj is None:
_usage = {} _usage = {}
elif isinstance(usage_obj, BaseModel): elif isinstance(usage_obj, BaseModel):

View file

@ -338,7 +338,7 @@ class DecisionsIRPredicateQuestion:
name: str | None name: str | None
instructions: DecisionsJSON | None instructions: DecisionsJSON | None
criteria: NoulCriteria | None = 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) @dataclass(frozen=True, slots=True)
@ -352,7 +352,7 @@ class DecisionsIRChoiceQuestion:
name: str | None name: str | None
instructions: DecisionsJSON | None instructions: DecisionsJSON | None
choices: tuple[DecisionsIRChoiceOption, ...] 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) @dataclass(frozen=True, slots=True)
@ -366,7 +366,7 @@ class DecisionsIRScoreQuestion:
name: str | None name: str | None
instructions: DecisionsJSON | None instructions: DecisionsJSON | None
levels: tuple[DecisionsIRScoreLevel, ...] 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 DecisionsIRQuestion: TypeAlias = DecisionsIRPredicateQuestion | DecisionsIRChoiceQuestion | DecisionsIRScoreQuestion
@ -382,7 +382,7 @@ class DecisionsIRRequest:
@dataclass(frozen=True, slots=True) @dataclass(frozen=True, slots=True)
class DecisionsIRPredicateAnswer: class DecisionsIRPredicateAnswer:
probability: float 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) @dataclass(frozen=True, slots=True)
@ -396,7 +396,7 @@ class DecisionsIRChoiceAnswer:
choice: str | bool choice: str | bool
confidence: float confidence: float
probabilities: tuple[DecisionsIRChoiceProbability, ...] 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) @dataclass(frozen=True, slots=True)
@ -411,7 +411,7 @@ class DecisionsIRScoreAnswer:
score: float score: float
confidence: float confidence: float
probabilities: tuple[DecisionsIRScoreProbability, ...] 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) @dataclass(frozen=True, slots=True)
@ -431,7 +431,7 @@ class DecisionsIRUsage:
cached_tokens: int = 0 cached_tokens: int = 0
cache_write_tokens: int = 0 cache_write_tokens: int = 0
reasoning_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) @dataclass(frozen=True, slots=True)
@ -439,7 +439,7 @@ class DecisionsIRResponse:
model: str | None model: str | None
answers: tuple[DecisionsIRAnswer, ...] answers: tuple[DecisionsIRAnswer, ...]
usage: DecisionsIRUsage 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) @dataclass(frozen=True, slots=True)

View file

@ -285,6 +285,46 @@ def test_decisions_cost_uses_litellm_token_pricing() -> None:
assert cost == pytest.approx(expected_cost) 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: def test_decisions_response_hidden_params_getter_preserves_mutable_identity() -> None:
response: Final = DecisionsResponse(model="decider", answers={}, usage=None) response: Final = DecisionsResponse(model="decider", answers={}, usage=None)