scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 25

fix(router): preserve superseded retry usage for billing

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-10-17 19:11:54 +00:00
parent aa5918c0cb
commit 4df47b58a7
6 changed files with 50 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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

5
litellm/router.py
View file

@ -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

10
litellm/router_utils/fallback_event_handlers.py
View file

@ -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(

13
tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py
View file

@ -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,

24
tests/unit/router_utils/test_fallback_event_handlers.py
View file

@ -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):

13
tests/unit/test_router/test_router.py
View file

@ -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