diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9b0f58e699c..f46efb74982 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1819,7 +1819,7 @@ def _failure_fields_to_lift( ) _billed: Final = provider_billed and not (_success_already_logged or _stream_success_bills_this_attempt) _usage_to_lift: Final = ( - None + (Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), 0.0) if _stream_success_bills_this_attempt else _failure_usage_to_lift( model_call_details=_model_call_details, diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 5242950d783..a0671da01c6 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -28,6 +28,7 @@ from litellm.llms.anthropic.messages.native.streaming_iterator import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import ProxyLogging, _failure_fields_to_lift from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Usage @@ -782,9 +783,11 @@ async def test_chat_stream_callback_error_after_its_success_dispatch_leaves_the_ pass await asyncio.sleep(0) - lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + failure_row: Final = get_logging_payload( + dict(_failure_fields_to_lift(request_data, provider_billed=True)), None, datetime.now(), datetime.now() + ) assert events == ["success_dispatched"] - assert (lifted.get("combined_usage_object"), lifted.get("response_cost")) == (None, None) + assert (failure_row["prompt_tokens"], failure_row["completion_tokens"], failure_row["spend"]) == (0, 0, 0.0) @pytest.mark.asyncio diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 2522301031b..14607a25289 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2868,6 +2868,7 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t async def test_a_failure_after_a_continued_stream_success_bills_nothing_the_success_row_already_billed( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter ) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import _failure_fields_to_lift monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -2880,7 +2881,12 @@ async def test_a_failure_after_a_continued_stream_success_bills_nothing_the_succ (success,) = recorder.successes assert success["cost_breakdown"]["additional_costs"]["mid_stream_partial_cost"] > 0 late_failure: Final = _failure_fields_to_lift({"litellm_logging_obj": shared_logging}, provider_billed=True) - assert (late_failure.get("combined_usage_object"), late_failure.get("response_cost")) == (None, None) + late_failure_row: Final = get_logging_payload(dict(late_failure), None, datetime.now(), datetime.now()) + assert (late_failure_row["prompt_tokens"], late_failure_row["completion_tokens"], late_failure_row["spend"]) == ( + 0, + 0, + 0.0, + ) _AGENTIC_FIRST_CALL_USAGE: Final = {"prompt_tokens": 50, "completion_tokens": 7, "total_tokens": 57} @@ -3286,6 +3292,7 @@ def _backup_responses_cost() -> float: async def test_responses_fallback_sharing_the_proxy_logging_object_bills_the_cut_partial_once( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, failed_usage: Mapping[str, int] | None ) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import _failure_fields_to_lift monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -3318,7 +3325,12 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_bills_the_cut success["response_cost"], ) late_failure: Final = _failure_fields_to_lift({"litellm_logging_obj": shared_logging}, provider_billed=True) - assert (late_failure.get("combined_usage_object"), late_failure.get("response_cost")) == (None, None) + late_failure_row: Final = get_logging_payload(dict(late_failure), None, datetime.now(), datetime.now()) + assert (late_failure_row["prompt_tokens"], late_failure_row["completion_tokens"], late_failure_row["spend"]) == ( + 0, + 0, + 0.0, + ) @pytest.mark.asyncio