mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(router): support anthropic messages streaming fallback
This commit is contained in:
parent
d86336a7c6
commit
f24a24547b
5 changed files with 790 additions and 5 deletions
|
|
@ -1,7 +1,8 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol, runtime_checkable
|
||||
|
||||
import httpx
|
||||
|
|
@ -23,6 +24,79 @@ INCOMPLETE_STREAM_ERROR_MESSAGE: Final = (
|
|||
"the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated."
|
||||
)
|
||||
|
||||
ANTHROPIC_STREAM_ERROR_STATUS_CODE_MAP: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{ # mutable-ok: immediately frozen into a read-only error lookup
|
||||
"invalid_request_error": 400,
|
||||
"authentication_error": 401,
|
||||
"permission_error": 403,
|
||||
"not_found_error": 404,
|
||||
"request_too_large": 413,
|
||||
"rate_limit_error": 429,
|
||||
"api_error": 500,
|
||||
"internal_server_error": 500,
|
||||
"overloaded_error": 529,
|
||||
"timeout_error": 504,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class AnthropicMessagesSSEDecoder:
|
||||
"""Incrementally separates SSE frames without retaining event state."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._buffer = b""
|
||||
|
||||
@staticmethod
|
||||
def _frame_boundary(buffer: bytes) -> tuple[int, int] | None:
|
||||
boundaries: Final = tuple(
|
||||
(position, len(separator))
|
||||
for separator in (b"\r\n\r\n", b"\n\n", b"\r\r")
|
||||
if (position := buffer.find(separator)) >= 0
|
||||
)
|
||||
return min(boundaries) if boundaries else None
|
||||
|
||||
@staticmethod
|
||||
def event_data(frame: bytes) -> tuple[str | None, object | None]:
|
||||
event: str | None = None # rebind-ok: the last event field defines this frame's event type
|
||||
data_lines: Final[list[str]] = [] # mutable-ok: SSE permits multiple data fields per event
|
||||
for raw_line in frame.splitlines():
|
||||
line = raw_line.decode("utf-8", errors="replace")
|
||||
if line.startswith(":"):
|
||||
continue
|
||||
field, separator, raw_value = line.partition(":")
|
||||
value = raw_value.removeprefix(" ") if separator else ""
|
||||
if field == "event":
|
||||
event = value # rebind-ok: the last event field defines this frame's event type
|
||||
elif field == "data":
|
||||
data_lines.append(value)
|
||||
if not data_lines:
|
||||
return event, None
|
||||
try:
|
||||
return event, json.loads("\n".join(data_lines))
|
||||
except json.JSONDecodeError:
|
||||
return event, None
|
||||
|
||||
def feed(self, chunk: bytes) -> tuple[bytes, ...]:
|
||||
self._buffer += chunk
|
||||
frames: Final[list[bytes]] = [] # mutable-ok: bounded to frames in the current transport chunk
|
||||
while (boundary := self._frame_boundary(self._buffer)) is not None:
|
||||
position, separator_length = boundary
|
||||
frame_end = position + separator_length
|
||||
frames.append(self._buffer[:frame_end])
|
||||
self._buffer = self._buffer[frame_end:]
|
||||
return tuple(frames)
|
||||
|
||||
def flush(self) -> bytes:
|
||||
remainder: Final = self._buffer
|
||||
self._buffer = b""
|
||||
return remainder
|
||||
|
||||
def flush_frame(self) -> bytes | None:
|
||||
remainder: Final = self.flush()
|
||||
if not remainder or not remainder.splitlines():
|
||||
return None
|
||||
return remainder
|
||||
|
||||
|
||||
def _is_message_stop_chunk(chunk: object) -> bool:
|
||||
if isinstance(chunk, dict):
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Coroutine
|
||||
from collections.abc import Coroutine, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final, Protocol
|
||||
|
||||
|
|
@ -46,6 +46,26 @@ class PassThroughStreamingHandler:
|
|||
if litellm_logging_obj.completion_start_time is None:
|
||||
litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now())
|
||||
|
||||
@staticmethod
|
||||
def _contains_provider_error_event(raw_bytes: Sequence[bytes]) -> bool:
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
AnthropicMessagesSSEDecoder,
|
||||
)
|
||||
|
||||
decoder: Final = AnthropicMessagesSSEDecoder()
|
||||
for chunk in raw_bytes:
|
||||
for frame in decoder.feed(chunk):
|
||||
event, payload = decoder.event_data(frame)
|
||||
if event == "error" or (isinstance(payload, dict) and payload.get("type") == "error"):
|
||||
return True
|
||||
trailing_frame: Final = decoder.flush_frame()
|
||||
if trailing_frame is None:
|
||||
return False
|
||||
trailing_event, trailing_payload = decoder.event_data(trailing_frame)
|
||||
return trailing_event == "error" or (
|
||||
isinstance(trailing_payload, dict) and trailing_payload.get("type") == "error"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def chunk_processor(
|
||||
response: httpx.Response,
|
||||
|
|
@ -62,6 +82,7 @@ class PassThroughStreamingHandler:
|
|||
)
|
||||
raw_bytes: Final[list[bytes]] = []
|
||||
logging_scheduled = False
|
||||
stream_failed = False # rebind-ok: the finally block suppresses success logging after transport failure
|
||||
model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
|
||||
request_body=request_body,
|
||||
url_route=url_route,
|
||||
|
|
@ -111,6 +132,7 @@ class PassThroughStreamingHandler:
|
|||
if pending:
|
||||
yield pending
|
||||
except Exception as e:
|
||||
stream_failed = True # rebind-ok: the stream ended with a transport failure
|
||||
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
|
||||
raise
|
||||
finally:
|
||||
|
|
@ -121,7 +143,13 @@ class PassThroughStreamingHandler:
|
|||
# the caller before this generator starts (see
|
||||
# _log_passthrough_upstream_failure); logging them again here as
|
||||
# a success would double-log the same request.
|
||||
if not logging_scheduled and raw_bytes and response.status_code < 400:
|
||||
if (
|
||||
not logging_scheduled
|
||||
and raw_bytes
|
||||
and response.status_code < 400
|
||||
and not stream_failed
|
||||
and not PassThroughStreamingHandler._contains_provider_error_event(raw_bytes)
|
||||
):
|
||||
logging_scheduled = True
|
||||
try:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import time
|
|||
import traceback
|
||||
import weakref
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, Generator, Mapping, Sequence
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
|
||||
|
|
@ -4776,6 +4776,305 @@ class Router:
|
|||
)
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _anthropic_messages_stream_error(
|
||||
frame: bytes,
|
||||
model: str,
|
||||
llm_provider: str,
|
||||
) -> Exception | None:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
ANTHROPIC_STREAM_ERROR_STATUS_CODE_MAP,
|
||||
AnthropicMessagesSSEDecoder,
|
||||
)
|
||||
|
||||
event, payload = AnthropicMessagesSSEDecoder.event_data(frame)
|
||||
if event != "error" and not (isinstance(payload, dict) and payload.get("type") == "error"):
|
||||
return None
|
||||
|
||||
error: Final = payload.get("error") if isinstance(payload, dict) else None
|
||||
error_type: Final = error.get("type") if isinstance(error, dict) else None
|
||||
raw_message: Final = error.get("message") if isinstance(error, dict) else None
|
||||
status_code: Final = ANTHROPIC_STREAM_ERROR_STATUS_CODE_MAP.get(
|
||||
error_type if isinstance(error_type, str) else "",
|
||||
500,
|
||||
)
|
||||
mapped_exception: Final = litellm.APIError(
|
||||
status_code=status_code,
|
||||
message=str(raw_message) if raw_message is not None else "Anthropic in-stream error",
|
||||
llm_provider=llm_provider,
|
||||
model=model,
|
||||
)
|
||||
if 400 <= status_code < 500 and status_code != 429:
|
||||
return mapped_exception
|
||||
return MidStreamFallbackError(
|
||||
message=str(mapped_exception),
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
original_exception=mapped_exception,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _anthropic_messages_frame_commits_stream(frame: bytes) -> bool:
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
AnthropicMessagesSSEDecoder,
|
||||
)
|
||||
|
||||
event, payload = AnthropicMessagesSSEDecoder.event_data(frame)
|
||||
if event is None and payload is None:
|
||||
return False
|
||||
if event in ("message_start", "ping"):
|
||||
return False
|
||||
if event != "content_block_start" or not isinstance(payload, dict):
|
||||
return True
|
||||
content_block: Final = payload.get("content_block")
|
||||
if not isinstance(content_block, dict):
|
||||
return True
|
||||
return content_block.get("type") not in ("text", "thinking") or bool(content_block.get("text"))
|
||||
|
||||
@staticmethod
|
||||
def _anthropic_messages_transport_error(
|
||||
error: Exception,
|
||||
model: str,
|
||||
llm_provider: str,
|
||||
) -> Exception:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
if isinstance(error, MidStreamFallbackError):
|
||||
return error
|
||||
try:
|
||||
mapped_exception: Exception = litellm.exception_type( # rebind-ok: exception mapper may return or raise
|
||||
model=model,
|
||||
original_exception=error,
|
||||
custom_llm_provider=llm_provider,
|
||||
completion_kwargs=MappingProxyType({}),
|
||||
extra_kwargs=MappingProxyType({}),
|
||||
)
|
||||
except Exception as mapped_error:
|
||||
mapped_exception = mapped_error # rebind-ok: exception mapper communicates most results by raising
|
||||
status_code: Final = getattr(mapped_exception, "status_code", None)
|
||||
if isinstance(status_code, int) and 400 <= status_code < 500 and status_code != 429:
|
||||
return mapped_exception
|
||||
return MidStreamFallbackError(
|
||||
message=str(mapped_exception),
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
original_exception=mapped_exception,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
|
||||
async def _aanthropic_messages_streaming_iterator(
|
||||
self,
|
||||
response: AsyncIterable[bytes],
|
||||
initial_kwargs: dict[str, object], # mutable-ok: Router fallback utilities enrich the request in place
|
||||
) -> AsyncIterator[bytes]:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
AnthropicMessagesSSEDecoder,
|
||||
AnthropicMessagesStreamHiddenParams,
|
||||
AnthropicMessagesStreamingResponse,
|
||||
aclose_if_supported,
|
||||
)
|
||||
|
||||
source_iterator: Final = response
|
||||
model: Final = str(initial_kwargs.get("model") or "")
|
||||
llm_provider: Final = str(initial_kwargs.get("custom_llm_provider") or "anthropic")
|
||||
|
||||
async def stream_with_fallbacks() -> AsyncGenerator[bytes]:
|
||||
decoder: Final = AnthropicMessagesSSEDecoder()
|
||||
buffered_frames: Final[list[bytes]] = [] # mutable-ok: holds protocol frames until content commits
|
||||
primary_frame_emitted = False # rebind-ok: tracks client-visible stream state
|
||||
fallback_response: object | None = None # rebind-ok: retained for shielded cleanup
|
||||
try:
|
||||
try:
|
||||
async for chunk in source_iterator:
|
||||
frames = decoder.feed(chunk)
|
||||
for frame in frames:
|
||||
frame_error = self._anthropic_messages_stream_error(
|
||||
frame=frame,
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
if frame_error is not None:
|
||||
raise frame_error
|
||||
if not primary_frame_emitted and not self._anthropic_messages_frame_commits_stream(frame):
|
||||
buffered_frames.append(frame)
|
||||
continue
|
||||
if not primary_frame_emitted:
|
||||
for buffered_frame in buffered_frames:
|
||||
yield buffered_frame
|
||||
buffered_frames.clear()
|
||||
primary_frame_emitted = True # rebind-ok: client-visible state is committed
|
||||
yield frame
|
||||
|
||||
remainder: Final = decoder.flush_frame()
|
||||
if remainder:
|
||||
remainder_error: Final = self._anthropic_messages_stream_error(
|
||||
frame=remainder,
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
if remainder_error is not None:
|
||||
raise remainder_error
|
||||
if not primary_frame_emitted:
|
||||
for buffered_frame in buffered_frames:
|
||||
yield buffered_frame
|
||||
primary_frame_emitted = True # rebind-ok: incomplete remainder is client-visible
|
||||
yield remainder
|
||||
elif not primary_frame_emitted:
|
||||
for buffered_frame in buffered_frames:
|
||||
yield buffered_frame
|
||||
except Exception as error:
|
||||
fallback_error: Final = self._anthropic_messages_transport_error(
|
||||
error=error,
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
if primary_frame_emitted or not isinstance(fallback_error, MidStreamFallbackError):
|
||||
original_exception: Final = getattr(fallback_error, "original_exception", None)
|
||||
if primary_frame_emitted and isinstance(original_exception, Exception):
|
||||
raise original_exception from fallback_error
|
||||
raise fallback_error
|
||||
|
||||
model_group_value: Final = initial_kwargs.get("model")
|
||||
if not isinstance(model_group_value, str):
|
||||
raise fallback_error
|
||||
model_group: Final = model_group_value
|
||||
initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper
|
||||
self._update_kwargs_before_fallbacks(
|
||||
model=model_group,
|
||||
kwargs=initial_kwargs,
|
||||
metadata_variable_name="litellm_metadata",
|
||||
)
|
||||
with anyio.CancelScope(shield=True):
|
||||
try:
|
||||
await aclose_if_supported(source_iterator)
|
||||
except BaseException as close_error:
|
||||
verbose_router_logger.debug(
|
||||
"stream_with_fallbacks(anthropic_messages): error closing source before fallback: %s",
|
||||
close_error,
|
||||
)
|
||||
try:
|
||||
fallback_response = cast( # rebind-ok: legacy result retained for cleanup # cast-ok: legacy return
|
||||
object,
|
||||
await self.async_function_with_fallbacks_common_utils( # rebind-ok: cleanup owns it
|
||||
e=fallback_error,
|
||||
disable_fallbacks=False,
|
||||
fallbacks=cast( # cast-ok: public Router fallback config is normalized to list-or-none
|
||||
list[object] | None,
|
||||
initial_kwargs.get("fallbacks", self.fallbacks),
|
||||
),
|
||||
context_window_fallbacks=cast( # cast-ok: Router config is normalized to list-or-none
|
||||
list[object] | None,
|
||||
initial_kwargs.get("context_window_fallbacks", self.context_window_fallbacks),
|
||||
),
|
||||
content_policy_fallbacks=cast( # cast-ok: Router config is normalized to list-or-none
|
||||
list[object] | None,
|
||||
initial_kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks),
|
||||
),
|
||||
model_group=model_group,
|
||||
args=(),
|
||||
kwargs=initial_kwargs,
|
||||
include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True,
|
||||
),
|
||||
)
|
||||
if isinstance(fallback_response, AsyncIterable):
|
||||
fallback_stream: Final = cast( # cast-ok: runtime AsyncIterable check narrows the protocol
|
||||
"AsyncIterable[bytes]", fallback_response
|
||||
)
|
||||
async for fallback_item in fallback_stream:
|
||||
yield fallback_item
|
||||
else:
|
||||
if not isinstance(fallback_response, bytes):
|
||||
raise TypeError("Anthropic streaming fallback must return an async iterator or bytes")
|
||||
yield fallback_response
|
||||
except MidStreamFallbackError as unresolved_error:
|
||||
if unresolved_error.original_exception is not None:
|
||||
raise unresolved_error.original_exception from unresolved_error
|
||||
raise
|
||||
finally:
|
||||
with anyio.CancelScope(shield=True):
|
||||
try:
|
||||
await aclose_if_supported(source_iterator)
|
||||
except BaseException as close_error:
|
||||
verbose_router_logger.debug(
|
||||
"stream_with_fallbacks(anthropic_messages): error closing source: %s",
|
||||
close_error,
|
||||
)
|
||||
if fallback_response is not None:
|
||||
try:
|
||||
await aclose_if_supported(fallback_response)
|
||||
except BaseException as close_error:
|
||||
verbose_router_logger.debug(
|
||||
"stream_with_fallbacks(anthropic_messages): error closing fallback: %s",
|
||||
close_error,
|
||||
)
|
||||
|
||||
source_hidden_params_dict: Final = get_hidden_params_dict(source_iterator)
|
||||
source_headers: Final = source_hidden_params_dict.get("additional_headers")
|
||||
additional_headers: Final[dict[str, str]] = ( # mutable-ok: response metadata is mutable by contract
|
||||
cast( # cast-ok: runtime dict check plus copy isolates provider response headers
|
||||
"dict[str, str]", source_headers.copy()
|
||||
)
|
||||
if isinstance(source_headers, dict)
|
||||
else {} # mutable-ok: response metadata is mutable by contract
|
||||
)
|
||||
hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers=additional_headers)
|
||||
wrapped: Final = AnthropicMessagesStreamingResponse(
|
||||
completion_stream=stream_with_fallbacks(),
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
wrapped_hidden_params: Final[dict[str, object]] = { # mutable-ok: response metadata is mutable by contract
|
||||
**source_hidden_params_dict,
|
||||
"additional_headers": additional_headers,
|
||||
}
|
||||
cast( # cast-ok: response wrapper owns this mutable metadata mapping
|
||||
_HiddenParamsHost, wrapped
|
||||
)._hidden_params = ( # pyright: ignore[reportPrivateUsage] # response owns its metadata mapping
|
||||
wrapped_hidden_params
|
||||
)
|
||||
return wrapped
|
||||
|
||||
async def _aanthropic_messages_with_streaming_fallbacks(
|
||||
self,
|
||||
original_function: Callable[..., object],
|
||||
**kwargs: object, # kwargs-ok: provider-specific fields are validated by the downstream API
|
||||
) -> object:
|
||||
"""Add consumption-time fallbacks to streaming Anthropic Messages calls."""
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import safe_deep_copy
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: fallback utilities enrich this snapshot
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
if isinstance(fallback_kwargs.get("metadata"), dict):
|
||||
fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"])
|
||||
fallback_kwargs["original_generic_function"] = original_function
|
||||
|
||||
model: Final = kwargs.get("model")
|
||||
if not isinstance(model, str):
|
||||
raise TypeError("Anthropic messages routing requires a string model")
|
||||
request_kwargs: Final[dict[str, object]] = { # mutable-ok: generic Router accepts provider-specific fields
|
||||
key: value for key, value in kwargs.items() if key != "model"
|
||||
}
|
||||
response: Final = cast( # cast-ok: generic Router has an untyped legacy return contract
|
||||
object,
|
||||
await self._ageneric_api_call_with_fallbacks(
|
||||
model=model,
|
||||
original_function=original_function,
|
||||
**request_kwargs,
|
||||
),
|
||||
)
|
||||
if kwargs.get("stream") and isinstance(response, AsyncIterable):
|
||||
return await self._aanthropic_messages_streaming_iterator(
|
||||
response=cast( # cast-ok: runtime AsyncIterable check narrows the streaming bytes protocol
|
||||
"AsyncIterable[bytes]", response
|
||||
),
|
||||
initial_kwargs=fallback_kwargs,
|
||||
)
|
||||
return response
|
||||
|
||||
def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
|
||||
"""
|
||||
Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router
|
||||
|
|
@ -5917,6 +6216,11 @@ class Router:
|
|||
original_function=original_function,
|
||||
**kwargs,
|
||||
)
|
||||
elif call_type == "anthropic_messages":
|
||||
return await self._aanthropic_messages_with_streaming_fallbacks(
|
||||
original_function=original_function,
|
||||
**kwargs,
|
||||
)
|
||||
elif call_type in (
|
||||
"acreate_realtime_client_secret",
|
||||
"arealtime_calls",
|
||||
|
|
@ -5928,7 +6232,6 @@ class Router:
|
|||
**kwargs,
|
||||
)
|
||||
elif call_type in (
|
||||
"anthropic_messages",
|
||||
"_arealtime",
|
||||
"_aresponses_websocket",
|
||||
"acreate_fine_tuning_job",
|
||||
|
|
|
|||
|
|
@ -130,6 +130,98 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er
|
|||
mock_route.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_does_not_log_provider_error_event_as_success():
|
||||
chunks = [
|
||||
b'event: message_start\ndata: {"type":"message_start"}\n\nevent:err',
|
||||
b'or\ndata:{"type":"error","error":{"type":"overloaded_error"}}\n\n',
|
||||
]
|
||||
response = _make_streaming_response(chunks)
|
||||
|
||||
with patch.object(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
new=AsyncMock(),
|
||||
) as mock_route:
|
||||
received = [
|
||||
chunk
|
||||
async for chunk in PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={"model": "claude-3-haiku"},
|
||||
litellm_logging_obj=MagicMock(),
|
||||
endpoint_type=EndpointType.ANTHROPIC,
|
||||
start_time=datetime.now(),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
]
|
||||
|
||||
assert received == chunks
|
||||
mock_route.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_logs_successful_anthropic_stream():
|
||||
chunks = [
|
||||
b'event: message_start\ndata: {"type":"message_start"}\n\n',
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
]
|
||||
response = _make_streaming_response(chunks)
|
||||
|
||||
with patch.object(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
new=AsyncMock(),
|
||||
) as mock_route:
|
||||
received = [
|
||||
chunk
|
||||
async for chunk in PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={"model": "claude-3-haiku"},
|
||||
litellm_logging_obj=MagicMock(),
|
||||
endpoint_type=EndpointType.ANTHROPIC,
|
||||
start_time=datetime.now(),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
]
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert received == chunks
|
||||
mock_route.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_does_not_log_transport_error_as_success():
|
||||
response = _make_streaming_response([b"partial"])
|
||||
|
||||
async def broken_stream():
|
||||
yield b"partial"
|
||||
raise httpx.ReadError("connection reset")
|
||||
|
||||
response.aiter_bytes = broken_stream
|
||||
with patch.object(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
new=AsyncMock(),
|
||||
) as mock_route:
|
||||
with pytest.raises(httpx.ReadError):
|
||||
_ = [
|
||||
chunk
|
||||
async for chunk in PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={"model": "claude-3-haiku"},
|
||||
litellm_logging_obj=MagicMock(),
|
||||
endpoint_type=EndpointType.ANTHROPIC,
|
||||
start_time=datetime.now(),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
]
|
||||
|
||||
mock_route.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
|
||||
response = _make_streaming_response([])
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import os
|
|||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
|
|
@ -2577,6 +2578,30 @@ class _AsyncList:
|
|||
return item
|
||||
|
||||
|
||||
class _AnthropicMessagesStream(_AsyncList):
|
||||
def __init__(self, items=(), error=None, hidden_params=None):
|
||||
super().__init__(items)
|
||||
self._error = error
|
||||
self._hidden_params = hidden_params or {}
|
||||
self.closed = False
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return await super().__anext__()
|
||||
except StopAsyncIteration:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
raise
|
||||
|
||||
async def aclose(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _UncopyableMetadata:
|
||||
def __deepcopy__(self, memo):
|
||||
raise TypeError("not copyable")
|
||||
|
||||
|
||||
def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"):
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -2593,6 +2618,269 @@ def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_falls_back_on_fragmented_sse_error():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[
|
||||
b'event: message_start\r\ndata: {"type":"message_start"}\r\n\r\nevent: error\r\n',
|
||||
b'data: {"type":"error","error":{"type":"overloaded_error","message":"busy"}}\r\n\r\n',
|
||||
],
|
||||
hidden_params={
|
||||
"model_id": "primary",
|
||||
"additional_headers": {"x-request-id": "one"},
|
||||
},
|
||||
)
|
||||
fallback_chunks = [
|
||||
b'event: message_start\ndata: {"type":"message_start"}\n\n',
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
]
|
||||
fallback = _AnthropicMessagesStream(fallback_chunks)
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
side_effect=lambda **kwargs: fallback if primary.closed else pytest.fail("source stream remained open"),
|
||||
) as mock_fallback:
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == fallback_chunks
|
||||
assert wrapped._hidden_params["model_id"] == "primary"
|
||||
assert primary.closed is True
|
||||
assert fallback.closed is True
|
||||
call = mock_fallback.call_args.kwargs
|
||||
assert call["e"].status_code == 529
|
||||
assert (
|
||||
call["kwargs"]["original_function"]
|
||||
== router._ageneric_api_call_with_fallbacks_helper
|
||||
)
|
||||
assert (
|
||||
call["kwargs"]["original_generic_function"]
|
||||
is litellm.anthropic_messages
|
||||
)
|
||||
assert call["kwargs"]["litellm_metadata"]["model_group"] == "anthropic/claude"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_falls_back_on_transport_error():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
request = httpx.Request("POST", "https://example.com/v1/messages")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[b'event: message_start\ndata: {"type":"message_start"}\n\n'],
|
||||
error=httpx.ReadError("connection reset", request=request),
|
||||
)
|
||||
fallback_chunks = [
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
|
||||
]
|
||||
fallback = _AnthropicMessagesStream(fallback_chunks)
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
return_value=fallback,
|
||||
) as mock_fallback:
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == fallback_chunks
|
||||
assert primary.closed is True
|
||||
assert fallback.closed is True
|
||||
assert mock_fallback.call_args.kwargs["e"].is_pre_first_chunk is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_comment_keepalive_does_not_block_fallback():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[
|
||||
b": keepalive\n\n",
|
||||
b'event: error\ndata: {"type":"error","error":'
|
||||
b'{"type":"rate_limit_error","message":"slow down"}}\n\n',
|
||||
]
|
||||
)
|
||||
fallback_chunks = [b'event: message_stop\ndata: {"type":"message_stop"}\n\n']
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
return_value=_AnthropicMessagesStream(fallback_chunks),
|
||||
):
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == fallback_chunks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_unterminated_error_frame_falls_back():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[
|
||||
b'event: error\ndata: {"type":"error","error":'
|
||||
b'{"type":"overloaded_error","message":"busy"}}'
|
||||
]
|
||||
)
|
||||
fallback_chunks = [b'event: message_stop\ndata: {"type":"message_stop"}\n\n']
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
return_value=_AnthropicMessagesStream(fallback_chunks),
|
||||
):
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == fallback_chunks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_non_retryable_sse_error_is_not_fallbacked():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[
|
||||
b'event: error\ndata: {"type":"error","error":'
|
||||
b'{"type":"invalid_request_error","message":"bad input"}}\n\n'
|
||||
]
|
||||
)
|
||||
|
||||
with patch.object(router, "async_function_with_fallbacks_common_utils") as mock_fallback:
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
with pytest.raises(litellm.APIError) as exc_info:
|
||||
_ = [chunk async for chunk in wrapped]
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
mock_fallback.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_unresolved_fallback_surfaces_provider_error():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream(
|
||||
[
|
||||
b'event: error\ndata: {"type":"error","error":'
|
||||
b'{"type":"internal_server_error","message":"failed"}}\n\n'
|
||||
]
|
||||
)
|
||||
|
||||
async def no_fallback(**kwargs):
|
||||
raise kwargs["e"]
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
side_effect=no_fallback,
|
||||
):
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
with pytest.raises(litellm.APIError) as exc_info:
|
||||
_ = [chunk async for chunk in wrapped]
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert not isinstance(exc_info.value, MidStreamFallbackError)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_does_not_fallback_after_content():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
content_frame = (
|
||||
b'event: content_block_delta\ndata: {"type":"content_block_delta",'
|
||||
b'"delta":{"type":"text_delta","text":"partial"}}\n\n'
|
||||
)
|
||||
error_frame = (
|
||||
b'event: error\ndata: {"type":"error","error":'
|
||||
b'{"type":"internal_server_error","message":"failed"}}\n\n'
|
||||
)
|
||||
primary = _AnthropicMessagesStream([content_frame, error_frame])
|
||||
|
||||
with patch.object(
|
||||
router, "async_function_with_fallbacks_common_utils"
|
||||
) as mock_fallback:
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=primary,
|
||||
initial_kwargs={
|
||||
"model": "anthropic/claude",
|
||||
"stream": True,
|
||||
"original_generic_function": litellm.anthropic_messages,
|
||||
},
|
||||
)
|
||||
iterator = wrapped.__aiter__()
|
||||
assert await iterator.__anext__() == content_frame
|
||||
with pytest.raises(litellm.APIError) as exc_info:
|
||||
await iterator.__anext__()
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert mock_fallback.called is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_factory_wraps_streaming_response():
|
||||
router = _make_router_with_fallback("anthropic/claude", "bedrock/claude")
|
||||
primary = _AnthropicMessagesStream()
|
||||
span = _UncopyableMetadata()
|
||||
original_metadata = {"trace_id": "test", "span": span}
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
return_value=primary,
|
||||
):
|
||||
factory = router.factory_function(
|
||||
litellm.anthropic_messages,
|
||||
call_type="anthropic_messages",
|
||||
)
|
||||
response = await factory(
|
||||
model="anthropic/claude",
|
||||
stream=True,
|
||||
litellm_metadata=original_metadata,
|
||||
)
|
||||
|
||||
assert response is not primary
|
||||
assert hasattr(response, "__aiter__")
|
||||
assert original_metadata == {"trace_id": "test", "span": span}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_streaming_iterator_fallback():
|
||||
"""Catches MidStreamFallbackError, re-enters the fallback chain via
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue