diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index bbd22968a9e..6656c3d7a3d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2305,8 +2305,8 @@ class CustomStreamWrapper: already delivered. Recover that partial usage from the chunks seen so far and stash it, with its cost, on the logging object so the failure handler records the real partial spend instead of zero. The router - discards the stash right before it opens the next attempt, so a request - that later recovers via a fallback logs that hop's own usage. + moves it into terminal failure usage before opening the next attempt, + so a recovered attempt logs its own usage. """ if self.logging_obj is None or not self.chunks: return diff --git a/litellm/router.py b/litellm/router.py index 529b3d55b70..6fbca759565 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -233,11 +233,11 @@ from litellm.router_utils.fallback_event_handlers import ( attempted_retries_for_request, carry_over_pre_routing_selection, carry_over_routed_deployment, + carry_superseded_attempt_usage, check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, continuation_fallback_targets, - discard_superseded_attempt_usage, fallback_lookup_groups, fallback_target_name, fallback_target_overrides_request_params, @@ -6163,7 +6163,6 @@ class Router: for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate - await close_source() async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, @@ -6407,7 +6406,7 @@ class Router: verbose_router_logger.debug( "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget ) - discard_superseded_attempt_usage(retry_kwargs) + carry_superseded_attempt_usage(retry_kwargs) try: response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e8248e9e429..31a664a38bc 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -441,13 +441,11 @@ def mid_stream_retry_kwargs( return {key: value for key, value in hop_kwargs.items() if key not in _MID_STREAM_RETRY_STRIPPED_KEYS} -def discard_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: - """The next attempt of this request is about to open. A stream the previous attempt dropped stashed the - usage it consumed on the request's logging object for its own failure log and for the failure row of a - request that runs out of attempts; the next attempt's log must carry its own usage instead.""" +def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: + """Carry the superseded stream's partial usage into billing if the request ultimately fails.""" logging_obj: Final = kwargs.get("litellm_logging_obj") if isinstance(logging_obj, Logging): - logging_obj.discard_partial_usage_for_failure() + logging_obj.carry_failed_attempt_usage() def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None: @@ -881,7 +879,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors - discard_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) + carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index 94df5984cdc..e12b68374be 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -397,6 +397,8 @@ def test_retry_rejected_before_it_opens_counts_against_the_same_budget(gateway: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) _assert_rejected_before_content(response, wire, status=500, attempts=2, error="error") + rows: Final = eventually(lambda: _request_rows(model), lambda values: len(values) >= 1, seconds=70) + assert rows == [{"status": "failure", "prompt_tokens": 5, "completion_tokens": 1}], rows def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_budget(gateway: Gateway) -> None: @@ -508,11 +510,18 @@ def _request_rows(model: str) -> list[dict[str, JsonValue]]: ) -def test_dropped_attempt_and_its_retry_each_log_their_own_usage(gateway: Gateway, tmp_path: Path) -> None: +@pytest.mark.parametrize( + "failure", + (_Failure("drop_after_message_start"), _Failure("error_frame_after_message_start", 500)), + ids=("transport-drop", "recoverable-error-frame"), +) +def test_superseded_attempt_and_its_retry_each_log_their_own_usage( + gateway: Gateway, tmp_path: Path, failure: _Failure +) -> None: prompt: Final = _prompt() attempts: Final = Attempts() with ( - wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts)) as wire, + wire_server(_upstream(prompt, failure, attempts)) as wire, wire_server(lambda request: Reply()) as sink, _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, candidate.scenario() as scenario, diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 72477279c5e..29aecc4b26d 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -199,7 +199,8 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ logging_obj.model_call_details["custom_llm_provider"] = "openai" logging_obj.model_call_details["litellm_params"] = {"acompletion": True} logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] - logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) await run_async_fallback( litellm_router=HopLoggingItsOwnSuccessRouter(), @@ -213,6 +214,27 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ assert [tokens for tokens, _ in recorder.logged] == [7] assert recorder.logged[0][1] != 0.01 + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + +@pytest.mark.asyncio +async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_before_opening() -> None: + logging_obj: Final = _request_logging_obj() + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) class _RateLimitedHop(Exception): diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 069b5572a50..93894efbcb5 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -16096,6 +16096,7 @@ class _AnthropicMessagesFakeByteStream: self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} self.closed = False self.closed_as_superseded = False + self.close_contexts: list[bool] = [] def __aiter__(self): return self @@ -16106,6 +16107,7 @@ class _AnthropicMessagesFakeByteStream: return self._chunks.pop(0) async def aclose(self) -> None: + self.close_contexts.append(closing_superseded_attempt()) self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() self.closed = True @@ -17350,10 +17352,8 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi @pytest.mark.asyncio -async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage(): - """The dropped attempt stashes the usage it consumed on the request's logging object, for its own failure - log and for the failure row of a request that runs out of retries. The router discards it before the retry - opens, so a retry that completes bills its own usage instead of the dropped attempt's.""" +async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: + """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" router = _anthropic_messages_retry_router(num_retries=1) logging_obj = LiteLLMLogging( model="glm", @@ -17364,7 +17364,8 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa litellm_call_id="retry-own-usage", function_id="retry-own-usage", ) - logging_obj.record_partial_usage_for_failure(Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53), 0.01) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) stash_when_retry_opened: list[object] = [] def retry_stream(): @@ -17379,6 +17380,7 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert stash_when_retry_opened == [None] assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) @pytest.mark.asyncio @@ -17403,6 +17405,7 @@ async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_o assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert superseded_closed_when_retry_opened == [True] + assert superseded.close_contexts[0] is True @pytest.mark.asyncio