diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 00c55b35182..d1423f04dd0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1758,6 +1758,12 @@ class PromptTokensDetailsWrapper( del self.cache_creation_token_details +def _non_negative_token_count(value: object) -> int | None: + if type(value) is int and value >= 0: + return value + return None + + class ServerToolUse(BaseModel): web_search_requests: int | None = None tool_search_requests: int | None = None @@ -1844,12 +1850,22 @@ class Usage(SafeAttributeModel, CompletionUsage): elif isinstance(prompt_tokens_details, PromptTokensDetailsWrapper): _prompt_tokens_details = prompt_tokens_details + deepseek_cache_hit_tokens: Final = _non_negative_token_count(params.get("prompt_cache_hit_tokens")) + deepseek_cache_miss_tokens: Final = _non_negative_token_count(params.get("prompt_cache_miss_tokens")) + ## DEEPSEEK MAPPING ## - if "prompt_cache_hit_tokens" in params and isinstance(params["prompt_cache_hit_tokens"], int): + if deepseek_cache_hit_tokens is not None: if _prompt_tokens_details is None: - _prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=params["prompt_cache_hit_tokens"]) + _prompt_tokens_details = PromptTokensDetailsWrapper( + cached_tokens=deepseek_cache_hit_tokens + ) # rebind-ok: builds prompt details from provider cache metadata else: - _prompt_tokens_details.cached_tokens = params["prompt_cache_hit_tokens"] + _prompt_tokens_details.cached_tokens = deepseek_cache_hit_tokens + + deepseek_prompt_tokens: Final = (deepseek_cache_hit_tokens or 0) + (deepseek_cache_miss_tokens or 0) + effective_prompt_tokens: Final = ( + deepseek_prompt_tokens if not prompt_tokens and deepseek_prompt_tokens else prompt_tokens or 0 + ) ## ANTHROPIC MAPPING ## if "cache_read_input_tokens" in params and isinstance(params["cache_read_input_tokens"], int): @@ -1867,7 +1883,7 @@ class Usage(SafeAttributeModel, CompletionUsage): _prompt_tokens_details.cache_write_tokens = params["cache_creation_input_tokens"] super().__init__( - prompt_tokens=prompt_tokens or 0, + prompt_tokens=effective_prompt_tokens, completion_tokens=completion_tokens or 0, total_tokens=total_tokens or 0, completion_tokens_details=_completion_tokens_details or None, @@ -1895,8 +1911,8 @@ class Usage(SafeAttributeModel, CompletionUsage): self._cache_read_input_tokens = params["cache_read_input_tokens"] ## DEEPSEEK MAPPING ## - if "prompt_cache_hit_tokens" in params and isinstance(params["prompt_cache_hit_tokens"], int): - self._cache_read_input_tokens = params["prompt_cache_hit_tokens"] + if deepseek_cache_hit_tokens is not None: + self._cache_read_input_tokens = deepseek_cache_hit_tokens for k, v in params.items(): setattr(self, k, v) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 5f44ba1773e..8bd321edce6 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -138,6 +138,69 @@ def test_usage_converts_server_tool_use_dict(): assert round_trip.server_tool_use.tool_search_requests == 1 +def test_usage_maps_deepseek_cache_hit_and_miss_tokens(): + from litellm.types.utils import Usage + + usage = Usage( + completion_tokens=50, + prompt_cache_hit_tokens=75, + prompt_cache_miss_tokens=25, + ) + + assert usage.prompt_tokens == 100 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 75 + assert usage._cache_read_input_tokens == 75 + + +def test_usage_does_not_override_reported_prompt_tokens(): + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + prompt_cache_hit_tokens=75, + prompt_cache_miss_tokens=25, + ) + + assert usage.prompt_tokens == 120 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 75 + + +def test_usage_ignores_malformed_deepseek_cache_bucket_tokens(): + from litellm.types.utils import Usage + + usage = Usage( + completion_tokens=50, + prompt_cache_hit_tokens=True, + prompt_cache_miss_tokens="25", + ) + + assert usage.prompt_tokens == 0 + assert usage.prompt_tokens_details is None + assert usage._cache_read_input_tokens == 0 + + +def test_usage_preserves_cache_read_and_write_mappings(): + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + cache_read_input_tokens=30, + cache_creation_input_tokens=10, + ) + + assert usage.prompt_tokens == 100 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 30 + assert usage.prompt_tokens_details.cache_write_tokens == 10 + assert usage.prompt_tokens_details.cache_creation_tokens == 10 + assert usage._cache_read_input_tokens == 30 + assert usage._cache_creation_input_tokens == 10 + + def test_usage_completion_tokens_details_text_tokens(): from litellm.types.utils import Usage