refactor(router): name the pre-content stream retry helpers for both stream paths and let the retry loop take its attempt function (#46474)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-14 11:25:15 -07:00 • committed by GitHub
parent 389403fede
commit 4f34856a65
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 67 additions and 57 deletions

View file

@ -24,6 +24,7 @@ from collections import defaultdict
from collections.abc import (
AsyncGenerator,
AsyncIterator,
Awaitable,
Callable,
Generator,
Iterable,
@ -628,7 +629,7 @@ def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error
return has_generated_content or not error.is_pre_first_chunk
def _anthropic_stream_raised_error_status(error: Exception) -> int | None:
def _stream_raised_error_status(error: Exception) -> int | None:
raw_status: Final = getattr(error, "status_code", None)
if isinstance(raw_status, int):
return raw_status
@ -638,27 +639,29 @@ def _anthropic_stream_raised_error_status(error: Exception) -> int | None:
return response_status if isinstance(response_status, int) else None
def _anthropic_stream_fallback_error_for_raised(
error: Exception, model: str, has_generated_content: bool
def _stream_fallback_error_for_raised(
error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic"
) -> "MidStreamFallbackError | None":
"""The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged."""
if has_generated_content:
return None
status_code: Final = _anthropic_stream_raised_error_status(error)
status_code: Final = _stream_raised_error_status(error)
if status_code is None:
return _anthropic_stream_pre_content_error(error, model)
return _stream_pre_content_error(error, model, llm_provider)
retriable: Final = litellm.should_retry(status_code)
return _anthropic_stream_pre_content_error(error, model) if retriable else None
return _stream_pre_content_error(error, model, llm_provider) if retriable else None
def _anthropic_stream_pre_content_error(error: Exception, model: str) -> "MidStreamFallbackError":
def _stream_pre_content_error(
error: Exception, model: str, llm_provider: str = "anthropic"
) -> "MidStreamFallbackError":
"""The envelope the fallback chain judges a failure by when the client has received no content yet."""
from litellm.exceptions import MidStreamFallbackError
return MidStreamFallbackError(
message=str(error),
model=model,
llm_provider="anthropic",
llm_provider=llm_provider,
original_exception=error,
is_pre_first_chunk=True,
)
@ -688,7 +691,7 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception:
return error.original_exception if error.original_exception is not None else error
def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
"""
Whether `chunk` should make Router._aanthropic_messages_streaming_iterator
commit to the primary Anthropic stream (real content arrived, or the
@ -705,18 +708,18 @@ def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, bu
@dataclass(frozen=True, slots=True)
class _AnthropicStreamRetryOpened:
class _StreamRetryOpened:
response: object
attempted_retries: int
max_retries: int
@dataclass(frozen=True, slots=True)
class _AnthropicStreamRetriesExhausted:
class _StreamRetriesExhausted:
error: "MidStreamFallbackError"
_AnthropicStreamRetryOutcome: TypeAlias = _AnthropicStreamRetryOpened | _AnthropicStreamRetriesExhausted
_StreamRetryOutcome: TypeAlias = _StreamRetryOpened | _StreamRetriesExhausted
MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200
@ -5688,7 +5691,7 @@ class Router:
if _anthropic_stream_forwards_ping_live(chunk, has_generated_content):
yield chunk
continue
if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
if _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
# detection parses the accumulated buffer plus the current chunk, never the
@ -5777,16 +5780,20 @@ class Router:
fallback_error: Final = (
stream_error
if isinstance(stream_error, MidStreamFallbackError)
else _anthropic_stream_fallback_error_for_raised(stream_error, model, has_generated_content)
else _stream_fallback_error_for_raised(stream_error, model, has_generated_content)
)
if fallback_error is None:
raise stream_error
outcome: Final = await self._aanthropic_messages_retry_same_group(fallback_error, initial_kwargs)
outcome: Final = await self._astream_retry_same_group(
fallback_error,
initial_kwargs,
attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt,
)
match outcome:
case _AnthropicStreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget):
case _StreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget):
async for item in self._aanthropic_messages_yield_recovered(retried, wrapper, (attempted, budget)):
yield item
case _AnthropicStreamRetriesExhausted(error=last_error):
case _StreamRetriesExhausted(error=last_error):
async for item in self._aanthropic_messages_fallback_attempt(last_error, initial_kwargs, wrapper):
yield item
case _:
@ -5840,7 +5847,7 @@ class Router:
return None
return _get_num_retries_from_retry_policy(exception=trigger, retry_policy=policy)
def _anthropic_messages_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]:
def _stream_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]:
"""
The budget an earlier retry of this request committed to, else the retry policy's grant when one
names this error, else the plain budget, with whether a policy governs the retry: a committed budget
@ -5890,14 +5897,14 @@ class Router:
if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk):
return None
frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group)
budget, _ = self._anthropic_messages_retry_budget(frame_error, kwargs)
budget, _ = self._stream_retry_budget(frame_error, kwargs)
if budget > attempted_retries_for_request(kwargs):
return frame_error
if self._anthropic_messages_stream_can_fall_back(model_group, kwargs):
return frame_error
return None
def _anthropic_messages_should_retry(
def _stream_should_retry(
self,
trigger: Exception,
healthy_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type
@ -5929,29 +5936,34 @@ class Router:
return False
return True
async def _aanthropic_messages_retry_same_group(
self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object]
) -> _AnthropicStreamRetryOutcome:
async def _astream_retry_same_group(
self,
e: "MidStreamFallbackError",
initial_kwargs: Mapping[str, object],
attempt_function: "Callable[..., Awaitable[object]]",
) -> _StreamRetryOutcome:
"""
Re-runs the attempt within the request's own model group, the way async_function_with_retries
would have for a failure raised before the stream opened, until a retry opens a stream or the
budget runs out. Each retry's stream carries its own wrapper with the remaining budget, so a
retry that drops before content again continues the same count instead of starting over.
`attempt_function` is the request's own attempt (the /v1/messages attempt, or a chat completions one).
"""
model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs)
healthy_deployments, all_deployments = await self._async_get_healthy_deployments(
model=model_group, parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs)
)
budget, policy_applies = self._anthropic_messages_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs)
budget, policy_applies = self._stream_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs)
last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see
for attempt in range(attempted_retries_for_request(initial_kwargs), budget):
trigger = _mid_stream_retry_trigger(last_error)
if not policy_applies and not self._anthropic_messages_should_retry(
if not policy_applies and not self._stream_should_retry(
trigger, healthy_deployments, all_deployments, initial_kwargs
):
return _AnthropicStreamRetriesExhausted(last_error)
self.log_retry(kwargs=retry_kwargs, e=trigger)
return _StreamRetriesExhausted(last_error)
if get_metadata_variable_name_from_kwargs(retry_kwargs) in retry_kwargs:
self.log_retry(kwargs=retry_kwargs, e=trigger)
await asyncio.sleep(
self._time_to_sleep_before_retry(
e=trigger,
@ -5963,20 +5975,18 @@ class Router:
)
record_retry_attempt(retry_kwargs, attempted_retries=attempt + 1, max_retries=budget)
verbose_router_logger.debug(
"Retrying anthropic_messages stream dropped before content, attempt %s of %s", attempt + 1, budget
"Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget
)
try:
response = await self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt(**retry_kwargs)
except Exception as retry_error: # noqa: BLE001 # every failure of a retry before its stream opens is the fallback chain's to judge
wrapped = _anthropic_stream_fallback_error_for_raised(retry_error, model_group, False)
response = await attempt_function(**retry_kwargs)
except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry
wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider)
if wrapped is None:
return _AnthropicStreamRetriesExhausted(
_anthropic_stream_pre_content_error(retry_error, model_group)
)
return _StreamRetriesExhausted(_stream_pre_content_error(retry_error, model_group, e.llm_provider))
last_error = wrapped
continue
return _AnthropicStreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget)
return _AnthropicStreamRetriesExhausted(last_error)
return _StreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget)
return _StreamRetriesExhausted(last_error)
async def _aanthropic_messages_yield_recovered(
self,
@ -6100,7 +6110,7 @@ class Router:
async def _ageneric_api_call_with_fallbacks_anthropic_messages_attempt(
self,
model: str,
original_generic_function: Callable,
original_generic_function: Callable[..., Awaitable[object]],
**kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site
) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]:
"""

View file

@ -97,7 +97,7 @@ ignored_function_names = [
"_async_get_available_deployment_for_pass_through", # Same, through async_get_available_deployment_for_pass_through in test_router.py
"_embedding",
"_aembedding",
"_anthropic_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py
"_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py
"_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py
"_request_fallback_list", # Tested through every mid-stream retry test in test_router.py
"_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call
@ -105,8 +105,8 @@ ignored_function_names = [
"_anthropic_messages_group_retry_policy", # Tested through the retry budget precedence test in test_router.py
"_anthropic_messages_resolved_retry_policy", # Tested through the malformed retry policy tests in test_router.py
"_anthropic_messages_plain_retry_budget", # Tested through the retry budget precedence test in test_router.py
"_anthropic_messages_should_retry", # Tested through every mid-stream retry test in test_router.py
"_aanthropic_messages_retry_same_group", # Tested through the dropped-before-content retry tests in test_router.py
"_stream_should_retry", # Tested through every mid-stream retry test in test_router.py
"_astream_retry_same_group", # Tested through the dropped-before-content retry tests in test_router.py
"_aanthropic_messages_yield_recovered", # Tested through every mid-stream retry and fallback test in test_router.py
"_anthropic_messages_policy_retries", # Tested through the retry budget precedence test in test_router.py
"_get_wildcard_deployments", # Tested through the get_model_list_of_routed_group wildcard test in test_router.py

View file

@ -46,11 +46,11 @@ from litellm.router import (
MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS,
FallbackAwareAnthropicMessagesStream,
_anthropic_stream_commits_now,
_stream_commits_now,
_anthropic_stream_error_is_gateway_verdict,
_anthropic_stream_fallback_error_for_raised,
_stream_fallback_error_for_raised,
_anthropic_stream_forwards_ping_live,
_anthropic_stream_raised_error_status,
_stream_raised_error_status,
_anthropic_stream_should_decline_fallback,
_is_retriable_anthropic_status,
_responses_stream_holds_event,
@ -15593,10 +15593,10 @@ def test_anthropic_stream_should_decline_fallback_direct_call():
def test_anthropic_stream_commits_now_direct_call():
content = _anthropic_messages_content_chunk("hi")
lifecycle_chunk = _anthropic_messages_message_start_chunk()
assert _anthropic_stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True
assert _anthropic_stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False
assert _stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True
assert _stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False
assert (
_anthropic_stream_commits_now(
_stream_commits_now(
lifecycle_chunk,
has_generated_content=False,
buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
@ -15604,7 +15604,7 @@ def test_anthropic_stream_commits_now_direct_call():
is True
)
assert (
_anthropic_stream_commits_now(
_stream_commits_now(
lifecycle_chunk,
has_generated_content=False,
buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - 1,
@ -16175,15 +16175,15 @@ def test_anthropic_messages_retry_budget_precedence_direct_call():
num_retries=3, retry_policy=RetryPolicy(InternalServerErrorRetries=4)
)
assert router._anthropic_messages_retry_budget(drop, {"model": "glm"}) == (3, False)
assert router._anthropic_messages_retry_budget(drop, routed) == (2, False)
assert router._anthropic_messages_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False)
assert policy_router._anthropic_messages_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True)
assert policy_router._anthropic_messages_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False)
assert policy_router._anthropic_messages_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False)
assert router._stream_retry_budget(drop, {"model": "glm"}) == (3, False)
assert router._stream_retry_budget(drop, routed) == (2, False)
assert router._stream_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False)
assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True)
assert policy_router._stream_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False)
assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False)
committed = {**routed, "litellm_metadata": {**routed["litellm_metadata"], "attempted_retries": 1, "max_retries": 5}}
assert router._anthropic_messages_retry_budget(drop, committed) == (5, False)
assert policy_router._anthropic_messages_retry_budget(reset, committed) == (5, True)
assert router._stream_retry_budget(drop, committed) == (5, False)
assert policy_router._stream_retry_budget(reset, committed) == (5, True)
def test_anthropic_messages_stream_can_retry_direct_call():
@ -16998,7 +16998,7 @@ async def test_anthropic_messages_raised_provider_error_after_content_propagates
ids=["int", "digit-str", "response-only", "none"],
)
def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, expected_status):
assert _anthropic_stream_raised_error_status(error) == expected_status
assert _stream_raised_error_status(error) == expected_status
@pytest.mark.parametrize(
@ -17014,7 +17014,7 @@ def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, ex
def test_anthropic_stream_fallback_error_for_raised_gates_like_a_detected_error_event(
error, has_generated_content, converts
):
converted = _anthropic_stream_fallback_error_for_raised(error, "primary", has_generated_content)
converted = _stream_fallback_error_for_raised(error, "primary", has_generated_content)
if not converts:
assert converted is None
return