diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index a92b5d25a37..7667d1bad84 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -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 diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 96ac2e2c345..09628cd4a76 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -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