diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 99ebf479b20..4fd88098123 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -812,6 +812,7 @@ class Logging(LiteLLMLoggingBaseClass): "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, + "current_attempt_api_call_start_time": None, "litellm_params": litellm_params, } new_call.streaming_chunks = [] @@ -1507,7 +1508,9 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.model_call_details["api_call_start_time"] = datetime.datetime.now() + api_call_start_time: Final = datetime.datetime.now() + self.model_call_details["api_call_start_time"] = api_call_start_time + self.model_call_details["current_attempt_api_call_start_time"] = api_call_start_time # Set-once first provider-handoff instant. api_call_start_time # is overwritten on every retry, so it can't measure one-time # preprocessing; pinning the first attempt excludes retry loops @@ -2330,6 +2333,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def discard_partial_usage_for_failure(self) -> None: + """The router is about to open the next attempt of this request: the usage a superseded attempt + stashed here belongs to that attempt's own failure log, and the next attempt's log carries its own.""" + self.model_call_details.pop("combined_usage_object", None) + self.model_call_details.pop("response_cost", None) + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" @@ -2340,6 +2349,7 @@ class Logging(LiteLLMLoggingBaseClass): def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" + self.model_call_details["current_attempt_api_call_start_time"] = None usage: Final = self.model_call_details.pop("combined_usage_object", None) response_cost: Final = self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 78cfb58bedf..6656c3d7a3d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2154,12 +2154,7 @@ class CustomStreamWrapper: traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback_exception) await self.aclose() self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: @@ -2286,15 +2281,18 @@ class CustomStreamWrapper: return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: - traceback_exception: Final = traceback.format_exc() - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback.format_exc()) self._handle_stream_fallback_error(e) + def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None: + if self.logging_obj is None: + return + self._record_partial_usage_for_failure() + attempt_logging_obj: Final = self.logging_obj.attempt_scoped_copy() + self._failure_logging = asyncio.create_task( + attempt_logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) + ) + async def wait_for_failure_logging(self, timeout: float) -> None: """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" if self._failure_logging is None: @@ -2306,9 +2304,9 @@ class CustomStreamWrapper: A stream that breaks mid-flight still billed the provider for the chunks 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. A request that - later recovers via a router fallback overwrites this with the combined - success log on the same request id, so this never double counts. + handler records the real partial spend instead of zero. The router + 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/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 6799b82bf06..c07f796e033 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -85,9 +85,10 @@ class PassThroughStreamingHandler: raw_bytes=raw_bytes, stream_context=stream_context, ) + attempt_logging_obj: Final = litellm_logging_obj.attempt_scoped_copy() try: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=litellm_logging_obj.dispatch_failure_handlers( + async_coroutine=attempt_logging_obj.dispatch_failure_handlers( exception, traceback.format_exc(), prefer_async_handlers=True ) ) @@ -294,10 +295,6 @@ class PassThroughStreamingHandler: @staticmethod def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: - """The provider `event: error` frame of a stream the router closed because a retry or fallback took the - request over. That attempt is logged as its own failure, on a copy taken now so the retry's rebinding of - the shared logging object never reaches it. A stream that ran to its end or that the client closed keeps - logging through the success route, so its usage is still billed.""" from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( closing_superseded_attempt, parse_anthropic_error_event, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 29749afb429..91928a72243 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1009,7 +1009,7 @@ def _billed_across_attempts( ) -def _failure_usage_to_lift( +def _current_failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched: bool, @@ -1023,9 +1023,6 @@ def _failure_usage_to_lift( chat-completions provider) without filling optional_params. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") - failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) - if failed_attempts: - return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): @@ -1049,6 +1046,19 @@ def _failure_usage_to_lift( return estimated_usage, 0.0 +def _failure_usage_to_lift( + model_call_details: Mapping[str, object], + request_body: Mapping[str, object], + dispatched: bool, +) -> tuple[object, object] | None: + current_attempt: Final = _current_failure_usage_to_lift(model_call_details, request_body, dispatched) + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if not failed_attempts: + return current_attempt + current_usage, current_cost = current_attempt if current_attempt is not None else (None, None) + return _billed_across_attempts(failed_attempts, current_usage, current_cost) + + _EMPTY_LIFT: Final = MappingProxyType({}) @@ -1176,7 +1186,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, - dispatched=_first_handoff is not None, + dispatched=_model_call_details.get("current_attempt_api_call_start_time") is not None, ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), diff --git a/litellm/router.py b/litellm/router.py index 8d77fd0dac5..3bc9bfa5f4d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -233,6 +233,7 @@ 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, @@ -6405,6 +6406,7 @@ class Router: verbose_router_logger.debug( "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget ) + 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 06bd3329cc1..31a664a38bc 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,12 +8,14 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -49,6 +51,7 @@ else: # (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. _REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) _NO_MODEL_CALL_DETAILS: Final[Mapping[str, object]] = MappingProxyType({}) +_FALLBACK_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) def _trigger_cooldown_for_failed_deployment( @@ -438,6 +441,20 @@ 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 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.carry_failed_attempt_usage() + + +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): + return None + return logging_obj if logging_obj.model_call_details.get("has_logged_async_failure", False) is True else None + + def _request_metadata_bucket(kwargs: Mapping[str, object]) -> Mapping[str, object] | None: bucket: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) return bucket if isinstance(bucket, Mapping) else None @@ -836,6 +853,9 @@ async def run_async_fallback( ) continue attempted.record(attempt_key) + failure_callbacks_blocked_logging_obj = _failure_callbacks_blocked_logging_object( + _FALLBACK_KWARGS_ADAPTER.validate_python(kwargs) + ) try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -859,6 +879,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors + 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( @@ -881,13 +902,12 @@ async def run_async_fallback( kwargs=kwargs, original_exception=original_exception, ) - logging_obj = kwargs.get("litellm_logging_obj") - if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False): + if failure_callbacks_blocked_logging_obj is not None: _trigger_cooldown_for_failed_deployment( litellm_router=litellm_router, kwargs=kwargs, exception=e, - model_call_details=logging_obj.model_call_details, + model_call_details=failure_callbacks_blocked_logging_obj.model_call_details, ) raise error_from_fallbacks 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 e0b2e4a336f..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 @@ -1,13 +1,16 @@ import json import os import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from dataclasses import dataclass +from pathlib import Path from typing import Final, Literal import anthropic import httpx import pytest +import yaml from integration._support.anthropic_sse import ( LIFECYCLE, Attempts, @@ -27,6 +30,7 @@ from integration._support.anthropic_sse import ( ) from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue from redis import Redis @@ -393,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: @@ -463,6 +469,94 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> assert rows == [{"attempted": "1", "budget": "1"}], rows +def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> AbstractContextManager[Gateway]: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url, "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret"}, + config=path, + ) + + +def _sink_events(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + for batch in batches: + yield from json.loads(batch.body) + + +class _SinkReader: + """Reads the callback sink across polls: drain() consumes the sink's queue, so earlier batches are kept here.""" + + def __init__(self, sink: Wire, model: str) -> None: + self._sink: Final = sink + self._model: Final = model + self._batches: tuple[Request, ...] = () + + def delivered_usage(self) -> tuple[tuple[str, int, int], ...]: + self._batches = (*self._batches, *self._sink.drain()) + return tuple( + (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) + for event in _sink_events(self._batches) + if event.get("model_group") == self._model + ) + + +def _request_rows(model: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s', (model,) + ) + + +@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, attempts)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_completed_by_retry(status, events, prompt, model, wire) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert sorted(delivered) == [("failure", 5, 1), ("success", 5, 3)], delivered + + +def test_every_dropped_attempt_logs_its_own_failure_and_the_request_row_still_bills_the_dropped_usage( + gateway: Gateway, tmp_path: Path +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts, failing_attempts=99)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_stream_failed_after_message_start(status, events, wire, attempts=2) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert delivered == (("failure", 5, 1), ("failure", 5, 1)), delivered + 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 _string_values(cache: Redis) -> tuple[bytes, ...]: keys: Final = tuple(key for key in cache.scan_iter(count=1000) if cache.type(key) == b"string") return tuple(value for value in cache.mget(keys) if value is not None) if keys else () diff --git a/tests/unit/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py index fc8daab3899..c387c9bf103 100644 --- a/tests/unit/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/unit/litellm_core_utils/test_max_streaming_duration.py @@ -70,8 +70,10 @@ class TestCustomStreamWrapperMaxDuration: """__anext__ should check the limit before iterating, dispatching the same failure-callback/logging path every other stream failure goes through (dispatch_failure_handlers is async on the real Logging class, - so the mock needs to be awaitable too).""" + so the mock needs to be awaitable too, and the wrapper dispatches it + through the attempt-scoped copy of the logging object).""" wrapper = _make_custom_stream_wrapper() + wrapper.logging_obj.attempt_scoped_copy.return_value = wrapper.logging_obj wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index c07fe2d1c31..5e83df872a4 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +import traceback import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var @@ -19,6 +20,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkEarlyReturn, _ProviderChunkParsed, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -3893,6 +3895,62 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): assert isinstance(logging_obj.model_call_details["response_cost"], float) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: list[tuple[int, int]] = [] + self.logged = asyncio.Event() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs["standard_logging_object"] + self.failures.append((payload["prompt_tokens"], payload["completion_tokens"])) + self.logged.set() + + +@pytest.mark.asyncio +async def test_log_stream_failure_bills_the_dropped_attempt_even_after_the_router_discards_its_usage(): + """The dropped attempt's failure log is queued, and the router discards the stashed usage as it opens the + fallback hop before that log runs. The log must still bill what the dropped attempt consumed, and the shared + logging object must not be marked as already failed, since the hop's own failure still has to reach the + integrations.""" + recorder = _FailureRecorder() + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-copy", + function_id="1245", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-copy", + created=1742056047, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The", role="assistant"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(httpx.ReadError("connection dropped")) + logging_obj.discard_partial_usage_for_failure() + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert recorder.failures == [(30, 1)] + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + + def test_record_partial_usage_for_failure_noop_without_chunks(): """With no chunks delivered there is nothing billed to recover, so the failure stash must stay absent and not force a zero-usage row. diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 5c0e9f1a305..49a74c4b240 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2831,6 +2831,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def dispatch_failure_handlers(self, *args, **kwargs): return None + def attempt_scoped_copy(self): + return self + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py index fdd21afbc67..b5e6c2a14f3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py @@ -218,6 +218,7 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop( logging_obj.model_call_details = {"model": model, "stream": True} logging_obj.litellm_params = {} logging_obj.get_router_model_id.return_value = None + logging_obj.attempt_scoped_copy.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() _, took, lags = await timed_with_loop_lags( diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 8ce41129aec..0c21729e893 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -750,6 +750,107 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +def _anthropic_stream_dropping_after(frames: tuple[bytes, ...]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for frame in frames: + yield frame + raise httpx.ReadError("upstream closed the connection") + + mock.aiter_bytes = _aiter_bytes + return mock + + +def _anthropic_messages_attempt(logging_obj: LiteLLMLoggingObj, response): + return PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + +def _anthropic_messages_logging_obj(recorder: _EventRecorder, call_id: str) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + +async def _drop_after_message_start(logging_obj: LiteLLMLoggingObj, message_id: str, input_tokens: int) -> None: + dropped = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_dropping_after((_anthropic_message_start_frame(message_id, input_tokens),)) + ) + with pytest.raises(httpx.ReadError): + async for _ in dropped: + pass + + +@pytest.mark.asyncio +async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logging_object(): + """The router retries a /v1/messages stream the provider drops after message_start on the same logging + object. The dropped attempt's failure log bills its own usage and leaves the shared object free: the + retry's success log carries the retry's usage and no stale exception, while the stash stays for the + failure row of a request that runs out of retries until the router opens the next attempt.""" + recorder = _EventRecorder() + logging_obj = _anthropic_messages_logging_obj(recorder, "test-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + + assert recorder.failure_kwargs[0]["standard_logging_object"]["prompt_tokens"] == 52 + assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadError) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + assert "exception" not in logging_obj.model_call_details + stashed = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 52 + + logging_obj.discard_partial_usage_for_failure() + retry = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + ) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + assert recorder.success_kwargs[0].get("exception") is None + assert len(recorder.failure_kwargs) == 1 + + +@pytest.mark.asyncio +async def test_every_dropped_attempt_logs_its_own_failure(): + """Two attempts of one /v1/messages request dropped in a row each log a failure billing their own usage; + the first attempt's log must not swallow the second's, and the shared object ends up holding the last + attempt's usage for the request's failure row.""" + recorder = _EventRecorder() + logging_obj = _anthropic_messages_logging_obj(recorder, "test-every-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + logging_obj.discard_partial_usage_for_failure() + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + assert [kwargs["standard_logging_object"]["prompt_tokens"] for kwargs in recorder.failure_kwargs] == [52, 9] + assert all(isinstance(kwargs["exception"], httpx.ReadError) for kwargs in recorder.failure_kwargs) + stashed = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", @@ -788,9 +889,9 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: ) -async def _settle(recorded: list[dict[str, object]]) -> None: +async def _settle(recorded: list[dict[str, object]], count: int = 1) -> None: for _ in range(300): - if recorded: + if len(recorded) >= count: return await asyncio.sleep(0.01) @@ -1036,8 +1137,6 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr success_payload = recorder.success_kwargs[0]["standard_logging_object"] assert success_payload["status"] == "success" assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) - - def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index ec7104fa459..ebe661cf43f 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1921,6 +1921,47 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): assert response_cost == 0.0 +@pytest.mark.parametrize( + "current_attempt_api_call_start_time", + [None, 1700000001.0], + ids=["retry-not-dispatched", "retry-dispatched"], +) +def test_a_retry_failure_bills_input_only_for_dispatched_attempts( + current_attempt_api_call_start_time: float | None, +) -> None: + from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift + from litellm.types.utils import Usage + + prior_usage: Final = Usage(prompt_tokens=4, completion_tokens=2, total_tokens=6) + estimated_retry_usage: Final = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert estimated_retry_usage is not None + + lifted: Final = _failure_fields_to_lift( + { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "current_attempt_api_call_start_time": current_attempt_api_call_start_time, + "failed_attempts_usage": ((prior_usage, 0.25),), + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + } + ) + } + ) + + usage: Final = lifted["combined_usage_object"] + assert isinstance(usage, Usage) + expected_retry_tokens: Final = ( + estimated_retry_usage.prompt_tokens if current_attempt_api_call_start_time is not None else 0 + ) + assert usage.prompt_tokens == prior_usage.prompt_tokens + expected_retry_tokens + assert usage.completion_tokens == prior_usage.completion_tokens + assert usage.total_tokens == prior_usage.total_tokens + expected_retry_tokens + assert lifted["response_cost"] == 0.25 + + @pytest.mark.parametrize( "model_call_details, dispatched", [ diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 637fdfa6873..b6117886dcb 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -16,7 +16,9 @@ from litellm import Router from litellm.integrations.custom_logger import CustomLogger from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import FakeAnthropicMessagesStreamIterator from litellm.litellm_core_utils import get_llm_provider_logic -from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import Choices, Message, ModelResponse, Usage +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, PRE_ROUTING_SELECTED_MODEL_KEY, @@ -144,6 +146,152 @@ class RecordingRouter: return StreamingWrapper() +def _request_logging_obj() -> Logging: + return Logging( + model="primary-model", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop", + function_id="fallback-hop", + ) + + +class _SuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[int, float]] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs["standard_logging_object"] + self.logged.append((payload["completion_tokens"], payload["response_cost"])) + + +class HopLoggingItsOwnSuccessRouter(FakeRouter): + async def async_function_with_fallbacks(self, *args, **kwargs): + hop_response = ModelResponse( + id="chatcmpl-hop", + model="gpt-4o-mini", + choices=[Choices(index=0, message=Message(role="assistant", content="served by the hop"))], + usage=Usage(prompt_tokens=52, completion_tokens=7, total_tokens=59), + ) + now = datetime(2026, 1, 1) + await kwargs["litellm_logging_obj"].async_success_handler(result=hop_response, start_time=now, end_time=now) + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_partial_usage(): + """A stream the failed attempt dropped stashed the usage it consumed on the request's logging object. The + hop that replaces it logs its own usage, so the integrations see the hop's tokens and cost, not the stash.""" + recorder = _SuccessRecorder() + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "ping"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop-success", + function_id="fallback-hop-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime(2026, 1, 1) + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + 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(), + 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 [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),) + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] is None + + +@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): + def __init__(self) -> None: + super().__init__("fallback deployment rate limited") + self.status_code = 429 + self.failed_deployment_id = "fallback-dep-1" + + +def _two_deployment_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_already_logged", [False, True]) +async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool): + """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has + been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed + with the flag already set gets the manual cooldown.""" + router = _two_deployment_router() + logging_obj = _request_logging_obj() + if failure_already_logged: + logging_obj.model_call_details["has_logged_async_failure"] = True + + async def rate_limited_hop(*args, **kwargs): + kwargs["litellm_logging_obj"].model_call_details["has_logged_async_failure"] = True + raise _RateLimitedHop() + + with pytest.raises(_RateLimitedHop): + await run_async_fallback( + litellm_router=router, + 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, + original_function=rate_limited_hop, + metadata={}, + ) + + cooled_down = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) + + @pytest.mark.asyncio async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): """A nested fallback (multi-hop) must keep collecting errors, so the opt-in @@ -1100,9 +1248,9 @@ class TestRunAsyncFallbackTriggersCooldown: async def async_function_with_fallbacks(self, *args, **kwargs): raise RuntimeError("fallback model also failed") - def _logging_obj(self, has_logged_async_failure: bool) -> MagicMock: - logging_obj = MagicMock() - logging_obj.model_call_details = {"has_logged_async_failure": has_logged_async_failure} + def _logging_obj(self, has_logged_async_failure: bool) -> Logging: + logging_obj: Final = _request_logging_obj() + logging_obj.model_call_details["has_logged_async_failure"] = has_logged_async_failure return logging_obj @pytest.mark.asyncio diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index bd3cf738af8..0bc6ebe3da8 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -93,7 +93,7 @@ from litellm.llms.base_llm.vector_store.transformation import( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import CallTypes, CredentialItem +from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -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 @@ -17349,6 +17351,38 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +@pytest.mark.asyncio +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", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="retry-own-usage", + function_id="retry-own-usage", + ) + 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(): + stash_when_retry_opened.append(logging_obj.model_call_details.get("combined_usage_object")) + return _anthropic_messages_retried_stream() + + provider = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) + + stream = await _anthropic_messages_stream_through_router(router, provider, litellm_logging_obj=logging_obj) + body = [chunk async for chunk in stream] + + 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 async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream @@ -17371,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