diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6fecc7fa97..9956a9d314 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1265,7 +1265,7 @@ class BaseTokenUsageProcessor: Combine multiple Usage objects into a single Usage object, checking model keys for nested values. """ from litellm.types.utils import ( - CompletionTokensDetails, + CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage, ) @@ -1320,7 +1320,7 @@ class BaseTokenUsageProcessor: not hasattr(combined, "completion_tokens_details") or not combined.completion_tokens_details ): - combined.completion_tokens_details = CompletionTokensDetails() + combined.completion_tokens_details = CompletionTokensDetailsWrapper() # Check what keys exist in the model's completion_tokens_details for attr in usage.completion_tokens_details.model_fields: diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fb919afd49..2f85c7aef6 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -527,7 +527,12 @@ class ChunkProcessor: returned_usage, "cache_read_input_tokens", cache_read_input_tokens ) # for anthropic if completion_tokens_details is not None: - returned_usage.completion_tokens_details = completion_tokens_details + if isinstance(completion_tokens_details, CompletionTokensDetails): + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( + **completion_tokens_details.model_dump() + ) + else: + returned_usage.completion_tokens_details = completion_tokens_details if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d1b2f57c72..c5af6e4ed0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -888,6 +888,9 @@ class Usage(CompletionUsage): server_tool_use: Optional[ServerToolUse] = None + completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None + """Breakdown of tokens used in a completion.""" + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None """Breakdown of tokens used in the prompt.""" diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 26df08ecc6..1f68e8a43a 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -2155,8 +2155,8 @@ async def test_caching_kwargs_input(sync_mode): Message, ModelResponse, Usage, - CompletionTokensDetails, - PromptTokensDetails, + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, ) from datetime import datetime @@ -2187,10 +2187,10 @@ async def test_caching_kwargs_input(sync_mode): completion_tokens=31, prompt_tokens=16, total_tokens=47, - completion_tokens_details=CompletionTokensDetails( + completion_tokens_details=CompletionTokensDetailsWrapper( audio_tokens=None, reasoning_tokens=0 ), - prompt_tokens_details=PromptTokensDetails( + prompt_tokens_details=PromptTokensDetailsWrapper( audio_tokens=None, cached_tokens=0 ), ), diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 1bf005db50..71b9247518 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -73,3 +73,56 @@ def test_usage_dump(): new_usage = Usage(**current_usage.model_dump()) assert new_usage.prompt_tokens_details.web_search_requests == 1 + + +def test_usage_completion_tokens_details_text_tokens(): + from litellm.types.utils import Usage + + # Test data from the reported issue + usage_data = { + 'completion_tokens': 77, + 'prompt_tokens': 11937, + 'total_tokens': 12014, + 'completion_tokens_details': { + 'accepted_prediction_tokens': None, + 'audio_tokens': None, + 'reasoning_tokens': 65, + 'rejected_prediction_tokens': None, + 'text_tokens': 12 + }, + 'prompt_tokens_details': { + 'audio_tokens': None, + 'cached_tokens': None, + 'text_tokens': 11937, + 'image_tokens': None + } + } + + # Create Usage object + u = Usage(**usage_data) + + # Verify the object has the text_tokens field + assert hasattr(u.completion_tokens_details, 'text_tokens') + assert u.completion_tokens_details.text_tokens == 12 + + # Get model_dump output + dump_result = u.model_dump() + + # Verify text_tokens is present in the model_dump output + assert 'completion_tokens_details' in dump_result + assert 'text_tokens' in dump_result['completion_tokens_details'] + assert dump_result['completion_tokens_details']['text_tokens'] == 12 + + # Verify the full completion_tokens_details structure + expected_completion_details = { + 'accepted_prediction_tokens': None, + 'audio_tokens': None, + 'reasoning_tokens': 65, + 'rejected_prediction_tokens': None, + 'text_tokens': 12 + } + assert dump_result['completion_tokens_details'] == expected_completion_details + + # Verify round-trip serialization works + new_usage = Usage(**dump_result) + assert new_usage.completion_tokens_details.text_tokens == 12