mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge 53cacd8e7c into d70e64d973
This commit is contained in:
commit
ef098d29b9
7 changed files with 510 additions and 15 deletions
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5570
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15281
|
||||
"limit": 15279
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -84,7 +84,7 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1804
|
||||
"limit": 1802
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
|
|
@ -99,7 +99,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44358
|
||||
"limit": 44357
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
|
|
@ -111,7 +111,7 @@
|
|||
"limit": 19584
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29814
|
||||
"limit": 29812
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 110
|
||||
|
|
|
|||
|
|
@ -786,9 +786,11 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
aws_decoder: Final = AmazonAnthropicClaudeMessagesStreamDecoder(
|
||||
model=model,
|
||||
)
|
||||
completion_stream: Final = aws_decoder.aiter_bytes(
|
||||
httpx_response.aiter_bytes(chunk_size=aws_decoder.DEFAULT_CHUNK_SIZE)
|
||||
)
|
||||
# No ``chunk_size``: httpx's ByteChunker withholds bytes until that many
|
||||
# accumulate, stranding a smaller ``message_start`` frame until the next
|
||||
# upstream event, which after a reasoning phase is tens of seconds later
|
||||
# (BerriAI/litellm#38689).
|
||||
completion_stream: Final = aws_decoder.aiter_bytes(httpx_response.aiter_bytes())
|
||||
# Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients.
|
||||
return self.bedrock_sse_wrapper(
|
||||
completion_stream=completion_stream,
|
||||
|
|
@ -943,7 +945,6 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
|
|||
Iterator to return Bedrock invoke response in anthropic /messages format
|
||||
"""
|
||||
super().__init__(model=model)
|
||||
self.DEFAULT_CHUNK_SIZE = 1024
|
||||
|
||||
def _chunk_parser(self, chunk_data: dict) -> GChunk | ModelResponseStream | dict:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -421,6 +421,35 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
|
|||
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
_FALLBACK_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_EXACT_KEY_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(dict[str, list[str]])
|
||||
|
||||
|
||||
def _exact_key_fallback_entries(
|
||||
fallbacks: object,
|
||||
) -> list[dict[str, list[str]]]: # mutable-ok: mirrors the exact-key resolver's contract
|
||||
"""
|
||||
The well-formed ``{model_group: [chain]}`` entries of an untyped fallback list, typed for
|
||||
_get_fallback_model_group_for_lookup_groups.
|
||||
|
||||
Entries of any other shape are dropped rather than rejecting the whole list, because the
|
||||
resolver walks entries one at a time and can return an earlier well-formed entry's chain
|
||||
without ever reading a malformed one.
|
||||
"""
|
||||
try:
|
||||
entries: Final = _FALLBACK_LIST_ADAPTER.validate_python(fallbacks)
|
||||
except ValidationError:
|
||||
return []
|
||||
return [typed for entry in entries if (typed := _as_exact_key_fallback_entry(entry)) is not None]
|
||||
|
||||
|
||||
def _as_exact_key_fallback_entry(entry: object) -> dict[str, list[str]] | None:
|
||||
try:
|
||||
return _EXACT_KEY_FALLBACK_ENTRY_ADAPTER.validate_python(entry)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
|
||||
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
|
||||
|
||||
|
|
@ -542,6 +571,30 @@ def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, bu
|
|||
return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS
|
||||
|
||||
|
||||
async def _anthropic_messages_stream_without_fallback_protection(
|
||||
source_iterator: AsyncIterator[bytes],
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
"""No fallback is configured for the requested model group, so there is nothing the
|
||||
buffer-until-content protection in Router._aanthropic_messages_streaming_iterator would
|
||||
protect: forward the source iterator live instead of wrapping it.
|
||||
|
||||
A client that disconnects mid-stream leaves this generator suspended at `yield`
|
||||
rather than exhausted, so the `finally` below - not the `async for` running to
|
||||
completion - is what closes the upstream connection; without it, a disconnect
|
||||
during a long adaptive-thinking pass would leak the request to the provider.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
)
|
||||
|
||||
try:
|
||||
async for chunk in source_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
|
||||
await aclose_if_supported(source_iterator)
|
||||
|
||||
|
||||
class FallbackAwareAnthropicMessagesStream:
|
||||
"""
|
||||
Bare async generators can't carry the `_hidden_params` attribute the
|
||||
|
|
@ -5425,6 +5478,17 @@ class Router:
|
|||
|
||||
source_iterator: Final = response
|
||||
|
||||
model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
|
||||
if fallbacks_disabled_for_request(initial_kwargs) or not self._has_any_configured_fallback(
|
||||
model_group, initial_kwargs
|
||||
):
|
||||
# Nothing to fall back to, so buffering lifecycle frames to protect a mid-stream
|
||||
# fallback attempt would only add latency for no benefit: forward the source
|
||||
# iterator live, exactly as it would stream without this wrapper.
|
||||
return FallbackAwareAnthropicMessagesStream(
|
||||
_anthropic_messages_stream_without_fallback_protection(source_iterator), source_iterator
|
||||
)
|
||||
|
||||
async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
|
|
@ -7319,12 +7383,7 @@ class Router:
|
|||
# Use wildcard-aware lookup so order-based fallback also works for model
|
||||
# groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`).
|
||||
all_deployments: Final = self.get_model_list(model_name=original_model_group, team_id=_request_team_id) or []
|
||||
_order_set: Final[set] = {
|
||||
litellm.utils._get_deployment_order(d)
|
||||
for d in all_deployments
|
||||
if litellm.utils._get_deployment_order(d) is not None
|
||||
}
|
||||
order_values: Final[list] = sorted(_order_set)
|
||||
order_values: Final = litellm.utils.get_distinct_deployment_orders(all_deployments)
|
||||
if len(order_values) > 1 and not _skip_order_fallback:
|
||||
# Determine which order levels have already been tried
|
||||
current_target: Final = kwargs.get("_target_order")
|
||||
|
|
@ -8415,6 +8474,71 @@ class Router:
|
|||
return True
|
||||
return False
|
||||
|
||||
def _has_any_configured_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
|
||||
"""
|
||||
Whether any fallback deployment - general, context-window, content-policy, or a
|
||||
catch-all default - could resolve for this model group.
|
||||
|
||||
Gates whether _aanthropic_messages_streaming_iterator's buffer-until-content
|
||||
protection is worth paying for: that protection exists so a mid-stream provider
|
||||
error can retry against a fallback deployment before any lifecycle frame commits
|
||||
the client to this attempt. With no fallback destination configured at all, a
|
||||
retry can never happen, so holding message_start/content_block_start hostage
|
||||
until real content arrives protects nothing and only adds latency (most visibly
|
||||
on adaptive-thinking models, where the first content_block_delta can lag
|
||||
message_start by well over a minute).
|
||||
|
||||
Matching mirrors what async_function_with_fallbacks_common_utils actually resolves at
|
||||
retry time, which is not one rule for all three lists. Generic ``fallbacks`` resolve
|
||||
through get_fallback_model_group_for_lookup_groups, which accepts a stripped model-group
|
||||
match (a fallback keyed by the bare model name still arming a request routed with a
|
||||
provider prefix) and a "*" chain on top of an exact key, and a client-supplied
|
||||
non-standard ``fallbacks`` list (a plain list of model names, or of full override params)
|
||||
applies to every model group unconditionally rather than being keyed by one at all.
|
||||
``context_window_fallbacks`` and ``content_policy_fallbacks`` instead resolve through
|
||||
self._get_fallback_model_group_for_lookup_groups, which matches an exact key only and
|
||||
raises the original exception on a miss. Using one resolver for both kinds gets it wrong
|
||||
in both directions: the permissive one arms the buffer on wildcard- or stripped-keyed
|
||||
special fallbacks the retry path would reject, paying the lifecycle delay for a retry that
|
||||
can never happen, and the strict one reports "nothing to fall back to" for a stripped or
|
||||
wildcard generic chain a real error would in fact retry.
|
||||
|
||||
Two more retry paths in the same dispatcher fire without any of `fallbacks` /
|
||||
`context_window_fallbacks` / `content_policy_fallbacks` configured at all: order-based
|
||||
fallback (deployments in the model group at more than one `order` level) and weighted
|
||||
intra-group failover (`enable_weighted_failover`), both of which pick a different
|
||||
deployment for the retry, not the one that already streamed lifecycle frames live.
|
||||
"""
|
||||
fallbacks: Final = kwargs.get("fallbacks", self.fallbacks)
|
||||
if _check_non_standard_fallback_format(fallbacks=fallbacks):
|
||||
return True
|
||||
team_id: Final = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id")
|
||||
all_deployments: Final = self.get_model_list(model_name=model_group, team_id=team_id) or []
|
||||
if self.enable_weighted_failover:
|
||||
strategy, _ = self._get_routing_context(model_group, kwargs) # pyright: ignore[reportArgumentType] # Mapping is read-only, safe for dict param
|
||||
if strategy == "simple-shuffle" and len(all_deployments) > 1:
|
||||
return True
|
||||
if len(litellm.utils.get_distinct_deployment_orders(all_deployments)) > 1:
|
||||
return True
|
||||
lookup_groups: Final = fallback_lookup_groups(kwargs, model_group)
|
||||
if (
|
||||
fallbacks is not None
|
||||
and get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=lookup_groups)[0]
|
||||
is not None
|
||||
):
|
||||
return True
|
||||
special_fallback_lists: Final = (
|
||||
kwargs.get("context_window_fallbacks", self.context_window_fallbacks),
|
||||
kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks),
|
||||
)
|
||||
if any(
|
||||
self._get_fallback_model_group_for_lookup_groups(fallbacks=entries, lookup_groups=lookup_groups) is not None
|
||||
for entries in map(_exact_key_fallback_entries, special_fallback_lists)
|
||||
if entries
|
||||
):
|
||||
return True
|
||||
return self._has_default_fallbacks()
|
||||
|
||||
def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
|
||||
"""
|
||||
Whether a content-policy fallback would resolve for this request, keyed the same way
|
||||
|
|
|
|||
|
|
@ -4940,6 +4940,17 @@ def _get_deployment_order(deployment: dict | Any) -> int | None:
|
|||
return order
|
||||
|
||||
|
||||
def get_distinct_deployment_orders(deployments: Sequence[Mapping[str, Any]]) -> tuple[int, ...]:
|
||||
"""
|
||||
The ascending distinct `order` levels present across `deployments`, ignoring those without one.
|
||||
|
||||
More than one level means the router can retry a failure against a different deployment in the
|
||||
same model group, so callers deciding whether an order-based retry is reachable read this rather
|
||||
than each deployment's order.
|
||||
"""
|
||||
return tuple(sorted({order for d in deployments for order in [_get_deployment_order(d)] if order is not None}))
|
||||
|
||||
|
||||
def get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list:
|
||||
if target_order is not None:
|
||||
return [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
|
||||
|
|
|
|||
|
|
@ -478,6 +478,152 @@ def test_has_content_policy_fallback_default_fallbacks_arm():
|
|||
assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_general_fallbacks_arm():
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
assert router._has_any_configured_fallback("other-group", {}) is False
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_context_window_fallbacks_arm():
|
||||
router = Router(
|
||||
model_list=[FABLE_TIER, OPUS_TARGET],
|
||||
context_window_fallbacks=[{"fable-tier": ["opus-target"]}],
|
||||
)
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
assert router._has_any_configured_fallback("other-group", {}) is False
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_content_policy_fallbacks_arm():
|
||||
router = Router(
|
||||
model_list=[FABLE_TIER, OPUS_TARGET],
|
||||
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
|
||||
)
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_default_fallbacks_arm():
|
||||
router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
|
||||
|
||||
assert router._has_any_configured_fallback("any-group", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_nothing_configured():
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET])
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is False
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_honors_per_request_kwargs_override():
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET])
|
||||
|
||||
assert (
|
||||
router._has_any_configured_fallback("fable-tier", {"fallbacks": [{"fable-tier": ["opus-target"]}]}) is True
|
||||
)
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_matches_stripped_model_group():
|
||||
"""Regression: async_function_with_fallbacks_common_utils resolves a fallback keyed by
|
||||
the bare model name even when the request was routed with a provider prefix (e.g. a
|
||||
fallback keyed "fable-tier" still arms "openai/fable-tier"); the gate must recognize
|
||||
that same stripped match instead of requiring an exact model-group key."""
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
assert router._has_any_configured_fallback("openai/fable-tier", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_matches_non_standard_client_fallbacks():
|
||||
"""Regression: a client-supplied non-standard `fallbacks` list (a plain list of model
|
||||
names, not keyed by model group at all) applies unconditionally at retry time via
|
||||
_check_non_standard_fallback_format, so the gate must arm for it too rather than only
|
||||
recognizing the dict-keyed `{"model_group": [...]}` shape."""
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET])
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {"fallbacks": ["opus-target"]}) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fallback_kind", ["context_window_fallbacks", "content_policy_fallbacks"])
|
||||
@pytest.mark.parametrize(
|
||||
"configured_key, requested_group",
|
||||
[("*", "fable-tier"), ("fable-tier", "openai/fable-tier")],
|
||||
)
|
||||
def test_has_any_configured_fallback_ignores_special_fallbacks_the_retry_path_rejects(
|
||||
fallback_kind: str, configured_key: str, requested_group: str
|
||||
):
|
||||
"""Regression: async_function_with_fallbacks_common_utils resolves context-window and
|
||||
content-policy chains through _get_fallback_model_group_for_lookup_groups, which matches an
|
||||
exact model-group key only and raises the original exception on a miss - it honors neither a
|
||||
"*" chain nor a stripped model-group match. Arming the buffer on those entries pays the
|
||||
buffer-until-content lifecycle delay for a retry that can never happen."""
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], **{fallback_kind: [{configured_key: ["opus-target"]}]})
|
||||
|
||||
assert router._has_any_configured_fallback(requested_group, {}) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fallback_kind", ["context_window_fallbacks", "content_policy_fallbacks"])
|
||||
def test_has_any_configured_fallback_arms_on_exact_keyed_special_fallbacks(fallback_kind: str):
|
||||
"""The flip side of the exact-key rule: a special chain keyed by the requested group is
|
||||
exactly what the retry path resolves, so the buffer must still arm for it."""
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], **{fallback_kind: [{"fable-tier": ["opus-target"]}]})
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_matches_wildcard_general_fallbacks():
|
||||
"""Counterpart to the special-fallback exact-key rule: generic `fallbacks` resolve through
|
||||
get_fallback_model_group_for_lookup_groups, which does honor a "*" chain, so tightening the
|
||||
special lists must not also stop the gate arming on a wildcard generic chain."""
|
||||
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_arms_on_order_based_deployments():
|
||||
"""Regression: async_function_with_fallbacks_common_utils retries against a higher-order
|
||||
deployment in the same model group whenever more than one `order` level is present, even
|
||||
with zero `fallbacks`/`context_window_fallbacks`/`content_policy_fallbacks` configured -
|
||||
the gate must recognize that retry path too, or a mid-stream error can still trigger an
|
||||
order-based retry that appends a second message_start onto a stream already forwarded live."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "fable-tier",
|
||||
"litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test", "order": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "fable-tier",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test", "order": 2},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
|
||||
|
||||
def test_has_any_configured_fallback_arms_on_weighted_failover():
|
||||
"""Regression: enable_weighted_failover lets a retryable failure re-pick across the
|
||||
model group's other deployments before any cross-group fallback runs, independent of
|
||||
`fallbacks` config entirely - the gate must arm for it too. It only has somewhere else to
|
||||
re-pick when the group itself holds more than one deployment, so a single-deployment group
|
||||
must not arm on enable_weighted_failover alone."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
FABLE_TIER,
|
||||
{
|
||||
"model_name": "fable-tier",
|
||||
"litellm_params": {"model": "anthropic/claude-fable-5-mini", "api_key": "sk-test"},
|
||||
},
|
||||
OPUS_TARGET,
|
||||
],
|
||||
enable_weighted_failover=True,
|
||||
)
|
||||
|
||||
assert router._has_any_configured_fallback("fable-tier", {}) is True
|
||||
assert router._has_any_configured_fallback("opus-target", {}) is False
|
||||
|
||||
|
||||
def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested():
|
||||
router = _router(content_policy_fallbacks=None)
|
||||
fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]
|
||||
|
|
|
|||
|
|
@ -3250,3 +3250,108 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo
|
|||
)
|
||||
|
||||
assert result.get("output_config") == {"format": schema_format}
|
||||
|
||||
|
||||
def _bedrock_invoke_event_frame(payload: dict) -> bytes:
|
||||
"""Encode one Bedrock invoke event-stream frame carrying an Anthropic event."""
|
||||
import base64
|
||||
import struct
|
||||
from binascii import crc32
|
||||
|
||||
encoded_event = base64.b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode()
|
||||
body = json.dumps({"bytes": encoded_event}).encode()
|
||||
|
||||
def _str_header(name: str, value: str) -> bytes:
|
||||
name_b = name.encode()
|
||||
value_b = value.encode()
|
||||
return (
|
||||
struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b
|
||||
)
|
||||
|
||||
headers = (
|
||||
_str_header(":event-type", "chunk")
|
||||
+ _str_header(":content-type", "application/json")
|
||||
+ _str_header(":message-type", "event")
|
||||
)
|
||||
prelude = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers))
|
||||
prelude_crc = crc32(prelude) & 0xFFFFFFFF
|
||||
prelude_crc_b = struct.pack("!I", prelude_crc)
|
||||
msg_crc_b = struct.pack("!I", crc32(prelude_crc_b + headers + body, prelude_crc) & 0xFFFFFFFF)
|
||||
return prelude + prelude_crc_b + headers + body + msg_crc_b
|
||||
|
||||
|
||||
_MESSAGE_START_EVENT = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_bdrk_01WxYzAbCdEfGhIjKlMnOpQr",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 41, "output_tokens": 1},
|
||||
},
|
||||
}
|
||||
|
||||
_CONTENT_BLOCK_START_EVENT = {
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "thinking", "thinking": "", "signature": ""},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_flushes_message_start_before_the_next_upstream_event():
|
||||
"""``message_start`` must reach the client as soon as Bedrock sends it.
|
||||
|
||||
Bedrock emits ``message_start`` immediately, then goes silent for the whole
|
||||
reasoning phase (tens of seconds at high effort) before the first content
|
||||
block. Reading the response with an ``httpx`` ``chunk_size`` stranded the
|
||||
preamble in httpx's ByteChunker until enough further bytes accumulated, so
|
||||
the client's first byte landed at first-content time and tripped its
|
||||
first-byte watchdog. Regression for BerriAI/litellm#38689.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
message_start_frame = _bedrock_invoke_event_frame(_MESSAGE_START_EVENT)
|
||||
# The stall only happens for a preamble smaller than the read threshold, so a
|
||||
# frame that grew past it would make this test pass without the fix.
|
||||
assert len(message_start_frame) < 1024
|
||||
|
||||
reasoning_finished = asyncio.Event()
|
||||
|
||||
class _ReasoningStall(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield message_start_frame
|
||||
await reasoning_finished.wait()
|
||||
yield _bedrock_invoke_event_frame(_CONTENT_BLOCK_START_EVENT)
|
||||
|
||||
cfg = AmazonAnthropicClaudeMessagesConfig()
|
||||
stream = cfg.get_async_streaming_response_iterator(
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
httpx_response=httpx.Response(200, stream=_ReasoningStall()),
|
||||
request_body={},
|
||||
litellm_logging_obj=LiteLLMLoggingObj(
|
||||
model="bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[{"role": "user", "content": "think hard"}],
|
||||
stream=True,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test_flush_message_start",
|
||||
function_id="test_flush_message_start",
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
# Fails by timing out while the upstream is still mid-reasoning if the
|
||||
# preamble is being held back.
|
||||
first_chunk = await asyncio.wait_for(stream.__anext__(), timeout=5)
|
||||
assert b"event: message_start" in first_chunk
|
||||
|
||||
reasoning_finished.set()
|
||||
second_chunk = await asyncio.wait_for(stream.__anext__(), timeout=5)
|
||||
assert b"event: content_block_start" in second_chunk
|
||||
finally:
|
||||
reasoning_finished.set()
|
||||
await stream.aclose()
|
||||
|
|
|
|||
|
|
@ -11842,6 +11842,13 @@ def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream:
|
|||
|
||||
|
||||
def _anthropic_messages_make_router() -> Router:
|
||||
"""A ``fallbacks`` entry mapping "primary" -> "fallback" is required, not just
|
||||
a second model_list entry: _has_any_configured_fallback gates the whole
|
||||
buffer-until-content mechanism on whether the router could actually resolve
|
||||
a fallback for the model group, so a router with nothing configured under
|
||||
``fallbacks``/``context_window_fallbacks``/``content_policy_fallbacks``
|
||||
(matching a real deployment with no fallback set up) skips buffering
|
||||
entirely and streams live - see test_anthropic_messages_streaming_iterator_skips_buffering_without_any_configured_fallback."""
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -11857,7 +11864,8 @@ def _anthropic_messages_make_router() -> Router:
|
|||
"model": "bedrock/anthropic.claude-sonnet-4-5",
|
||||
},
|
||||
},
|
||||
]
|
||||
],
|
||||
fallbacks=[{"primary": ["fallback"]}],
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -11905,6 +11913,30 @@ class _AnthropicMessagesRaisingByteStream:
|
|||
self.closed = True
|
||||
|
||||
|
||||
class _AnthropicMessagesHangingByteStream:
|
||||
"""Yields the given chunks, then hangs forever on the next `__anext__()`
|
||||
instead of raising StopAsyncIteration - simulates a real upstream stuck
|
||||
mid-thinking-pass, so a test can prove a chunk reached the caller without
|
||||
waiting for the rest of the stream (which, here, never arrives)."""
|
||||
|
||||
def __init__(self, chunks: list) -> None:
|
||||
self._chunks = list(chunks)
|
||||
self._hidden_params: dict = {}
|
||||
self.closed = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
if self._chunks:
|
||||
return self._chunks.pop(0)
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable") # pragma: no cover
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _AnthropicMessagesFallbackByteStream:
|
||||
def __init__(self, chunks: list, hidden_params: dict | None = None) -> None:
|
||||
self._chunks = list(chunks)
|
||||
|
|
@ -12017,6 +12049,82 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_
|
|||
await wrapped.__anext__()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_iterator_skips_buffering_without_any_configured_fallback():
|
||||
"""Regression: with no fallback configured for the model group (no
|
||||
``fallbacks``, ``context_window_fallbacks``, ``content_policy_fallbacks``,
|
||||
or catch-all default), lifecycle frames like message_start must reach the
|
||||
caller as soon as the source produces them, not be buffered until real
|
||||
content arrives. Buffering exists to protect a mid-stream fallback
|
||||
attempt; with nothing to fall back to, there is nothing to protect, and
|
||||
the delay is pure added latency (most visible on adaptive-thinking models,
|
||||
where the first content_block_delta can lag message_start by a minute or
|
||||
more). The source here hangs forever after its first chunk, so this can
|
||||
only pass if that chunk was forwarded live rather than buffered."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"},
|
||||
},
|
||||
]
|
||||
)
|
||||
source = _AnthropicMessagesHangingByteStream([_anthropic_messages_message_start_chunk()])
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source, initial_kwargs={"model": "primary"}
|
||||
)
|
||||
|
||||
first_chunk = await asyncio.wait_for(wrapped.__anext__(), timeout=1.0)
|
||||
assert first_chunk == _anthropic_messages_message_start_chunk()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_iterator_closes_upstream_on_disconnect_without_fallback():
|
||||
"""Regression: without any fallback configured, the live-forwarding generator must
|
||||
still close the upstream stream when the caller disconnects mid-stream (aclose()
|
||||
on the wrapper), not just when the source iterator runs to exhaustion on its own -
|
||||
otherwise a client that disconnects during a long thinking pass leaks the upstream
|
||||
request/connection to the provider indefinitely."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"},
|
||||
},
|
||||
]
|
||||
)
|
||||
source = _AnthropicMessagesHangingByteStream([_anthropic_messages_message_start_chunk()])
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source, initial_kwargs={"model": "primary"}
|
||||
)
|
||||
|
||||
await asyncio.wait_for(wrapped.__anext__(), timeout=1.0)
|
||||
assert source.closed is False
|
||||
|
||||
await wrapped.aclose()
|
||||
assert source.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_streaming_iterator_still_buffers_lifecycle_frames_when_fallback_configured():
|
||||
"""Confirms the buffer-until-content protection is still intact once a
|
||||
fallback IS configured for the model group - only the no-fallback case
|
||||
added by this fix skips it. The source hangs forever after message_start,
|
||||
so if it were forwarded live this would resolve within the timeout instead
|
||||
of raising."""
|
||||
router = _anthropic_messages_make_router()
|
||||
source = _AnthropicMessagesHangingByteStream([_anthropic_messages_message_start_chunk()])
|
||||
|
||||
wrapped = await router._aanthropic_messages_streaming_iterator(
|
||||
response=source, initial_kwargs={"model": "primary"}
|
||||
)
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(wrapped.__anext__(), timeout=0.2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback():
|
||||
"""Greptile review round: transport-level buffering can coalesce a real
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue