Non-destructive retrying on nemotron-parse API

packages/paper-qa-nemotron/src/paperqa_nemotron/api.py

@@ -26,6 +26,7 @@
     cast,
     overload,
 )
+from unittest.mock import patch
 
 import litellm
 from aviary.core import Message, ToolCall
@@ -39,8 +40,10 @@
     ValidationInfo,
 )
 from tenacity import (
+    RetryCallState,
     before_sleep_log,
     retry,
+    retry_any,
     retry_if_exception,
     retry_if_exception_type,
     stop_after_attempt,
@@ -320,11 +323,33 @@ def merge_with_detection(
 
 def _is_litellm_timeout_with_408(exc: BaseException) -> bool:
     return (
-        isinstance(exc, litellm.exceptions.Timeout)
+        isinstance(exc, litellm.Timeout)
         and exc.status_code == http.HTTPStatus.REQUEST_TIMEOUT
     )
 
 
+# Exponential backoff for when hitting rate limits on Nvidia's API
+_NVIDIA_API_RETRY_WAIT = wait_exponential(multiplier=2, min=GLOBAL_RATE_LIMITER_TIMEOUT)
+
+
+def _wait_exponential_for_nvidia_api_retry(retry_state: RetryCallState) -> float:
+    """Only apply exponential backoff for rate limit failure mode.
+
+    Uses a Nvidia API rate limit-specific counter so backoff is based on consecutive
+    rate limit failures, not overall attempt number.
+    """
+    if retry_state.outcome is not None and isinstance(
+        retry_state.outcome.exception(), TimeoutError
+    ):
+        # Track TimeoutError (Nvidia API rate limit) count separately
+        timeout_count: int = getattr(retry_state, "_timeout_error_count", 0) + 1
+        retry_state._timeout_error_count = timeout_count  # type: ignore[attr-defined]
+        # Temporarily override attempt_number for wait_exponential calculation
+        with patch.object(retry_state, "attempt_number", timeout_count):
+            return _NVIDIA_API_RETRY_WAIT(retry_state)
+    return 0
+
+
 @overload
 async def _call_nvidia_api(
     image: "np.ndarray",
@@ -359,19 +384,13 @@ async def _call_nvidia_api(
 
 
 @retry(
-    retry=retry_if_exception_type(NemotronBBoxError),
-    stop=stop_after_attempt(3),
-    before_sleep=before_sleep_log(logger, logging.WARNING),
-)
-@retry(
-    retry=retry_if_exception_type(TimeoutError),  # Hitting rate limits
-    stop=stop_after_attempt(3),
-    wait=wait_exponential(multiplier=2, min=GLOBAL_RATE_LIMITER_TIMEOUT),
-    before_sleep=before_sleep_log(logger, logging.WARNING),
-)
-@retry(
-    retry=retry_if_exception(_is_litellm_timeout_with_408),  # Inference timeout
+    retry=retry_any(
+        retry_if_exception_type(NemotronBBoxError),
+        retry_if_exception_type(TimeoutError),  # Hitting rate limits
+        retry_if_exception(_is_litellm_timeout_with_408),  # Inference timeout
+    ),
     stop=stop_after_attempt(3),
+    wait=_wait_exponential_for_nvidia_api_retry,
     before_sleep=before_sleep_log(logger, logging.WARNING),
 )
 async def _call_nvidia_api(

packages/paper-qa-nemotron/tests/test_api.py

@@ -1,9 +1,12 @@
+import http
 from pathlib import Path
 
+import litellm
 import numpy as np
 import pypdfium2 as pdfium
 import pytest
 from pydantic import ValidationError
+from tenacity import Future, RetryCallState
 
 from paperqa_nemotron.api import (
     NemotronParseAnnotatedBBox,
@@ -13,6 +16,7 @@
     NemotronParseMarkdownBBox,
     _call_nvidia_api,
     _call_sagemaker_api,
+    _wait_exponential_for_nvidia_api_retry,
 )
 
 REPO_ROOT = Path(__file__).parents[3]
@@ -270,6 +274,44 @@ def test_merge_multiple_items(self) -> None:
         assert merged[2].text == "Table content"
 
 
+@pytest.mark.parametrize(
+    ("exc", "expected_backoff"),
+    [
+        pytest.param(TimeoutError(), True, id="TimeoutError-has-backoff"),
+        pytest.param(
+            litellm.Timeout(
+                message="Request timed out",
+                model="stub",
+                llm_provider=litellm.LlmProviders.CUSTOM,
+                exception_status_code=http.HTTPStatus.REQUEST_TIMEOUT,
+            ),
+            False,
+            id="litellm-Timeout-no-backoff",
+        ),
+    ],
+)
+def test_wait_exponential_for_nvidia_api_retry(
+    exc: BaseException, expected_backoff: bool
+) -> None:
+    """Test we properly handle exponential backoff for Nvidia API."""
+    retry_state = RetryCallState(
+        retry_object=None,  # type: ignore[arg-type]
+        fn=None,
+        args=(),
+        kwargs={},
+    )
+    retry_state.outcome = Future.construct(
+        attempt_number=1, value=exc, has_exception=True
+    )
+    retry_state.attempt_number = 1
+
+    wait_time = _wait_exponential_for_nvidia_api_retry(retry_state)
+    if expected_backoff:
+        assert wait_time > 0
+    else:
+        assert wait_time == 0
+
+
 @pytest.fixture(name="pdf_page_np")
 def fixture_pdf_page_np() -> np.ndarray:
     pdf_doc = pdfium.PdfDocument(STUB_DATA_DIR / "pasa.pdf")