fix(autofix) Retry stream from scratch (#1748)

[Investigation](https://github.com/getsentry/ml-models/blob/main/autofix/retry_stream/investigate_errors.ipynb)
found that backoff would address almost all of the overloaded errors.

Fix #1671 by retrying the completion from scratch, unlike #1675.

## How has this been tested?

New unit tests and local runs described below

<details>
<summary>Local autofix runs</summary>

Repro issue locally:

1. `git checkout main`
2. Replace all `AnthropicProvider`s with `AnthropicProviderFlaky`s in
`coding/components.py`.
3. Run the Autofix coding step on a local issue, confirm you see the
overloaded error and the run fails.
4. `git stash` the stuff in step 2 b/c we'll use it again when testing
the fix.

Test fix:

1. `git checkout kddubey/autofix/retry-stream-from-scratch`
2. `git stash pop` to replace all `AnthropicProvider`s with
`AnthropicProviderFlaky`s in `coding/components.py`.
3. Run Autofix on a local issue, confirm you see the overloaded error in
your logs, but it quickly tries again and finishes the run.
4. `git stash` to undo step (2) to test an unflaky API
5. Run Autofix on the same issue, confirm the run finishes fine.

</summary>

Author: kddubey

src/seer/automation/agent/client.py

@@ -90,6 +90,10 @@ def _get_config(cls, model_name: str):
                 return config
         return None
 
+    @staticmethod
+    def is_completion_exception_retryable(exception: Exception) -> bool:
+        return False
+
     def generate_text(
         self,
         *,
@@ -408,6 +412,15 @@ def _get_config(cls, model_name: str):
                 return config
         return None
 
+    @staticmethod
+    def is_completion_exception_retryable(exception: Exception) -> bool:
+        if isinstance(exception, anthropic.APIStatusError):
+            return exception.status_code == 529
+            # https://docs.anthropic.com/en/api/errors#http-errors
+        return isinstance(exception, anthropic.AnthropicError) and (
+            "overloaded_error" in str(exception)
+        )
+
     @observe(as_type="generation", name="Anthropic Generation")
     @inject
     def generate_text(
@@ -708,6 +721,10 @@ def search_the_web(self, prompt: str, temperature: float | None = None) -> str:
             answer += each.text
         return answer
 
+    @staticmethod
+    def is_completion_exception_retryable(exception: Exception) -> bool:
+        return False
+
     @observe(as_type="generation", name="Gemini Generation")
     def generate_structured(
         self,

src/seer/automation/autofix/autofix_agent.py

@@ -1,7 +1,7 @@
 import contextlib
 import logging
 from concurrent.futures import Executor, Future, ThreadPoolExecutor
-from typing import Optional
+from typing import Callable, Optional
 
 from seer.automation.agent.agent import AgentConfig, LlmAgent, RunConfig
 from seer.automation.agent.models import (
@@ -17,6 +17,7 @@
 from seer.automation.autofix.models import AutofixContinuation, AutofixStatus, DefaultStep
 from seer.automation.state import State
 from seer.dependency_injection import copy_modules_initializer
+from seer.utils import backoff_on_exception
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +70,7 @@ def _check_prompt_for_help(self, run_config: RunConfig):
                 "You're taking a while. If you need help, ask me a concrete question using the tool provided."
             )
 
-    def get_completion(self, run_config: RunConfig):
+    def _get_completion(self, run_config: RunConfig):
         """
         Streams the preliminary output to the current step and only returns when output is complete
         """
@@ -118,6 +119,27 @@ def get_completion(self, run_config: RunConfig):
             ),
         )
 
+    def get_completion(
+        self,
+        run_config: RunConfig,
+        max_tries: int = 4,
+        sleep_sec_scaler: Callable[[int], float] = lambda num_tries: 2**num_tries,
+    ):
+        """
+        Streams the preliminary output to the current step and only returns when output is complete.
+
+        The completion request is retried `max_tries - 1` times if a retryable exception was just
+        raised, e.g, Anthropic's API is overloaded.
+        """
+        is_exception_retryable = getattr(
+            run_config.model, "is_completion_exception_retryable", lambda _: False
+        )
+        retrier = backoff_on_exception(
+            is_exception_retryable, max_tries=max_tries, sleep_sec_scaler=sleep_sec_scaler
+        )
+        get_completion_retryable = retrier(self._get_completion)
+        return get_completion_retryable(run_config)
+
     def run_iteration(self, run_config: RunConfig):
         logger.debug(f"----[{self.name}] Running Iteration {self.iterations}----")
 

src/seer/utils.py

@@ -1,13 +1,18 @@
 import contextlib
 import functools
 import json
+import logging
+import random
+import time
 import weakref
 from enum import Enum
 from queue import Empty, Full, Queue
-from typing import Sequence
+from typing import Callable, Sequence
 
 from sqlalchemy.orm import DeclarativeBase, Session
 
+logger = logging.getLogger(__name__)
+
 
 def class_method_lru_cache(*lru_args, **lru_kwargs):
     def decorator(func):
@@ -68,3 +73,44 @@ def closing_queue(*queues: Queue):
                 queue.get_nowait()
             except Empty:
                 pass
+
+
+def backoff_on_exception(
+    is_exception_retryable: Callable[[Exception], bool],
+    max_tries: int = 2,
+    sleep_sec_scaler: Callable[[int], float] = lambda num_tries: 2**num_tries,
+    jitterer: Callable[[], float] = lambda: random.uniform(0, 0.5),
+):
+    """
+    Returns a decorator which retries a function on exception iff `is_exception_retryable(exception)`.
+    Defaults to exponential backoff with random jitter and one retry.
+    """
+
+    if max_tries < 1:
+        raise ValueError("max_tries must be at least 1")  # pragma: no cover
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapped_func(*args, **kwargs):
+            num_tries = 0
+            last_exception = None
+            while num_tries < max_tries:
+                try:
+                    return func(*args, **kwargs)
+                except Exception as exception:
+                    num_tries += 1
+                    last_exception = exception
+                    if is_exception_retryable(exception):
+                        sleep_sec = sleep_sec_scaler(num_tries) + jitterer()
+                        logger.info(
+                            f"Encountered {type(exception).__name__}: {exception}. Sleeping for "
+                            f"{sleep_sec} seconds before attempting retry {num_tries}/{max_tries}."
+                        )
+                        time.sleep(sleep_sec)
+                    else:
+                        raise exception
+            raise last_exception
+
+        return wrapped_func
+
+    return decorator