diff --git a/python/tests/core/llm/test_llm_retry.py b/python/tests/core/llm/test_llm_retry.py index a98a7ab0..e095c7f7 100644 --- a/python/tests/core/llm/test_llm_retry.py +++ b/python/tests/core/llm/test_llm_retry.py @@ -279,6 +279,103 @@ async def mock_sleep(delay): assert delays == [2.5] + @pytest.mark.asyncio + async def test_retry_after_capped_at_max_delay(self): + """A huge Retry-After must not turn into a multi-minute in-place sleep.""" + from timbal.core.llm.retry import MAX_RETRY_DELAY + + response = MagicMock() + response.headers = {"Retry-After": "999"} + delays = [] + + async def rate_limited_stream(): + raise OpenAIRateLimitError("Rate limit exceeded", response=response, body=None) + yield + + async def mock_sleep(delay): + delays.append(delay) + + with patch("asyncio.sleep", side_effect=mock_sleep): + with patch("timbal.core.llm.retry.random.uniform", return_value=0.25): + with pytest.raises(OpenAIRateLimitError): + async for _ in _retry_on_error(rate_limited_stream, max_retries=1, retry_delay=1.0, context="Test"): + pass + + assert delays == [MAX_RETRY_DELAY] + + @pytest.mark.asyncio + async def test_fail_fast_rate_limit_raises_without_retry(self): + """With fail_fast_rate_limit, rate limits raise immediately — no sleep, no retry.""" + attempt_count = 0 + response = MagicMock() + response.headers = {"Retry-After": "999"} + delays = [] + + async def rate_limited_stream(): + nonlocal attempt_count + attempt_count += 1 + raise OpenAIRateLimitError("Rate limit exceeded", response=response, body=None) + yield + + async def mock_sleep(delay): + delays.append(delay) + + with patch("asyncio.sleep", side_effect=mock_sleep): + with pytest.raises(OpenAIRateLimitError): + async for _ in _retry_on_error( + rate_limited_stream, max_retries=3, retry_delay=1.0, context="Test", fail_fast_rate_limit=True, + ): + pass + + assert attempt_count == 1 + assert delays == [] + + @pytest.mark.asyncio + async def test_fail_fast_rate_limit_applies_to_429_status(self): + """429 APIStatusError (not just RateLimitError) also fails fast.""" + attempt_count = 0 + response = MagicMock() + response.status_code = 429 + response.headers = {} + + async def rate_limited_stream(): + nonlocal attempt_count + attempt_count += 1 + raise OpenAIAPIStatusError("Too many requests", response=response, body=None) + yield + + with pytest.raises(OpenAIAPIStatusError): + async for _ in _retry_on_error( + rate_limited_stream, max_retries=3, retry_delay=0.01, context="Test", fail_fast_rate_limit=True, + ): + pass + + assert attempt_count == 1 + + @pytest.mark.asyncio + async def test_fail_fast_rate_limit_still_retries_server_errors(self): + """fail_fast_rate_limit only affects rate limits — 5xx still retries in place.""" + attempt_count = 0 + response = MagicMock() + response.status_code = 503 + response.headers = {} + + async def flaky_stream(): + nonlocal attempt_count + attempt_count += 1 + if attempt_count < 2: + raise OpenAIAPIStatusError("Service unavailable", response=response, body=None) + yield "success" + + chunks = [] + async for chunk in _retry_on_error( + flaky_stream, max_retries=2, retry_delay=0.01, context="Test", fail_fast_rate_limit=True, + ): + chunks.append(chunk) + + assert chunks == ["success"] + assert attempt_count == 2 + @pytest.mark.asyncio async def test_invalid_retry_after_uses_jitter_delay(self): """Test that malformed Retry-After headers are ignored.""" diff --git a/python/tests/core/test_fallback_model.py b/python/tests/core/test_fallback_model.py index 7511e931..2401c2f3 100644 --- a/python/tests/core/test_fallback_model.py +++ b/python/tests/core/test_fallback_model.py @@ -50,6 +50,28 @@ async def router(**kwargs): assert [call["model"] for call in calls] == ["openai/primary", "openai/backup"] assert all(call["temperature"] == 0.2 for call in calls) + @pytest.mark.asyncio + async def test_fail_fast_rate_limit_set_for_all_but_last_entry(self): + """Every entry with a fallback behind it must fail over on 429 instead of + sleeping through Retry-After in place; the last entry retries normally.""" + model = FallbackModel("openai/primary", "openai/middle", "openai/last") + calls = [] + + async def router(**kwargs): + calls.append(kwargs) + if kwargs["model"] != "openai/last": + raise _status_error(429) + yield "ok" + + chunks = [chunk async for chunk in model.route(router)] + + assert chunks == ["ok"] + assert [(call["model"], call["fail_fast_rate_limit"]) for call in calls] == [ + ("openai/primary", True), + ("openai/middle", True), + ("openai/last", False), + ] + @pytest.mark.asyncio async def test_uses_per_entry_retry_and_auth_overrides(self): model = FallbackModel( diff --git a/python/timbal/core/fallback_model.py b/python/timbal/core/fallback_model.py index 11e3f82c..a8ef6f80 100644 --- a/python/timbal/core/fallback_model.py +++ b/python/timbal/core/fallback_model.py @@ -75,11 +75,17 @@ async def route( for index, entry in enumerate(self.entries): started = False + has_fallback = index + 1 < len(self.entries) kwargs = { **llm_router_kwargs, "model": entry.model, "max_retries": entry.max_retries, "retry_delay": entry.retry_delay, + # A rate limit means the provider is unavailable for a while by + # definition — while another model is still available, fail over + # immediately instead of sleeping through Retry-After in place. + # The last entry has nowhere to go, so it retries normally. + "fail_fast_rate_limit": has_fallback, } if entry.api_key is not None: kwargs["api_key"] = entry.api_key @@ -98,7 +104,7 @@ async def route( raise errors.append((entry.model, exc)) - next_model = self.entries[index + 1].model if index + 1 < len(self.entries) else None + next_model = self.entries[index + 1].model if has_fallback else None logger.warning( "Falling back to next LLM model", failed_model=entry.model, diff --git a/python/timbal/core/llm/retry.py b/python/timbal/core/llm/retry.py index a0a5a69f..4090d270 100644 --- a/python/timbal/core/llm/retry.py +++ b/python/timbal/core/llm/retry.py @@ -12,7 +12,9 @@ MAX_RETRY_DELAY = 30.0 -async def _retry_on_error(async_gen_func, max_retries: int, retry_delay: float, context: str): +async def _retry_on_error( + async_gen_func, max_retries: int, retry_delay: float, context: str, fail_fast_rate_limit: bool = False, +): """Helper to retry an async generator function on transient failures. Retryable errors (using SDK exception types): @@ -28,11 +30,20 @@ async def _retry_on_error(async_gen_func, max_retries: int, retry_delay: float, - Invalid requests (400, 404) - Other 4xx client errors + Retry delays honor a ``Retry-After`` response header as a floor, but are + always capped at ``MAX_RETRY_DELAY`` — a provider-side cooldown of minutes + must never turn into an in-place sleep of minutes. + Args: async_gen_func: Async callable that returns an async generator max_retries: Maximum number of retry attempts retry_delay: Base delay for exponential backoff context: Description for logging (e.g., "Anthropic API") + fail_fast_rate_limit: If True, rate-limit errors (RateLimitError / 429) + are raised immediately instead of retried in place. Set by + FallbackModel for every entry that still has a fallback behind it: + a rate-limited provider is unavailable for a while by definition, + so the chain should move on rather than sleep. Yields: Items from the async generator @@ -103,13 +114,21 @@ async def _retry_on_error(async_gen_func, max_retries: int, retry_delay: float, ) raise + if fail_fast_rate_limit and error_type == "rate_limit": + logger.warning( + "Rate limited and a fallback model is configured, failing over without retrying", + context=context, + error=error_msg, + ) + raise + # Retry logic for retryable errors if attempt < max_retries: cap = min(retry_delay * (2**attempt), MAX_RETRY_DELAY) delay = random.uniform(0, cap) retry_after = _retry_after_seconds(last_error) if retry_after is not None: - delay = max(delay, retry_after) + delay = min(max(delay, retry_after), MAX_RETRY_DELAY) logger.warning( "Retryable error from LLM provider, retrying...", context=context, diff --git a/python/timbal/core/llm/router.py b/python/timbal/core/llm/router.py index 98528c79..98a1109e 100644 --- a/python/timbal/core/llm/router.py +++ b/python/timbal/core/llm/router.py @@ -29,6 +29,7 @@ async def _llm_router( api_key: str | SecretStr | None = None, max_retries: int = 0, retry_delay: float = 1.0, + fail_fast_rate_limit: bool = False, provider_params: dict[str, Any] | None = None, ) -> Message: # type: ignore[misc] # Declared as Message for framework schema generation; runtime is an async generator of provider-specific chunks. """ @@ -141,5 +142,7 @@ async def _llm_router( provider=provider, config=config, **request_kwargs, ) - async for res_chunk in _retry_on_error(create_stream, max_retries, retry_delay, context): + async for res_chunk in _retry_on_error( + create_stream, max_retries, retry_delay, context, fail_fast_rate_limit=fail_fast_rate_limit, + ): yield res_chunk # type: ignore[return-type]