scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Move cache token normalization to the transformation layer

Browse source 
Address review feedback:
- Move provider-specific field-name knowledge out of proxy/db/ and into
  get_usage_as_dict (the transformation layer), so the spend writer
  receives a uniform usage dict regardless of provider.
- Remove dead code: cache_creation_tokens fallback referenced a field
  that no provider's API actually returns.
- Use `not in` key check instead of `if value:` to correctly preserve
  an explicit 0 from the top-level field.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-11 20:43:01 -07:00
parent 3203cd9895
commit 8182d66aae
4 changed files with 145 additions and 107 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4751,27 +4751,51 @@ class StandardLoggingPayloadSetup:
"""
_empty: dict = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
if combined_usage_object is not None:
return combined_usage_object.model_dump()
if not response_obj:
result = combined_usage_object.model_dump()
elif not response_obj:
return _empty
_raw = response_obj.get("usage", None)
if _raw is None:
return _empty
if isinstance(_raw, ResponseAPIUsage):
return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
_raw
).model_dump()
if isinstance(_raw, dict):
if ResponseAPILoggingUtils._is_response_api_usage(_raw):
return (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
_raw
).model_dump()
)
return _raw
if isinstance(_raw, Usage):
return _raw.model_dump()
return _empty
else:
_raw = response_obj.get("usage", None)
if _raw is None:
return _empty
if isinstance(_raw, ResponseAPIUsage):
result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
_raw
).model_dump()
elif isinstance(_raw, dict):
if ResponseAPILoggingUtils._is_response_api_usage(_raw):
result = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
_raw
).model_dump()
)
else:
result = _raw
elif isinstance(_raw, Usage):
result = _raw.model_dump()
else:
return _empty
StandardLoggingPayloadSetup._normalize_cache_tokens(result)
return result
@staticmethod
def _normalize_cache_tokens(usage_dict: dict) -> None:
"""
Ensure cache_read_input_tokens is set on the usage dict.
OpenAI models report cached tokens only under
prompt_tokens_details.cached_tokens. Anthropic, DeepSeek, and
Gemini set cache_read_input_tokens at the top level. This
normalizes the dict so downstream consumers (spend writer, etc.)
can always read cache_read_input_tokens directly.
"""
if "cache_read_input_tokens" not in usage_dict:
prompt_details = usage_dict.get("prompt_tokens_details")
if isinstance(prompt_details, dict):
cached = prompt_details.get("cached_tokens")
if cached is not None:
usage_dict["cache_read_input_tokens"] = cached
@staticmethod
def get_model_cost_information(

44
litellm/proxy/db/db_spend_update_writer.py
View file

@ -1852,38 +1852,6 @@ class DBSpendUpdateWriter:
unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
def _get_cache_read_input_tokens(usage_obj: dict) -> int:
"""
Extract cache read input tokens from a usage object dict.
Anthropic/DeepSeek store this as a top-level `cache_read_input_tokens` field.
OpenAI stores it nested under `prompt_tokens_details.cached_tokens`.
"""
value = usage_obj.get("cache_read_input_tokens", 0) or 0
if value:
return value
prompt_details = usage_obj.get("prompt_tokens_details") or {}
if isinstance(prompt_details, dict):
return prompt_details.get("cached_tokens", 0) or 0
return 0
@staticmethod
def _get_cache_creation_input_tokens(usage_obj: dict) -> int:
"""
Extract cache creation input tokens from a usage object dict.
Anthropic stores this as a top-level `cache_creation_input_tokens` field.
It may also appear under `prompt_tokens_details.cache_creation_tokens`.
"""
value = usage_obj.get("cache_creation_input_tokens", 0) or 0
if value:
return value
prompt_details = usage_obj.get("prompt_tokens_details") or {}
if isinstance(prompt_details, dict):
return prompt_details.get("cache_creation_tokens", 0) or 0
return 0
async def _common_add_spend_log_transaction_to_daily_transaction(
self,
payload: Union[dict, SpendLogsPayload],
@ -1964,12 +1932,12 @@ class DBSpendUpdateWriter:
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
cache_read_input_tokens=self._get_cache_read_input_tokens(
usage_obj
),
cache_creation_input_tokens=self._get_cache_creation_input_tokens(
usage_obj
),
cache_read_input_tokens=usage_obj.get("cache_read_input_tokens", 0)
or 0,
cache_creation_input_tokens=usage_obj.get(
"cache_creation_input_tokens", 0
)
or 0,
)
return daily_transaction
except Exception as e:

97
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -13,8 +13,11 @@ import time
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.litellm_core_utils.litellm_logging import set_callbacks
from litellm.types.utils import ModelResponse, TextCompletionResponse
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
set_callbacks,
)
from litellm.types.utils import ModelResponse, TextCompletionResponse, Usage
@pytest.fixture
@ -1160,6 +1163,96 @@ def test_get_usage_as_dict():
assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
class TestNormalizeCacheTokens:
"""Tests for _normalize_cache_tokens called via get_usage_as_dict.
OpenAI models store cached tokens only under prompt_tokens_details.cached_tokens.
Anthropic/DeepSeek/Gemini set cache_read_input_tokens at the top level.
The normalization ensures cache_read_input_tokens is always present when
cached tokens are reported by any provider.
"""
def test_openai_cached_tokens_promoted_to_top_level(self):
"""OpenAI-style usage gets cache_read_input_tokens set from prompt_tokens_details."""
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj={
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"prompt_tokens_details": {"cached_tokens": 300},
}
}
)
assert result["cache_read_input_tokens"] == 300
def test_anthropic_top_level_field_preserved(self):
"""Anthropic-style usage already has cache_read_input_tokens — no overwrite."""
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj={
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"cache_read_input_tokens": 500,
"prompt_tokens_details": {"cached_tokens": 500},
}
}
)
assert result["cache_read_input_tokens"] == 500
def test_top_level_zero_not_overwritten(self):
"""An explicit 0 at top level must not be replaced by a nested value."""
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj={
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"cache_read_input_tokens": 0,
"prompt_tokens_details": {"cached_tokens": 300},
}
}
)
assert result["cache_read_input_tokens"] == 0
def test_no_cached_tokens_anywhere(self):
"""No cache fields at all — key should not be added."""
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj={
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
}
}
)
assert "cache_read_input_tokens" not in result
def test_combined_usage_object_normalized(self):
"""Usage passed via combined_usage_object is also normalized."""
usage = Usage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150,
prompt_tokens_details={"cached_tokens": 42},
)
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj=None,
combined_usage_object=usage,
)
assert result["cache_read_input_tokens"] == 42
def test_none_cached_tokens_not_promoted(self):
"""cached_tokens: None should not be promoted."""
result = StandardLoggingPayloadSetup.get_usage_as_dict(
response_obj={
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"prompt_tokens_details": {"cached_tokens": None},
}
}
)
assert "cache_read_input_tokens" not in result
def test_append_system_prompt_messages():
"""
Test append_system_prompt_messages prepends system message from kwargs to messages list.

