mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
57f6de7953
commit
0df4ec60bf
4 changed files with 78 additions and 12 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue