diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ba5bf11ebd6..2d640f79d4a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2273,8 +2273,14 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: total_tokens=prompt_tokens + completion_tokens, ) - if latest_usage_chunk and hasattr(latest_usage_chunk, "cost") and latest_usage_chunk.cost is not None: - returned_usage_chunk.cost = latest_usage_chunk.cost + if latest_usage_chunk is not None: + latest_cost = ( + latest_usage_chunk.get("cost") + if isinstance(latest_usage_chunk, dict) + else getattr(latest_usage_chunk, "cost", None) + ) + if latest_cost is not None: + returned_usage_chunk.cost = latest_cost return returned_usage_chunk diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 6108d853286..41c72c9593e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1428,6 +1428,27 @@ def test_calculate_total_usage_with_cost(): assert usage.completion_tokens == 5 +def test_calculate_total_usage_with_dict_usage_cost(): + """Regression: dict-shaped `usage` with a `cost` key must still surface + provider cost even though `hasattr` on a dict does not consult its keys.""" + from litellm.litellm_core_utils.streaming_handler import calculate_total_usage + + chunk = { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "cost": 0.00025, + } + } + + usage = calculate_total_usage([chunk]) + + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 + assert getattr(usage, "cost", None) == 0.00025 + + @pytest.mark.asyncio async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging): from litellm.utils import ModelResponseListIterator