From 0df4ec60bf9cbf96a3c6034d2f2ee5ce925eb2f6 Mon Sep 17 00:00:00 2001 From: James Hargreaves Date: Fri, 19 Jun 2026 09:21:05 +0100 Subject: [PATCH] address CI: black format, fix empty-dict test, read prompt_tokens_details off Usage, drop redundant or-0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - black-format posthog.py + litellm_logging.py (lint gate) - update test_get_usage_as_dict for the new cache keys in the _empty sentinel (core-utils) - _inject_cache_tokens reads prompt_tokens_details off the Usage object, not the passed dict — correct regardless of caller; tolerates wrapper or dict - add direct _inject_cache_tokens unit test (satisfies logging-coverage gate) - drop redundant 'or 0' in posthog.py per review --- litellm/integrations/posthog.py | 8 +++-- litellm/litellm_core_utils/litellm_logging.py | 26 +++++++++++---- .../test_standard_logging_payload.py | 32 ++++++++++++++++++- .../test_litellm_logging.py | 24 ++++++++++++-- 4 files changed, 78 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index 1651e3b5ba7..ba1c746bef7 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -229,8 +229,12 @@ class PostHogLogger(CustomBatchLogger): properties["$ai_output_tokens"] = self._safe_get( standard_logging_object, "completion_tokens", 0 ) - cache_read = self._safe_get(standard_logging_object, "cache_read_input_tokens", 0) or 0 - cache_creation = self._safe_get(standard_logging_object, "cache_creation_input_tokens", 0) or 0 + cache_read = self._safe_get( + standard_logging_object, "cache_read_input_tokens", 0 + ) + cache_creation = self._safe_get( + standard_logging_object, "cache_creation_input_tokens", 0 + ) if cache_read: properties["$ai_cache_read_input_tokens"] = cache_read if cache_creation: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f2dbf6974b2..66b98489126 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5197,8 +5197,10 @@ class StandardLoggingPayloadSetup: ) if isinstance(_raw, dict): if ResponseAPILoggingUtils._is_response_api_usage(_raw): - usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - _raw + usage = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + _raw + ) ) return StandardLoggingPayloadSetup._inject_cache_tokens( usage.model_dump(), usage @@ -5223,13 +5225,23 @@ class StandardLoggingPayloadSetup: 2. prompt_tokens_details.cached_tokens — set by OpenAI, Gemini, DeepSeek cache_creation_input_tokens: Anthropic and Bedrock only via private attr. + + Reads prompt_tokens_details off the Usage object (not ``d``) so the + result is correct regardless of what the caller passes as ``d``. + prompt_tokens_details may be a Pydantic wrapper (attribute access) or a + dict, so tolerate both. """ - ptd = d.get("prompt_tokens_details") or {} - ptd_cached = (ptd.get("cached_tokens") or 0) if isinstance(ptd, dict) else 0 + ptd = getattr(usage, "prompt_tokens_details", None) + if isinstance(ptd, dict): + ptd_cached = ptd.get("cached_tokens") or 0 + else: + ptd_cached = getattr(ptd, "cached_tokens", 0) or 0 d["cache_read_input_tokens"] = ( getattr(usage, "_cache_read_input_tokens", 0) or ptd_cached ) - d["cache_creation_input_tokens"] = getattr(usage, "_cache_creation_input_tokens", 0) or 0 + d["cache_creation_input_tokens"] = ( + getattr(usage, "_cache_creation_input_tokens", 0) or 0 + ) return d @staticmethod @@ -5927,7 +5939,9 @@ def get_standard_logging_object_payload( prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), cache_read_input_tokens=usage_dict.get("cache_read_input_tokens", 0), - cache_creation_input_tokens=usage_dict.get("cache_creation_input_tokens", 0), + cache_creation_input_tokens=usage_dict.get( + "cache_creation_input_tokens", 0 + ), request_tags=request_tags, end_user=end_user_id or "", api_base=StandardLoggingPayloadSetup.strip_trailing_slash( diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index 3a75994d233..34098d63a90 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -1105,8 +1105,12 @@ def test_merge_litellm_metadata_bedrock_passthrough_scenario(): # --- cache token tests --- -def _make_usage(cache_read: int = 0, cache_creation: int = 0, ptd_cached: int = 0) -> Usage: + +def _make_usage( + cache_read: int = 0, cache_creation: int = 0, ptd_cached: int = 0 +) -> Usage: from litellm.types.utils import PromptTokensDetailsWrapper + u = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) u._cache_read_input_tokens = cache_read u._cache_creation_input_tokens = cache_creation @@ -1154,3 +1158,29 @@ def test_get_usage_as_dict_from_response_obj(): d = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj={"usage": usage}) assert d["cache_read_input_tokens"] == 300 assert d["cache_creation_input_tokens"] == 100 + + +def test_inject_cache_tokens_prefers_private_attr_then_prompt_tokens_details(): + """Direct coverage for the _inject_cache_tokens helper: the private + Anthropic/Bedrock attrs win, with a prompt_tokens_details.cached_tokens + fallback for OpenAI/Gemini/DeepSeek.""" + from litellm.types.utils import PromptTokensDetailsWrapper + + # Private attrs present → used directly. + anthropic = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + anthropic._cache_read_input_tokens = 200 + anthropic._cache_creation_input_tokens = 500 + d = StandardLoggingPayloadSetup._inject_cache_tokens({}, anthropic) + assert d["cache_read_input_tokens"] == 200 + assert d["cache_creation_input_tokens"] == 500 + + # No private attr → fall back to prompt_tokens_details.cached_tokens. + openai = Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=80), + ) + d = StandardLoggingPayloadSetup._inject_cache_tokens({}, openai) + assert d["cache_read_input_tokens"] == 80 + assert d["cache_creation_input_tokens"] == 0 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f0db0409bd7..f9af463e7be 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1726,11 +1726,23 @@ def test_get_usage_as_dict(): # Test case 1: None response_obj returns empty usage dict result = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj=None) - assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + assert result == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } # Test case 2: Empty response_obj returns empty usage dict result = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj={}) - assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + assert result == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } # Test case 3: combined_usage_object takes priority combined = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) @@ -1752,7 +1764,13 @@ def test_get_usage_as_dict(): result = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj={"id": "resp-1", "choices": []} ) - assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + assert result == { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } def test_append_system_prompt_messages():