diff --git a/litellm/router.py b/litellm/router.py index 313c404cabe..ee8a58f21e2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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]]: """ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index e7c2e4234f5..af0e24038ac 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4a7a23ff57b..64f3f88cd89 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -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