mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
Merge 5d7a1b8b28 into 4273996985
This commit is contained in:
Mateo Wang
GitHub
commit
40655d279f
15 changed files with 559 additions and 41 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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={})
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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}],
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue