diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9b6ec8d28d6..817d45db518 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass): self._async_success_scheduled: bool = False self._async_success_enqueued: bool = False self._sync_success_scheduled: bool = False + self._superseded_success_logs_in_flight: frozenset[str] = frozenset() self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self._stream_failure_claimed: bool = False self._stash_owed_on_success: bool = False @@ -2904,6 +2905,9 @@ class Logging(LiteLLMLoggingBaseClass): Ignore check on stream, as there can be multiple chunks """ return + if event_type in self._superseded_success_logs_in_flight: + self._superseded_success_logs_in_flight = self._superseded_success_logs_in_flight - frozenset((event_type,)) + return self.model_call_details[f"has_logged_{event_type}"] = True if event_type == "async_failure": self.model_call_details["async_failure_logged_attempt"] = self.model_call_details.get( @@ -3511,6 +3515,26 @@ class Logging(LiteLLMLoggingBaseClass): self._reset_success_emission_dedupe() await self.async_success_handler(result=result) + def reopen_success_logging_for_next_attempt(self) -> None: + """The router discarded an earlier attempt's response, so the next attempt logs and prices its own + response. A success log the discarded attempt parked for the proxy to dispatch later is dropped, since + it describes a response the caller never receives. One it already dispatched still runs, and its late + completion must not mark the next attempt as logged.""" + self._superseded_success_logs_in_flight = frozenset( + event_type + for event_type, scheduled in ( + ("async_success", self._async_success_scheduled), + ("sync_success", self._sync_success_scheduled), + ) + if scheduled + and self.deferred_success_result is None + and self.model_call_details.get(f"has_logged_{event_type}") is not True + ) + self._reset_success_emission_dedupe() + self._async_success_scheduled = False + self.deferred_success_result = None + self.enqueue_deferred_logging = None + def _reset_success_emission_dedupe(self) -> None: """ Success callbacks dedupe per request, because the sync and async diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 3e222313f32..52ad47411a9 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -4584,10 +4584,15 @@ class _Driver: self._fallback_depth = depth if isinstance(depth, int) else self._fallback_depth + 1 def _carry_superseded_usage(self) -> None: - """``carry_superseded_attempt_usage``, before each same-group retry and each fallback hop.""" - from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage + """``carry_superseded_attempt_usage`` and ``reopen_success_logging_for_next_attempt``, before each + same-group retry and each fallback hop.""" + from litellm.router_utils.fallback_event_handlers import ( + carry_superseded_attempt_usage, + reopen_success_logging_for_next_attempt, + ) carry_superseded_attempt_usage(self._request) + reopen_success_logging_for_next_attempt(self._request) def _order_chain_failed(self, event: Mapping[str, object]) -> None: """The order path of ``async_function_with_fallbacks_common_utils`` raised: the chain's last diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 20f9e114747..100675e3d14 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -466,6 +466,12 @@ def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: logging_obj.carry_failed_attempt_usage() +def reopen_success_logging_for_next_attempt(kwargs: Mapping[str, object]) -> None: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.reopen_success_logging_for_next_attempt() + + def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None: logging_obj: Final = kwargs.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): @@ -900,6 +906,7 @@ async def run_async_fallback( if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) + reopen_success_logging_for_next_attempt(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) response = await ctx.with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( diff --git a/tests/unit/test_router/test_router_retries.py b/tests/unit/test_router/test_router_retries.py index 857b5052d6c..399c78f03e1 100644 --- a/tests/unit/test_router/test_router_retries.py +++ b/tests/unit/test_router/test_router_retries.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +from collections.abc import Callable +from datetime import datetime from typing import Final from unittest.mock import AsyncMock @@ -13,6 +15,7 @@ from pydantic import TypeAdapter import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.router import AllowedFailsPolicy, RetryPolicy import os @@ -334,6 +337,100 @@ async def test_authentication_errors_are_not_retried( assert len(route.calls) == 1 +class _SuccessLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[tuple[object, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload: Final = TypeAdapter(dict[str, object]).validate_python(kwargs["standard_logging_object"]) + self.logged = (*self.logged, (payload["id"], payload["response_cost"])) + + +def _chat_completion(response_id: str, finish_reason: str, completion_tokens: int) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": response_id, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "finish_reason": finish_reason, "message": {"role": "assistant", "content": "hi"}} + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": completion_tokens, + "total_tokens": 10 + completion_tokens, + }, + }, + ) + + +async def _wait_until(condition: Callable[[], bool]) -> None: + for _ in range(200): + if condition(): + return + await asyncio.sleep(0.01) + + +@pytest.mark.parametrize("deferred", [False, True]) +@pytest.mark.parametrize("redispatch", ["retry", "content_policy_fallback"]) +@pytest.mark.asyncio +async def test_attempt_after_a_discarded_response_logs_its_own_cost( + deferred: bool, + redispatch: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[ + _chat_completion("chatcmpl-filtered", "content_filter", completion_tokens=40), + _chat_completion("chatcmpl-served", "stop", completion_tokens=4), + ] + ) + recorder: Final = _SuccessLogRecorder() + messages: Final = [{"role": "user", "content": "hi"}] + logging_obj: Final = Logging( + model="retry-test-model", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="discarded-attempt", + function_id="discarded-attempt", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.defer_async_logging = deferred + retry_policy: Final = ( + {"retry_policy": RetryPolicy(ContentPolicyViolationErrorRetries=1)} if redispatch == "retry" else {} + ) + router: Final = Router( + model_list=[ + {"model_name": group, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}} + for group in ("retry-test-model", "fallback-model") + ], + content_policy_fallbacks=[{"retry-test-model": ["fallback-model"]}], + **retry_policy, + ) + + response: Final = await router.acompletion( + model="retry-test-model", messages=messages, litellm_logging_obj=logging_obj + ) + if deferred: + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + expected_ids: Final = ("chatcmpl-served",) if deferred else ("chatcmpl-filtered", "chatcmpl-served") + await _wait_until(lambda: len(recorder.logged) >= len(expected_ids)) + + assert response.id == "chatcmpl-served" + assert tuple(logged_id for logged_id, _ in recorder.logged) == expected_ids + assert recorder.logged[-1][1] == response._hidden_params["response_cost"] + + def test_retry_rate_limit_error_with_healthy_deployments(): """ Test 1. It SHOULD retry when there is a rate limit error and len(healthy_deployments) > 0