mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(proxy): keep a retried /v1/messages stream's spend log by logging the superseded attempt as a failure (#45136)
* fix(proxy): keep a retried /v1/messages stream's spend log by logging the superseded attempt as a failure A /v1/messages stream the router retries before any content shares one logging object across both attempts. The superseded attempt's end-of-stream logging ran as a success (a usage-only or placeholder result) and claimed the request's one-shot success log, so the recovered stream's spend row was either missing or written under the superseded attempt's message id. The passthrough logging path now logs an Anthropic stream whose frames carry an `event: error` as a failure with whatever usage the provider reported, the way a dropped stream already does, and the router closes the superseded stream before it opens the retry so that logging is queued first. * fix(proxy): log a provider error frame through an attempt-scoped logging copy The error-frame failure stashed the superseded attempt's usage on the logging object the retry shares, so the retry's success log carried the failed attempt's token counts. The failure now logs through a copy of the logging object with its own model_call_details, and the helper moves into the Anthropic passthrough logging handler so the proxy streaming handler only routes by endpoint type. * fix(proxy): scope the error-frame failure log to streams closed at the frame A stream that ran to its end at the provider's error frame kept logging through the success route before this PR, and the attempt-scoped failure log is now taken only for a stream the router closed at that frame, with the copy made at the close so the retry's deployment never reaches it * chore: drop duplicate litellm import and mark the stream_ended rebind * fix(proxy): log only router-superseded attempts as error-frame failures A client that hangs up at a forwarded error frame closed the stream the same way the router closes a superseded attempt, so a non-retried stream lost its spend. The router now closes a superseded attempt through aclose_superseded_attempt, and only that close routes the attempt to the failure log * test: type the _settle helper's recorded events --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
ee6ca049b3
commit
8bb8ac7f18
10 changed files with 516 additions and 12 deletions
|
|
@ -2300,6 +2300,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.model_call_details["combined_usage_object"] = usage
|
||||
self.model_call_details["response_cost"] = response_cost
|
||||
|
||||
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."""
|
||||
attempt: Final = copy.copy(self)
|
||||
attempt.model_call_details = dict(self.model_call_details)
|
||||
return attempt
|
||||
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import json
|
||||
import re
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
|
||||
from contextvars import ContextVar
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
|
||||
|
||||
|
|
@ -352,6 +353,22 @@ async def aclose_if_supported(stream: object) -> None:
|
|||
await stream.aclose()
|
||||
|
||||
|
||||
_CLOSING_SUPERSEDED_ATTEMPT: Final[ContextVar[bool]] = ContextVar("closing_superseded_attempt", default=False)
|
||||
|
||||
|
||||
async def aclose_superseded_attempt(stream: object) -> None:
|
||||
"""Close an attempt a retry or fallback is taking over, so its end-of-stream logging records a failure."""
|
||||
token: Final = _CLOSING_SUPERSEDED_ATTEMPT.set(True)
|
||||
try:
|
||||
await aclose_if_supported(stream)
|
||||
finally:
|
||||
_CLOSING_SUPERSEDED_ATTEMPT.reset(token)
|
||||
|
||||
|
||||
def closing_superseded_attempt() -> bool:
|
||||
return _CLOSING_SUPERSEDED_ATTEMPT.get()
|
||||
|
||||
|
||||
_RESPONSE_HEADERS_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,12 +10,14 @@ from pydantic import ConfigDict, TypeAdapter
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import ANTHROPIC_BATCHES_ROUTE
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_content_from_model_response,
|
||||
)
|
||||
from litellm.llms.anthropic import get_anthropic_config
|
||||
from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import (
|
||||
build_complete_streaming_response,
|
||||
build_complete_streaming_response_legacy,
|
||||
|
|
@ -254,6 +256,26 @@ class AnthropicPassthroughLoggingHandler:
|
|||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def log_error_frame_as_failure(
|
||||
attempt_logging_obj: LiteLLMLoggingObj,
|
||||
request_body: Mapping[str, object],
|
||||
all_chunks: Sequence[str | bytes],
|
||||
error_event: tuple[str, str, int],
|
||||
) -> None:
|
||||
"""A stream closed at the provider's `event: error` frame is one failed attempt of the request: it logs
|
||||
through the attempt-scoped copy its caller took at the close, so the usage it consumed never reaches
|
||||
the success log of the attempt the router opens in its place."""
|
||||
error_type, message, status_code = error_event
|
||||
await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)(
|
||||
litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks
|
||||
)
|
||||
await attempt_logging_obj.dispatch_failure_handlers(
|
||||
anthropic_error_frame_exception(error_type, message, status_code, str(attempt_logging_obj.model)),
|
||||
"",
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _cost_partial_stream_or_zero(
|
||||
partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj
|
||||
|
|
|
|||
|
|
@ -165,6 +165,7 @@ class PassThroughStreamingHandler:
|
|||
)
|
||||
|
||||
logging_scheduled = False
|
||||
stream_ended = False # rebind-ok: flipped once the upstream ran to its end, so the finally can tell a close-at-frame from a completed stream
|
||||
model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
|
||||
request_body=request_body,
|
||||
url_route=url_route,
|
||||
|
|
@ -224,6 +225,7 @@ class PassThroughStreamingHandler:
|
|||
)
|
||||
if pending:
|
||||
yield pending
|
||||
stream_ended = True # rebind-ok: the finally reads it to spot a close-at-frame
|
||||
# Stream completed cleanly. When the proxy armed deferred
|
||||
# dispatch (post-call guardrails active), park the logging
|
||||
# coroutine on logging_obj instead of enqueueing now, so
|
||||
|
|
@ -267,13 +269,46 @@ class PassThroughStreamingHandler:
|
|||
# a success would double-log the same request.
|
||||
if not logging_scheduled and raw_bytes and response.status_code < 400:
|
||||
logging_scheduled = True
|
||||
superseded_error_event: Final = (
|
||||
None
|
||||
if stream_ended
|
||||
else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes)
|
||||
)
|
||||
try:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine())
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
async_coroutine=(
|
||||
_build_logging_coroutine()
|
||||
if superseded_error_event is None
|
||||
else AnthropicPassthroughLoggingHandler.log_error_frame_as_failure(
|
||||
attempt_logging_obj=litellm_logging_obj.attempt_scoped_copy(),
|
||||
request_body=resolved_request_body,
|
||||
all_chunks=tuple(raw_bytes),
|
||||
error_event=superseded_error_event,
|
||||
)
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)
|
||||
else:
|
||||
bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
|
||||
|
||||
@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,
|
||||
)
|
||||
|
||||
if endpoint_type != EndpointType.ANTHROPIC or not closing_superseded_attempt():
|
||||
return None
|
||||
|
||||
complete_frames, _ = split_complete_sse_frames(b"".join(raw_bytes))
|
||||
return parse_anthropic_error_event(complete_frames)
|
||||
|
||||
@staticmethod
|
||||
async def route_streaming_logging_to_handler(
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
|
|
|
|||
|
|
@ -6035,6 +6035,7 @@ class Router:
|
|||
"""
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
aclose_superseded_attempt,
|
||||
parse_anthropic_error_event,
|
||||
parse_anthropic_refusal_stop_details,
|
||||
)
|
||||
|
|
@ -6044,6 +6045,14 @@ class Router:
|
|||
|
||||
source_iterator: Final = response
|
||||
|
||||
async def close_source() -> None:
|
||||
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
|
||||
await aclose_if_supported(source_iterator)
|
||||
|
||||
async def close_superseded() -> None:
|
||||
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
|
||||
await aclose_superseded_attempt(source_iterator)
|
||||
|
||||
async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
|
|
@ -6130,11 +6139,11 @@ class Router:
|
|||
model,
|
||||
initial_kwargs,
|
||||
wrapper,
|
||||
close_superseded,
|
||||
):
|
||||
yield item
|
||||
finally:
|
||||
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
|
||||
await aclose_if_supported(source_iterator)
|
||||
await close_source()
|
||||
|
||||
# Referenced by stream_with_fallbacks via closure - assigned here, before
|
||||
# the generator body ever runs, so the reference resolves fine despite
|
||||
|
|
@ -6150,6 +6159,7 @@ class Router:
|
|||
model: str,
|
||||
initial_kwargs: dict[str, Any], # mutable-ok: handed to _aanthropic_messages_fallback_attempt, which mutates it
|
||||
wrapper: "FallbackAwareAnthropicMessagesStream",
|
||||
close_superseded: Callable[[], Awaitable[None]],
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
"""Turns a source-iterator failure into a fallback attempt or the error reaching the caller."""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
|
@ -6169,6 +6179,7 @@ class Router:
|
|||
)
|
||||
if fallback_error is None:
|
||||
raise stream_error
|
||||
await close_superseded()
|
||||
outcome: Final = await self._astream_retry_same_group(
|
||||
fallback_error,
|
||||
initial_kwargs,
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ FirstAttempt = Literal[
|
|||
"drop_after_message_start",
|
||||
"error_frame",
|
||||
"error_frame_after_message_start",
|
||||
"error_frame_after_content",
|
||||
"http_status",
|
||||
"drop_after_content",
|
||||
]
|
||||
|
|
@ -66,6 +67,8 @@ class _Failure:
|
|||
return status_reply(self.status)
|
||||
case "drop_after_content":
|
||||
return dropping_reply(chunks, abort_after=3)
|
||||
case "error_frame_after_content":
|
||||
return stream_reply((*chunks[:3], error_frame(self.status, f"scripted {self.status}")))
|
||||
|
||||
|
||||
_MID_STREAM_FAILURES: Final = (
|
||||
|
|
@ -124,12 +127,12 @@ def _stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[int, tuple
|
|||
return response.status_code, parse_sse(response.text)
|
||||
|
||||
|
||||
def _rows(request_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,))
|
||||
|
||||
|
||||
def _success_rows(request_id: str) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)),
|
||||
lambda rows: len(rows) >= 1,
|
||||
seconds=70,
|
||||
)
|
||||
return eventually(lambda: _rows(request_id), lambda rows: len(rows) >= 1, seconds=70)
|
||||
|
||||
|
||||
def _assert_completed_by_retry(
|
||||
|
|
@ -141,6 +144,8 @@ def _assert_completed_by_retry(
|
|||
assert delta_text(events) == _TEXT, events
|
||||
assert [request.target for request in wire.drain()] == ["/v1/messages"] * attempts
|
||||
assert _success_rows(_served_id(prompt, attempts)) == [{"status": "success", "model_group": model}]
|
||||
superseded_rows: Final = {attempt: _rows(_served_id(prompt, attempt)) for attempt in range(1, attempts)}
|
||||
assert not any(superseded_rows.values()), superseded_rows
|
||||
|
||||
|
||||
def _assert_rejected_before_content(
|
||||
|
|
@ -302,6 +307,23 @@ def test_stream_dropping_after_content_is_not_retried(gateway: Gateway) -> None:
|
|||
assert [request.target for request in wire.drain()] == ["/v1/messages"]
|
||||
|
||||
|
||||
def test_stream_error_frame_after_content_is_not_retried_and_keeps_its_spend_row(gateway: Gateway) -> None:
|
||||
prompt: Final = _prompt()
|
||||
attempts: Final = Attempts()
|
||||
with (
|
||||
wire_server(_upstream(prompt, _Failure("error_frame_after_content", 529), attempts)) as wire,
|
||||
gateway.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(gateway, _body(model, prompt, "deployment"))
|
||||
assert status == 200, events
|
||||
assert event_types(events) == (*LIFECYCLE[:3], "error"), events
|
||||
assert message_id(events) == _served_id(prompt, 1), events
|
||||
assert delta_text(events) == _TEXT, events
|
||||
assert [request.target for request in wire.drain()] == ["/v1/messages"]
|
||||
assert [row["model_group"] for row in _success_rows(_served_id(prompt, 1))] == [model]
|
||||
|
||||
|
||||
def test_stream_rejected_with_401_before_it_opens_is_not_retried(gateway: Gateway) -> None:
|
||||
prompt: Final = _prompt()
|
||||
attempts: Final = Attempts()
|
||||
|
|
|
|||
|
|
@ -201,6 +201,10 @@ class RecordingObserver(PromptCacheObserver):
|
|||
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
self.finished.set()
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await super().async_log_failure_event(kwargs, response_obj, start_time, end_time)
|
||||
self.finished.set()
|
||||
|
||||
|
||||
def native_response():
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -5,11 +5,10 @@ from typing import Any, Dict, Final, List
|
|||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data
|
||||
|
|
@ -2644,6 +2643,58 @@ class TestRecordPartialUsageForFailure:
|
|||
assert "response_cost" not in logging_obj.model_call_details
|
||||
|
||||
|
||||
class _FailureRecorder(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.failure_kwargs = []
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.failure_kwargs.append(kwargs)
|
||||
|
||||
|
||||
class TestLogErrorFrameAsFailure:
|
||||
"""A stream closed at the provider's `event: error` frame is one failed attempt: its failure log bills the
|
||||
usage streamed before the frame on the attempt-scoped copy its caller hands it, and the request's shared
|
||||
logging state stays as it was, so the attempt the router opens next logs with nothing left behind."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self):
|
||||
recorder = _FailureRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-error-frame-failure",
|
||||
function_id="test-error-frame-failure",
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
chunks = [
|
||||
*TestRecordPartialUsageForFailure()._interrupted_chunks()[:1],
|
||||
TestRecordPartialUsageForFailure._sse(
|
||||
"error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}
|
||||
),
|
||||
]
|
||||
|
||||
await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure(
|
||||
attempt_logging_obj=logging_obj.attempt_scoped_copy(),
|
||||
request_body={"model": "claude-sonnet-5", "stream": True},
|
||||
all_chunks=chunks,
|
||||
error_event=("overloaded_error", "Overloaded", 529),
|
||||
)
|
||||
|
||||
assert len(recorder.failure_kwargs) == 1
|
||||
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
|
||||
assert failure_payload["status"] == "failure"
|
||||
assert failure_payload["prompt_tokens"] == 52
|
||||
assert failure_payload["response_cost"] > 0
|
||||
assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError)
|
||||
assert "combined_usage_object" not in logging_obj.model_call_details
|
||||
assert "response_cost" not in logging_obj.model_call_details
|
||||
assert logging_obj.should_run_logging("async_failure")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("all_chunks", "interrupted"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthr
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import aclose_superseded_attempt
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import (
|
||||
PassThroughStreamingHandler,
|
||||
)
|
||||
|
|
@ -749,6 +750,294 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception(
|
|||
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)
|
||||
|
||||
|
||||
def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes:
|
||||
return _anthropic_sse(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": message_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": input_tokens, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_completed_stream_tail() -> tuple[bytes, ...]:
|
||||
return (
|
||||
_anthropic_sse(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
_anthropic_sse(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}},
|
||||
),
|
||||
_anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
_anthropic_sse(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}},
|
||||
),
|
||||
_anthropic_sse("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
async def _settle(recorded: list[dict[str, object]]) -> None:
|
||||
for _ in range(300):
|
||||
if recorded:
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
def _anthropic_overloaded_error_frame() -> bytes:
|
||||
return _anthropic_sse("error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}})
|
||||
|
||||
|
||||
def _anthropic_stream_of(frames: tuple[bytes, ...]):
|
||||
mock = MagicMock(spec=httpx.Response)
|
||||
mock.status_code = 200
|
||||
|
||||
async def _aiter_bytes():
|
||||
for frame in frames:
|
||||
yield frame
|
||||
|
||||
mock.aiter_bytes = _aiter_bytes
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"frames, billed",
|
||||
[
|
||||
pytest.param(
|
||||
(_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start"
|
||||
),
|
||||
pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_closed_at_the_frame(frames, billed):
|
||||
"""The router retries a /v1/messages stream whose provider sent `event: error` before any content by closing
|
||||
it at the frame. That attempt never completed: it logs once as a failure carrying whatever usage the provider
|
||||
had reported, and the request's one-shot success log stays free for the retry that shares the logging object."""
|
||||
recorder = _EventRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-error-frame",
|
||||
function_id="test-error-frame",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
stream = PassThroughStreamingHandler.chunk_processor(
|
||||
response=_anthropic_stream_of(frames),
|
||||
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",
|
||||
)
|
||||
received = []
|
||||
async for chunk in stream:
|
||||
received.append(chunk)
|
||||
if b"event: error" in chunk:
|
||||
break
|
||||
await aclose_superseded_attempt(stream)
|
||||
await _settle(recorder.failure_kwargs)
|
||||
|
||||
assert received == list(frames)
|
||||
assert recorder.success_kwargs == []
|
||||
assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True
|
||||
assert "combined_usage_object" not in logging_obj.model_call_details
|
||||
assert "response_cost" not in logging_obj.model_call_details
|
||||
assert logging_obj.should_run_logging("async_failure")
|
||||
assert len(recorder.failure_kwargs) == 1
|
||||
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
|
||||
assert failure_payload["status"] == "failure"
|
||||
assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError)
|
||||
if billed:
|
||||
assert failure_payload["prompt_tokens"] == 52
|
||||
assert failure_payload["response_cost"] > 0
|
||||
else:
|
||||
assert failure_payload["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_logs_a_stream_that_ended_at_the_provider_error_frame_through_the_success_route():
|
||||
"""A provider error frame that ends a stream nobody closed early is the request's final answer, and that
|
||||
stream keeps logging the way it always has: once, through the success route, with the usage it carried."""
|
||||
recorder = _EventRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-final-error-frame",
|
||||
function_id="test-final-error-frame",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
frames = (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())
|
||||
received = [
|
||||
chunk
|
||||
async for chunk in PassThroughStreamingHandler.chunk_processor(
|
||||
response=_anthropic_stream_of(frames),
|
||||
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",
|
||||
)
|
||||
]
|
||||
await _settle(recorder.success_kwargs)
|
||||
|
||||
assert received == list(frames)
|
||||
assert recorder.failure_kwargs == []
|
||||
success_payload = recorder.success_kwargs[0]["standard_logging_object"]
|
||||
assert success_payload["status"] == "success"
|
||||
assert success_payload["prompt_tokens"] == 52
|
||||
assert logging_obj.model_call_details["prompt_cache_response_complete"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_bills_a_stream_the_client_closed_at_the_provider_error_frame():
|
||||
"""The Anthropic SDK raises on a provider `event: error` frame and drops the connection, so a non-retried
|
||||
stream is closed at the frame by the client, not by a retry taking over. Its usage must still reach the
|
||||
success route, or the tokens the provider billed never reach the key's spend."""
|
||||
recorder = _EventRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-client-closed-at-error-frame",
|
||||
function_id="test-client-closed-at-error-frame",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
stream = PassThroughStreamingHandler.chunk_processor(
|
||||
response=_anthropic_stream_of((_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())),
|
||||
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",
|
||||
)
|
||||
async for chunk in stream:
|
||||
if b"event: error" in chunk:
|
||||
break
|
||||
await stream.aclose()
|
||||
await _settle(recorder.success_kwargs)
|
||||
|
||||
assert recorder.failure_kwargs == []
|
||||
success_payload = recorder.success_kwargs[0]["standard_logging_object"]
|
||||
assert success_payload["status"] == "success"
|
||||
assert success_payload["prompt_tokens"] == 52
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_superseded_attempt_failure_log_keeps_its_own_deployment_when_the_retry_rebinds_the_logging_object():
|
||||
"""The router rebinds the shared logging object to the retry's deployment right after closing the superseded
|
||||
stream, before the logging worker runs. The failure log of the superseded attempt must still name the
|
||||
deployment that served it, or the router's cooldown lands on the deployment serving the retry."""
|
||||
recorder = _EventRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-superseded-deployment",
|
||||
function_id="test-superseded-deployment",
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="claude-sonnet-5", litellm_params={"model_info": {"id": "deployment-1"}}, optional_params={}
|
||||
)
|
||||
superseded = PassThroughStreamingHandler.chunk_processor(
|
||||
response=_anthropic_stream_of((_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())),
|
||||
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",
|
||||
)
|
||||
async for chunk in superseded:
|
||||
if b"event: error" in chunk:
|
||||
break
|
||||
await aclose_superseded_attempt(superseded)
|
||||
logging_obj.update_environment_variables(
|
||||
model="claude-sonnet-5-retry", litellm_params={"model_info": {"id": "deployment-2"}}, optional_params={}
|
||||
)
|
||||
await _settle(recorder.failure_kwargs)
|
||||
|
||||
failure_kwargs = recorder.failure_kwargs[0]
|
||||
assert failure_kwargs["litellm_params"]["model_info"]["id"] == "deployment-1"
|
||||
assert failure_kwargs["model"] == "claude-sonnet-5"
|
||||
assert logging_obj.model_call_details["litellm_params"]["model_info"]["id"] == "deployment-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame():
|
||||
"""Both attempts of a retried /v1/messages stream log through the request's one logging object. The
|
||||
superseded attempt's failure log must not leave the usage it consumed behind, or the retry's success
|
||||
log bills the failed attempt's tokens instead of its own."""
|
||||
recorder = _EventRecorder()
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="claude-sonnet-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-retry-own-usage",
|
||||
function_id="test-retry-own-usage",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
|
||||
def attempt(frames: tuple[bytes, ...]):
|
||||
return PassThroughStreamingHandler.chunk_processor(
|
||||
response=_anthropic_stream_of(frames),
|
||||
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",
|
||||
)
|
||||
|
||||
superseded = attempt((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame()))
|
||||
async for chunk in superseded:
|
||||
if b"event: error" in chunk:
|
||||
break
|
||||
await aclose_superseded_attempt(superseded)
|
||||
await _settle(recorder.failure_kwargs)
|
||||
retry = attempt((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail()))
|
||||
async for _ in retry:
|
||||
pass
|
||||
await _settle(recorder.success_kwargs)
|
||||
|
||||
failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
|
||||
assert failure_payload["prompt_tokens"] == 52
|
||||
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}],
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.otel.plumbing.context import set_request_root_span
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import closing_superseded_attempt
|
||||
from litellm.llms.anthropic.pass_through.messages.agentic_streaming_iterator import (
|
||||
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES,
|
||||
)
|
||||
|
|
@ -16004,6 +16005,7 @@ class _AnthropicMessagesFakeByteStream:
|
|||
self._chunks = list(chunks)
|
||||
self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}}
|
||||
self.closed = False
|
||||
self.closed_as_superseded = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
|
@ -16014,6 +16016,7 @@ class _AnthropicMessagesFakeByteStream:
|
|||
return self._chunks.pop(0)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt()
|
||||
self.closed = True
|
||||
|
||||
|
||||
|
|
@ -17256,6 +17259,30 @@ 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_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
|
||||
logging must be queued before the retry starts producing its own: the router closes the dropped stream
|
||||
before it asks the group for another."""
|
||||
router = _anthropic_messages_retry_router(num_retries=1)
|
||||
superseded = _AnthropicMessagesFakeByteStream(
|
||||
[_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()]
|
||||
)
|
||||
superseded_closed_when_retry_opened: list[bool] = []
|
||||
|
||||
def retry_stream():
|
||||
superseded_closed_when_retry_opened.append(superseded.closed_as_superseded)
|
||||
return _anthropic_messages_retried_stream()
|
||||
|
||||
provider = _AnthropicMessagesScriptedProvider(lambda: superseded, retry_stream)
|
||||
|
||||
stream = await _anthropic_messages_stream_through_router(router, provider)
|
||||
body = [chunk async for chunk in stream]
|
||||
|
||||
assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")]
|
||||
assert superseded_closed_when_retry_opened == [True]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_ptu_share_holders_messages_stream_dropped_before_content_is_retried_on_its_shared_group(monkeypatch):
|
||||
"""The same-group retry has to look the group's health up as the calling team; looked up with no
|
||||
|
|
@ -18387,7 +18414,7 @@ async def test_aanthropic_messages_recover_stream_error_flushes_buffered_frames_
|
|||
|
||||
with patch.object(router, "_aanthropic_messages_fallback_attempt") as mock_attempt:
|
||||
recovery = router._aanthropic_messages_recover_stream_error(
|
||||
declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper()
|
||||
declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock()
|
||||
)
|
||||
with pytest.raises(BedrockError) as exc_info:
|
||||
await drain(recovery)
|
||||
|
|
@ -18408,7 +18435,7 @@ async def test_aanthropic_messages_recover_stream_error_hands_converted_raise_to
|
|||
|
||||
with patch.object(router, "_aanthropic_messages_fallback_attempt", new=fake_attempt):
|
||||
recovery = router._aanthropic_messages_recover_stream_error(
|
||||
raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper()
|
||||
raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock()
|
||||
)
|
||||
collected = [chunk async for chunk in recovery]
|
||||
assert collected == [b"fallback"]
|
||||
|
|
@ -18466,6 +18493,25 @@ async def test_anthropic_messages_post_first_chunk_error_skips_fallback():
|
|||
mock_fallback.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_stream_the_client_closes_at_a_forwarded_error_is_not_closed_as_superseded():
|
||||
"""An error frame forwarded to the client ends that request, and the SDK hangs up on it. No retry took the
|
||||
request over, so the source must not be closed as a superseded attempt, which would log it as a failure and
|
||||
drop the usage the provider billed for it."""
|
||||
router = _anthropic_messages_make_router()
|
||||
error_chunk = _anthropic_messages_overloaded_error_chunk()
|
||||
source = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("partial answer"), error_chunk])
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"})
|
||||
async for chunk in wrapped:
|
||||
if chunk == error_chunk:
|
||||
break
|
||||
await wrapped.aclose()
|
||||
|
||||
assert source.closed is True
|
||||
assert source.closed_as_superseded is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames():
|
||||
"""A non-retriable error arriving while lifecycle frames are still
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue