mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
fix(router): stream /v1/messages lifecycle frames live when no fallback can take over (#43600)
* fix(router): stream anthropic messages lifecycle frames live when no fallback can take over The /v1/messages streaming wrapper buffered message_start and content_block_start until the first content_block_delta and dropped pings behind buffered frames unconditionally, even for requests no fallback could ever recover. With adaptive thinking on Bedrock or Vertex the client saw no bytes for the whole thinking pass and hit read timeouts. Buffering now applies only while a fallback can still take over (generic or refusal chain resolving), and a ping is always forwarded live since it carries no lifecycle and keeps the connection alive. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): mirror every dispatcher fallback path in the anthropic stream gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): skip already-tried order levels in the anthropic stream gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): keep a transport-split ping behind buffered lifecycle frames instead of forwarding its head live Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(router): credit the #39566 branch this fix supersedes Co-authored-by: Radu Swigler <radu.porumba@gmail.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Radu Swigler <radu.porumba@gmail.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
66db132627
commit
d46304900f
5 changed files with 496 additions and 88 deletions
|
|
@ -40,22 +40,31 @@ def _is_message_stop_chunk(chunk: object) -> bool:
|
|||
|
||||
def is_anthropic_ping_chunk(chunk: object) -> bool:
|
||||
"""
|
||||
Whether a chunk is a pure ``ping`` keepalive frame. It carries no content
|
||||
and can recur indefinitely on a slow-starting or idle connection, so a
|
||||
mid-stream fallback wrapper drops it outright while still deciding
|
||||
whether to commit to the primary stream, rather than buffering it.
|
||||
Whether a chunk is made only of whole ``ping`` keepalive frames. A ping
|
||||
carries no content or lifecycle, so a mid-stream fallback wrapper can
|
||||
forward it live while still deciding whether to commit to the primary
|
||||
stream, without risking two overlapping message lifecycles on the wire.
|
||||
|
||||
A physical transport chunk that coalesces a ping with any other SSE
|
||||
event (``message_start``, ``content_block_delta``, ``event: error``, ...)
|
||||
is NOT a pure ping - dropping it whole would discard those events - so
|
||||
only a chunk whose every ``event:`` line is ``event: ping`` qualifies.
|
||||
is NOT a pure ping, and neither is a fragment of a ping frame split
|
||||
across two reads, or a chunk that opens with the tail of an earlier
|
||||
frame: forwarding either live would interleave it with frames still
|
||||
held back for a fallback. Only a chunk that begins with ``event: ping``,
|
||||
ends on a frame boundary, and whose every ``event:`` line is
|
||||
``event: ping`` qualifies.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
return chunk.get("type") == "ping"
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:"))
|
||||
return bool(event_lines) and all(line == b"event: ping" for line in event_lines)
|
||||
return False
|
||||
if not isinstance(chunk, (bytes, bytearray)):
|
||||
return False
|
||||
event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:"))
|
||||
return (
|
||||
bool(event_lines)
|
||||
and all(line == b"event: ping" for line in event_lines)
|
||||
and chunk.startswith(b"event: ping")
|
||||
and chunk.endswith((b"\n\n", b"\r\n\r\n"))
|
||||
)
|
||||
|
||||
|
||||
def is_anthropic_content_delta_chunk(chunk: object) -> bool:
|
||||
|
|
|
|||
|
|
@ -518,31 +518,20 @@ def _with_router_resolved_session_model(session: object, model_name: str) -> Map
|
|||
|
||||
|
||||
# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks
|
||||
# until real content commits the primary stream; a hostile or slow-starting
|
||||
# upstream that never emits content or an error could otherwise grow that
|
||||
# buffer without bound, so hitting this cap forces an early commit instead.
|
||||
# until real content commits the primary stream, and only while a fallback
|
||||
# can still take over; a hostile or slow-starting upstream that never emits
|
||||
# content or an error could otherwise grow that buffer without bound, so
|
||||
# hitting this cap forces an early commit instead.
|
||||
MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200
|
||||
|
||||
|
||||
def _anthropic_stream_should_drop_pre_content_ping(chunk: object, has_generated_content: bool) -> bool:
|
||||
"""A `ping` keepalive seen before any real content is dropped outright - it recurs indefinitely on a
|
||||
slow-starting connection and carries nothing worth buffering toward a possible fallback."""
|
||||
def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool:
|
||||
"""A `ping` keepalive reaches the client live whenever the stream has not committed: it carries no
|
||||
lifecycle, so it cannot create overlapping lifecycles on the wire, and it keeps the connection alive
|
||||
while lifecycle frames sit buffered for a possible fallback during a long thinking pass."""
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import is_anthropic_ping_chunk
|
||||
|
||||
if has_generated_content:
|
||||
return False
|
||||
return is_anthropic_ping_chunk(chunk)
|
||||
|
||||
|
||||
def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
|
||||
"""A `ping` that no lifecycle frame precedes reaches the client live: a fallback's own message_start can still
|
||||
follow it without overlapping lifecycles, and AgenticAnthropicStreamingIterator's hold-back keepalive is exactly
|
||||
such a ping."""
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import is_anthropic_ping_chunk
|
||||
|
||||
if has_generated_content or buffered_chunk_count:
|
||||
return False
|
||||
return is_anthropic_ping_chunk(chunk)
|
||||
return not has_generated_content and is_anthropic_ping_chunk(chunk)
|
||||
|
||||
|
||||
def _is_retriable_anthropic_status(status_code: int) -> bool:
|
||||
|
|
@ -5457,14 +5446,19 @@ class Router:
|
|||
|
||||
Lifecycle/bookkeeping frames (message_start, content_block_start,
|
||||
ping, ...) do not by themselves disqualify a fallback attempt -
|
||||
Anthropic routinely sends message_start before an overload error -
|
||||
but they are BUFFERED rather than forwarded immediately, since
|
||||
forwarding one and then appending a fallback attempt's own
|
||||
message_start would produce two overlapping message lifecycles on
|
||||
one SSE stream. Buffered frames are flushed, in order, the moment
|
||||
real content arrives (the primary attempt has committed by then
|
||||
anyway) or once the stream ends without ever producing content or
|
||||
an error.
|
||||
Anthropic routinely sends message_start before an overload error.
|
||||
When a fallback can still take over they are BUFFERED rather than
|
||||
forwarded immediately, since forwarding one and then appending a
|
||||
fallback attempt's own message_start would produce two overlapping
|
||||
message lifecycles on one SSE stream; a `ping` carries no lifecycle,
|
||||
so it is forwarded live even while lifecycle frames sit buffered,
|
||||
keeping the connection alive during a long thinking pass. Buffered
|
||||
frames are flushed, in order, the moment real content arrives (the
|
||||
primary attempt has committed by then anyway) or once the stream
|
||||
ends without ever producing content or an error. When no fallback
|
||||
can take over the request is already committed, so every frame,
|
||||
including pings and provider error frames, is forwarded live and
|
||||
verbatim instead.
|
||||
"""
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
|
|
@ -5481,34 +5475,33 @@ class Router:
|
|||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
# Lifecycle/bookkeeping frames (message_start, content_block_start,
|
||||
# ping, ...) are held back rather than forwarded immediately:
|
||||
# Anthropic routinely sends message_start before an overload
|
||||
# error, and once a byte reaches the client a fallback attempt
|
||||
# can only append its OWN message_start, producing two
|
||||
# overlapping message lifecycles on one SSE stream. Buffered
|
||||
# frames are flushed the moment real content (content_block_delta)
|
||||
# ...) are held back rather than forwarded immediately, but only
|
||||
# while a fallback can still take over: Anthropic routinely sends
|
||||
# message_start before an overload error, and once a byte reaches
|
||||
# the client a fallback attempt can only append its OWN
|
||||
# message_start, producing two overlapping message lifecycles on
|
||||
# one SSE stream. A `ping` keepalive carries no lifecycle, so it
|
||||
# is forwarded live even behind buffered frames, keeping the
|
||||
# connection alive through a long thinking pass. Buffered frames
|
||||
# are flushed the moment real content (content_block_delta)
|
||||
# arrives - at that point the primary attempt has committed and a
|
||||
# clean retry is no longer possible anyway - or once the primary
|
||||
# stream ends without ever producing content. A `ping` keepalive
|
||||
# that nothing precedes is forwarded live (it is how a hold-back
|
||||
# turn keeps its connection alive); one behind buffered frames is
|
||||
# dropped outright rather than buffered, since it can recur
|
||||
# indefinitely on a slow-starting connection and carries nothing
|
||||
# worth preserving; hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS
|
||||
# forces the same early commit as real content arriving, so a
|
||||
# hostile or pathological upstream can't grow the buffer forever.
|
||||
has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit
|
||||
buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline
|
||||
# stream ends without ever producing content. Hitting
|
||||
# MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS forces the same early
|
||||
# commit as real content arriving, so a hostile or pathological
|
||||
# upstream can't grow the buffer forever. With no fallback able
|
||||
# to take over there is nothing to buffer for, so every frame,
|
||||
# including pings and provider error frames, is forwarded live.
|
||||
model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
|
||||
has_generated_content = not self._anthropic_messages_stream_can_fall_back( # rebind-ok: set once real content is seen, the buffer cap is hit, or no fallback can take over
|
||||
model, initial_kwargs
|
||||
)
|
||||
buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline
|
||||
try:
|
||||
async for chunk in source_iterator:
|
||||
if _anthropic_stream_forwards_ping_live(
|
||||
chunk, has_generated_content, len(buffered_lifecycle_chunks)
|
||||
):
|
||||
if _anthropic_stream_forwards_ping_live(chunk, has_generated_content):
|
||||
yield chunk
|
||||
continue
|
||||
if _anthropic_stream_should_drop_pre_content_ping(chunk, has_generated_content):
|
||||
continue
|
||||
if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
|
||||
has_generated_content = True
|
||||
# A transport can split one SSE data line across byte chunks, so pre-content
|
||||
|
|
@ -8447,6 +8440,56 @@ class Router:
|
|||
)
|
||||
return has_unattempted_fallback_target(resolved, kwargs)
|
||||
|
||||
def _anthropic_messages_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]:
|
||||
"""
|
||||
The distinct deployment order levels the fallback dispatcher would see for this request,
|
||||
computed the same way: the tier a pre-routing hook selected wins over the requested group.
|
||||
"""
|
||||
request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id")
|
||||
order_model_group: Final = get_pre_routing_selection(kwargs) or model_group
|
||||
all_deployments: Final = self.get_model_list(model_name=order_model_group, team_id=request_team_id) or ()
|
||||
return tuple(
|
||||
sorted(
|
||||
{
|
||||
litellm.utils._get_deployment_order(d)
|
||||
for d in all_deployments
|
||||
if litellm.utils._get_deployment_order(d) is not None
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
def _anthropic_messages_stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
|
||||
"""
|
||||
Whether async_function_with_fallbacks_common_utils could still route a
|
||||
MidStreamFallbackError somewhere for this request (order levels, weighted
|
||||
failover, content-policy or generic fallbacks), which is the only case where
|
||||
holding lifecycle frames back from the client buys a clean retry. Errs toward
|
||||
True whenever a dispatcher path might reach a fallback.
|
||||
"""
|
||||
if fallbacks_disabled_for_request(kwargs):
|
||||
return False
|
||||
if self.enable_weighted_failover:
|
||||
return True
|
||||
order_levels: Final = self._anthropic_messages_order_levels(model_group, kwargs)
|
||||
if len(order_levels) > 1:
|
||||
current_target: Final = kwargs.get("_target_order")
|
||||
skip_up_to: Final = current_target if current_target is not None else order_levels[0]
|
||||
if any(o > skip_up_to for o in order_levels):
|
||||
return True
|
||||
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
|
||||
if content_policy_fallbacks is not None and self._has_content_policy_fallback(model_group, kwargs):
|
||||
return True
|
||||
fallbacks: Final = kwargs.get("fallbacks", self.fallbacks)
|
||||
if not fallbacks:
|
||||
return False
|
||||
if _check_non_standard_fallback_format(fallbacks=fallbacks):
|
||||
return True
|
||||
resolved, _ = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks,
|
||||
lookup_groups=fallback_lookup_groups(kwargs, model_group),
|
||||
)
|
||||
return has_unattempted_fallback_target(resolved, kwargs)
|
||||
|
||||
def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
|
||||
"""
|
||||
Determines if a content policy error should be raised.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,82 @@
|
|||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from typing import Final
|
||||
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
_API_KEY: Final = "synthetic-anthropic-key"
|
||||
|
||||
|
||||
def _sse(event: str, payload: dict[str, object]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def test_messages_stream_message_start_reaches_client_before_content_without_fallback(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
"""With no fallback able to take over, the proxy must not hold lifecycle
|
||||
frames back for a retry that cannot happen: message_start reaches the
|
||||
client while the upstream is still thinking."""
|
||||
gate: Final = threading.Event()
|
||||
head: Final = _sse("message_start", {"type": "message_start", "message": {"id": "msg_live_1"}}) + _sse(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
)
|
||||
tail: Final = (
|
||||
_sse(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}},
|
||||
)
|
||||
+ _sse("content_block_stop", {"type": "content_block_stop", "index": 0})
|
||||
+ _sse(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 3}},
|
||||
)
|
||||
+ _sse("message_stop", {"type": "message_stop"})
|
||||
)
|
||||
prompt: Final = "live-lifecycle-" + uuid.uuid4().hex
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/v1/messages"
|
||||
assert request.headers["x-api-key"] == _API_KEY
|
||||
body: Final = json.loads(request.body)
|
||||
assert body["model"] == _MODEL
|
||||
assert body["stream"] is True
|
||||
assert body["messages"] == [{"role": "user", "content": prompt}]
|
||||
return Reply(content_type="text/event-stream", chunks=(head, tail), gate_after_first=gate)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
with gateway.client.stream(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
json={
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
},
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
) as response:
|
||||
assert response.status_code == 200, response.read().decode()
|
||||
lines = response.iter_lines()
|
||||
first_event: Final = next(
|
||||
json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ")
|
||||
)
|
||||
assert first_event["type"] == "message_start"
|
||||
gate.set()
|
||||
events: Final = (first_event,) + tuple(
|
||||
json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ")
|
||||
)
|
||||
assert tuple(event["type"] for event in events) == (
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
), f"observed events: {events!r}"
|
||||
assert [request.target for request in wire.drain()] == ["/v1/messages"]
|
||||
|
|
@ -19,6 +19,7 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
|
|||
_is_provider_error_chunk,
|
||||
anthropic_messages_response_as_sse_events,
|
||||
is_anthropic_content_delta_chunk,
|
||||
is_anthropic_ping_chunk,
|
||||
parse_anthropic_error_event,
|
||||
)
|
||||
|
||||
|
|
@ -171,6 +172,25 @@ def test_is_message_stop_chunk():
|
|||
assert _is_message_stop_chunk("message_stop") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("chunk", "expected"),
|
||||
[
|
||||
(b'event: ping\ndata: {"type": "ping"}\n\n', True),
|
||||
(b'event: ping\r\ndata: {"type": "ping"}\r\n\r\n', True),
|
||||
(b'event: ping\ndata: {"type": "ping"}\n\nevent: ping\ndata: {"type": "ping"}\n\n', True),
|
||||
({"type": "ping"}, True),
|
||||
(b'event: ping\ndata: {"ty', False),
|
||||
(b'pe": "ping"}\n\n', False),
|
||||
(b'pe": "message_start"}}\n\nevent: ping\ndata: {"type": "ping"}\n\n', False),
|
||||
(b'event: ping\ndata: {"type": "ping"}\n\nevent: content_block_delta\ndata: {}\n\n', False),
|
||||
({"type": "message_start"}, False),
|
||||
("event: ping", False),
|
||||
],
|
||||
)
|
||||
def test_is_anthropic_ping_chunk_only_matches_whole_ping_frames(chunk: object, expected: bool):
|
||||
assert is_anthropic_ping_chunk(chunk) is expected, chunk
|
||||
|
||||
|
||||
def test_is_message_stop_chunk_ignores_substring_in_payload():
|
||||
"""
|
||||
Regression: a `content_block_delta` frame whose payload happens to contain
|
||||
|
|
|
|||
|
|
@ -45,7 +45,6 @@ from litellm.router import (
|
|||
_anthropic_stream_forwards_ping_live,
|
||||
_anthropic_stream_raised_error_status,
|
||||
_anthropic_stream_should_decline_fallback,
|
||||
_anthropic_stream_should_drop_pre_content_ping,
|
||||
_is_retriable_anthropic_status,
|
||||
_responses_stream_holds_event,
|
||||
)
|
||||
|
|
@ -4170,7 +4169,7 @@ def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"):
|
|||
|
||||
class _InjectedFallbackRouter(Router):
|
||||
def __init__(self, fallback_response: object) -> None:
|
||||
super().__init__(model_list=[])
|
||||
super().__init__(model_list=[], fallbacks=[{"primary": ["fallback"]}])
|
||||
self._fallback_response: Final = fallback_response
|
||||
|
||||
async def async_function_with_fallbacks_common_utils(
|
||||
|
|
@ -13696,7 +13695,8 @@ def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream:
|
|||
return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object())
|
||||
|
||||
|
||||
def _anthropic_messages_make_router() -> Router:
|
||||
def _anthropic_messages_make_router(**router_kwargs) -> Router:
|
||||
router_kwargs.setdefault("fallbacks", [{"primary": ["fallback"]}])
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -13712,7 +13712,8 @@ def _anthropic_messages_make_router() -> Router:
|
|||
"model": "bedrock/anthropic.claude-sonnet-4-5",
|
||||
},
|
||||
},
|
||||
]
|
||||
],
|
||||
**router_kwargs,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -13900,24 +13901,286 @@ async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_c
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_dropped():
|
||||
"""Bugbot regression: a `ping` keepalive behind buffered lifecycle frames
|
||||
carries no content and is dropped outright rather than buffered -
|
||||
otherwise a slow-starting connection sending many pings could grow the
|
||||
pre-content buffer without bound."""
|
||||
router = _anthropic_messages_make_router()
|
||||
async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwarded_live():
|
||||
"""A `ping` behind buffered lifecycle frames still reaches the client
|
||||
live: it carries no lifecycle, so it cannot create overlapping
|
||||
lifecycles, and it keeps the connection alive while a fallback-able
|
||||
stream holds message_start back through a long thinking pass."""
|
||||
router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}])
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
yield _anthropic_messages_ping_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"})
|
||||
|
||||
assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk()
|
||||
content_released.set()
|
||||
assert [chunk async for chunk in wrapped] == [
|
||||
_anthropic_messages_message_start_chunk(),
|
||||
_anthropic_messages_content_chunk("hi"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_split_ping_stays_in_order_behind_buffered_lifecycle_frame():
|
||||
"""A ping the transport splits across two reads is not a whole frame, so
|
||||
neither fragment may jump ahead of the buffered message_start: yielding
|
||||
the head live and flushing the tail behind message_start would splice a
|
||||
lifecycle frame into the middle of the ping on the wire."""
|
||||
router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}])
|
||||
ping_head, ping_tail = b'event: ping\ndata: {"ty', b'pe": "ping"}\n\n'
|
||||
source = _AnthropicMessagesFakeByteStream(
|
||||
[
|
||||
_anthropic_messages_message_start_chunk(),
|
||||
_anthropic_messages_ping_chunk(),
|
||||
_anthropic_messages_content_chunk("hi"),
|
||||
]
|
||||
[_anthropic_messages_message_start_chunk(), ping_head, ping_tail, _anthropic_messages_content_chunk("hi")]
|
||||
)
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"})
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi")]
|
||||
assert [chunk async for chunk in wrapped] == [
|
||||
_anthropic_messages_message_start_chunk(),
|
||||
ping_head,
|
||||
ping_tail,
|
||||
_anthropic_messages_content_chunk("hi"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_no_fallback_message_start_reaches_client_before_content():
|
||||
"""With no fallback able to take over, the stream is committed from the
|
||||
first frame: message_start reaches the client live instead of waiting
|
||||
behind the buffer for content that may be a whole thinking pass away."""
|
||||
router = _anthropic_messages_make_router(fallbacks=None)
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"})
|
||||
|
||||
assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk()
|
||||
content_released.set()
|
||||
assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_disabled_fallbacks_message_start_reaches_client_before_content():
|
||||
"""A router with fallbacks configured cannot take over a request that
|
||||
opted out with disable_fallbacks=True, so its lifecycle frames reach
|
||||
the client live exactly like a no-fallback router's."""
|
||||
router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}])
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source(), initial_kwargs={"model": "primary", "disable_fallbacks": True}
|
||||
)
|
||||
|
||||
assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk()
|
||||
content_released.set()
|
||||
assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_no_fallback_error_frame_reaches_client_verbatim():
|
||||
"""With no fallback able to take over, a retriable provider error frame
|
||||
is forwarded verbatim instead of triggering a fallback that does not
|
||||
exist, and the frames already received stay in order ahead of it."""
|
||||
router = _anthropic_messages_make_router(fallbacks=None)
|
||||
source = _AnthropicMessagesFakeByteStream(
|
||||
[_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()]
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])),
|
||||
) as mock_fallback:
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source, initial_kwargs={"model": "primary"}
|
||||
)
|
||||
collected = [chunk async for chunk in wrapped]
|
||||
|
||||
assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()]
|
||||
mock_fallback.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecycle_frames():
|
||||
"""A "*" default fallback can take over for any group, so lifecycle
|
||||
frames are still held back until real content commits the primary."""
|
||||
router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}])
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"})
|
||||
|
||||
pending = asyncio.ensure_future(wrapped.__anext__())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not pending.done()
|
||||
content_released.set()
|
||||
assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk()
|
||||
assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")]
|
||||
|
||||
|
||||
def _anthropic_messages_two_order_primary_model_list() -> list:
|
||||
return [
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5", "order": 2},
|
||||
},
|
||||
{
|
||||
"model_name": "fallback",
|
||||
"litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router_kwargs,request_kwargs,expected",
|
||||
[
|
||||
pytest.param({"fallbacks": None}, {"model": "primary"}, False, id="no-fallbacks"),
|
||||
pytest.param({"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary"}, True, id="group-fallback"),
|
||||
pytest.param({"fallbacks": [{"other": ["fallback"]}]}, {"model": "primary"}, False, id="unrelated-group"),
|
||||
pytest.param(
|
||||
{"fallbacks": [{"*": ["fallback"]}]},
|
||||
{"model": "primary", "fallbacks": None},
|
||||
False,
|
||||
id="wildcard-overridden-by-request-none",
|
||||
),
|
||||
pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"),
|
||||
pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"),
|
||||
pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"),
|
||||
pytest.param(
|
||||
{"fallbacks": [{"primary": ["fallback"]}]},
|
||||
{"model": "primary", "disable_fallbacks": True},
|
||||
False,
|
||||
id="disable-fallbacks",
|
||||
),
|
||||
pytest.param(
|
||||
{"fallbacks": None, "content_policy_fallbacks": [{"primary": ["fallback"]}]},
|
||||
{"model": "primary"},
|
||||
True,
|
||||
id="content-policy-fallback",
|
||||
),
|
||||
pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"),
|
||||
],
|
||||
)
|
||||
def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected):
|
||||
router = _anthropic_messages_make_router(**router_kwargs)
|
||||
assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"orders,expected",
|
||||
[
|
||||
pytest.param([1, 2], True, id="distinct-orders-can-fall-back"),
|
||||
pytest.param([1, 1], False, id="same-order-cannot-fall-back"),
|
||||
],
|
||||
)
|
||||
def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order},
|
||||
}
|
||||
for order in orders
|
||||
],
|
||||
fallbacks=None,
|
||||
)
|
||||
assert router._anthropic_messages_stream_can_fall_back("primary", {"model": "primary"}) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_kwargs,expected",
|
||||
[
|
||||
pytest.param({"model": "primary"}, True, id="no-target-order"),
|
||||
pytest.param({"model": "primary", "_target_order": 1}, True, id="higher-order-remains"),
|
||||
pytest.param({"model": "primary", "_target_order": 2}, False, id="top-order-no-order-fallback"),
|
||||
pytest.param(
|
||||
{"model": "primary", "_target_order": 2, "fallbacks": [{"primary": ["fallback"]}]},
|
||||
True,
|
||||
id="top-order-external-fallback",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_anthropic_messages_stream_can_fall_back_order_target(request_kwargs, expected):
|
||||
router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None)
|
||||
assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected
|
||||
|
||||
|
||||
def test_anthropic_messages_order_levels_direct_call():
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order},
|
||||
}
|
||||
for order in (2, 1, None)
|
||||
],
|
||||
fallbacks=None,
|
||||
)
|
||||
assert router._anthropic_messages_order_levels("primary", {"model": "primary"}) == (1, 2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames():
|
||||
"""Two order levels in one group are a real fallback target for the
|
||||
dispatcher, so lifecycle frames stay buffered until content commits."""
|
||||
router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None)
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"})
|
||||
|
||||
pending = asyncio.ensure_future(wrapped.__anext__())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not pending.done()
|
||||
content_released.set()
|
||||
assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk()
|
||||
assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_request_fallbacks_none_forwards_message_start_live():
|
||||
"""A per-request fallbacks=None override disables the router's wildcard
|
||||
fallback, so lifecycle frames reach the client live before content."""
|
||||
router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}])
|
||||
content_released = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
yield _anthropic_messages_message_start_chunk()
|
||||
await content_released.wait()
|
||||
yield _anthropic_messages_content_chunk("hi")
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source(), initial_kwargs={"model": "primary", "fallbacks": None}
|
||||
)
|
||||
|
||||
assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk()
|
||||
content_released.set()
|
||||
assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -14320,21 +14583,12 @@ def test_merge_fallback_hidden_params_direct_call():
|
|||
}
|
||||
|
||||
|
||||
def test_anthropic_stream_should_drop_pre_content_ping_direct_call():
|
||||
ping = _anthropic_messages_ping_chunk()
|
||||
content = _anthropic_messages_content_chunk("hi")
|
||||
assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=False) is True
|
||||
assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=True) is False
|
||||
assert _anthropic_stream_should_drop_pre_content_ping(content, has_generated_content=False) is False
|
||||
|
||||
|
||||
def test_anthropic_stream_forwards_ping_live_direct_call():
|
||||
ping = _anthropic_messages_ping_chunk()
|
||||
content = _anthropic_messages_content_chunk("hi")
|
||||
assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False, buffered_chunk_count=0) is True
|
||||
assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False, buffered_chunk_count=1) is False
|
||||
assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=True, buffered_chunk_count=0) is False
|
||||
assert _anthropic_stream_forwards_ping_live(content, has_generated_content=False, buffered_chunk_count=0) is False
|
||||
assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False) is True
|
||||
assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=True) is False
|
||||
assert _anthropic_stream_forwards_ping_live(content, has_generated_content=False) is False
|
||||
|
||||
|
||||
def test_anthropic_stream_error_is_gateway_verdict_direct_call():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue