scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

Tally total_tokens in response if missing (#18468) (#18445)

Browse source 
Calculate `total_tokens` in usage data in Response manually if:

    - `total_tokens` is missing
    - `total_tokens` can be calculated from input and output tokens

Run the test for this feature with:

`poetry run pytest tests/test_litellm/responses/test_responses_utils.py -k "test_transform_response_api_usage_calculates_total_from_input_and_output_tokens_if_available" -v`
This commit is contained in:
Wiley Kestner 2026-01-06 16:04:02 +01:00 • committed by GitHub
parent 8d40335684
commit 9e6714fe1b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 31 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/responses/utils.py
View file

@ -443,11 +443,18 @@ class ResponseAPILoggingUtils:
completion_tokens=0,
total_tokens=0,
)
response_api_usage: ResponseAPIUsage = (
ResponseAPIUsage(**usage_input)
if isinstance(usage_input, dict)
else usage_input
)
response_api_usage: ResponseAPIUsage
if isinstance(usage_input, dict):
total_tokens = usage_input.get("total_tokens")
if total_tokens is None:
input_tokens = usage_input.get("input_tokens")
output_tokens = usage_input.get("output_tokens")
if input_tokens is not None and output_tokens is not None:
total_tokens = input_tokens + output_tokens
usage_input["total_tokens"] = total_tokens
response_api_usage = ResponseAPIUsage(**usage_input)
else:
response_api_usage = usage_input
prompt_tokens: int = response_api_usage.input_tokens or 0
completion_tokens: int = response_api_usage.output_tokens or 0
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None

19
tests/test_litellm/responses/test_responses_utils.py
View file

@ -203,3 +203,22 @@ class TestResponseAPILoggingUtils:
assert result.prompt_tokens == 0
assert result.completion_tokens == 20
assert result.total_tokens == 20
def test_transform_response_api_usage_calculates_total_from_input_and_output_tokens_if_available(self):
"""Test transformation calculates total_tokens when it's None and input / output tokens are present"""
# Setup
usage = {
"input_tokens": 15,
"output_tokens": 25,
"total_tokens": None,
}
# Execute
result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
usage
)
# Assert
assert result.prompt_tokens == 15
assert result.completion_tokens == 25
assert result.total_tokens == 40 # 15 + 25