diff --git a/litellm/router.py b/litellm/router.py index 6e8127110cc..70ef79ae0c7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1899,11 +1899,7 @@ class Router: usage = cast(Optional[Usage], getattr(fallback_item, "usage", None)) usage_objects = [usage] if usage is not None else [] - if ( - complete_response_object_usage is not None - and hasattr(complete_response_object_usage, "usage") - and complete_response_object_usage.usage is not None # type: ignore - ): + if complete_response_object_usage is not None: usage_objects.append(complete_response_object_usage) combined_usage = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) setattr(fallback_item, "usage", combined_usage) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9c4d83ff7ea..186e98d83dc 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4025,12 +4025,10 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): assert model_info["output_cost_per_token"] == 0.0015 -def test_combine_fallback_usage(): - """Test that _combine_fallback_usage merges partial and fallback usage.""" +def test_combine_fallback_usage_without_partial_usage(): from litellm.router import Router from litellm.types.utils import Usage - # Create a stream chunk with usage chunk = litellm.ModelResponseStream( id="test", model="gpt-4o", @@ -4038,7 +4036,6 @@ def test_combine_fallback_usage(): usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), ) - # Call _combine_fallback_usage with no extra usage Router._combine_fallback_usage(chunk, None) assert chunk.usage is not None assert chunk.usage.prompt_tokens == 10 @@ -4046,6 +4043,26 @@ def test_combine_fallback_usage(): assert chunk.usage.total_tokens == 15 +def test_combine_fallback_usage(): + from litellm.router import Router + from litellm.types.utils import Usage + + chunk = litellm.ModelResponseStream( + id="test", + model="gpt-4o", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + partial_usage = Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10) + + Router._combine_fallback_usage(chunk, partial_usage) + + assert chunk.usage is not None + assert chunk.usage.prompt_tokens == 17 + assert chunk.usage.completion_tokens == 8 + assert chunk.usage.total_tokens == 25 + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): """A mid-stream failure with no successful fallback raises and is logged as