address CI: black format, fix empty-dict test, read prompt_tokens_details off Usage, drop redundant or-0

- 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
This commit is contained in:
James Hargreaves 2026-06-19 09:21:05 +01:00
parent 57f6de7953
commit 0df4ec60bf
4 changed files with 78 additions and 12 deletions

View file

@ -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:

View file

@ -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(

View file

@ -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

View file

@ -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():