47
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -1430,50 +1430,3 @@ async def test_commit_spend_updates_uses_pipeline():
mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called()
class TestGetCacheTokenHelpers:
"""Tests for _get_cache_read_input_tokens and _get_cache_creation_input_tokens.
These helpers extract cache token counts from usage objects. Anthropic/DeepSeek
store them as top-level fields, while OpenAI nests them under prompt_tokens_details.
"""
def test_cache_read_from_top_level_field(self):
"""Anthropic-style: cache_read_input_tokens at top level."""
usage = {"cache_read_input_tokens": 500}
assert DBSpendUpdateWriter._get_cache_read_input_tokens(usage) == 500
def test_cache_read_from_prompt_tokens_details(self):
"""OpenAI-style: cached_tokens nested in prompt_tokens_details."""
usage = {"prompt_tokens_details": {"cached_tokens": 300}}
assert DBSpendUpdateWriter._get_cache_read_input_tokens(usage) == 300
def test_cache_read_top_level_takes_precedence(self):
"""When both are present, top-level field wins."""
usage = {
"cache_read_input_tokens": 500,
"prompt_tokens_details": {"cached_tokens": 300},
}
assert DBSpendUpdateWriter._get_cache_read_input_tokens(usage) == 500
def test_cache_read_empty_usage(self):
assert DBSpendUpdateWriter._get_cache_read_input_tokens({}) == 0
def test_cache_read_none_values(self):
usage = {
"cache_read_input_tokens": None,
"prompt_tokens_details": None,
}
assert DBSpendUpdateWriter._get_cache_read_input_tokens(usage) == 0
def test_cache_creation_from_top_level_field(self):
"""Anthropic-style: cache_creation_input_tokens at top level."""
usage = {"cache_creation_input_tokens": 200}
assert DBSpendUpdateWriter._get_cache_creation_input_tokens(usage) == 200
def test_cache_creation_from_prompt_tokens_details(self):
"""cache_creation_tokens nested in prompt_tokens_details."""
usage = {"prompt_tokens_details": {"cache_creation_tokens": 150}}
assert DBSpendUpdateWriter._get_cache_creation_input_tokens(usage) == 150
def test_cache_creation_empty_usage(self):
assert DBSpendUpdateWriter._get_cache_creation_input_tokens({}) == 0