unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py (1)

254-317: Extract the sync retry loop instead of maintaining two copies.

retry_with_exponential_backoff() now reimplements the same attempt/sleep/decision flow as call_with_retry(), which is the SonarCloud failure on Line 254. Pulling the common loop into one helper will lower the branching here and keep the decorator path from drifting away from the SDK path.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 254 - 317,
The retry logic in retry_with_exponential_backoff duplicates the
attempt/sleep/decision flow from call_with_retry; extract that shared loop into
a single helper function (e.g., _run_retry_loop or reuse call_with_retry
signature) that encapsulates attempts, calling the target, invoking
_get_retry_delay, sleeping, logging, and raising, then have
retry_with_exponential_backoff's decorator/wrapper call this helper (passing
func, max_retries, base_delay, multiplier, jitter, exceptions, logger_instance,
prefix, retry_predicate) to eliminate the duplicated flow and keep both the
decorator path and call_with_retry using the same implementation.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py (2)

257-324: Consider extracting the inner wrapper function to reduce cognitive complexity.

SonarCloud flags cognitive complexity of 20 (allowed: 15). The nested decorator structure with try/except and conditional logging contributes to this. You could extract the "giving up" logging block or the entire wrapper body into a helper function.

Possible refactor approach

+def _log_final_failure(
+    logger_instance: logging.Logger,
+    func_name: str,
+    attempt: int,
+    prefix: str,
+) -> None:
+    """Log when retry attempts are exhausted."""
+    if attempt > 0:
+        logger_instance.exception(
+            "Giving up '%s' after %d attempt(s) for %s",
+            func_name,
+            attempt + 1,
+            prefix,
+        )
+
 def retry_with_exponential_backoff(...) -> Callable:
     ...
     def decorator(func: Callable) -> Callable:
         `@wraps`(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             for attempt in range(max_retries + 1):
                 try:
                     result = func(*args, **kwargs)
                     ...
                     return result
                 except exceptions as e:
                     delay = _get_retry_delay(...)
                     if delay is None:
-                        if attempt > 0:
-                            logger_instance.exception(
-                                "Giving up '%s' after %d attempt(s) for %s",
-                                func.__name__,
-                                attempt + 1,
-                                prefix,
-                            )
+                        _log_final_failure(logger_instance, func.__name__, attempt, prefix)
                         raise
                     time.sleep(delay)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 257 - 324,
The wrapper inside retry_with_exponential_backoff is too complex; extract the
retry loop or the "giving up" logging and sleep/raise logic into a helper
function to reduce cognitive complexity. Create a private helper (e.g.,
_execute_with_retries or _handle_attempt) that accepts func, args, kwargs,
max_retries, base_delay, multiplier, jitter, exceptions, retry_predicate,
prefix, and logger_instance and returns the result or raises; have wrapper call
this helper and keep wrapper only responsible for `@wraps` and invoking the
helper; ensure the helper uses _get_retry_delay and preserves the existing
logging messages (including the success log in wrapper when attempt > 0 or move
that logging into the helper) so behavior is unchanged.

---

138-145: Consider tightening type hints for fn parameter.

The current Callable[[], object] is imprecise:

• acall_with_retry awaits the result, so fn should return an awaitable
• iter_with_retry iterates the result, so fn should return an iterable

This won't cause runtime issues but weakens static type checking.

Suggested type hints

from collections.abc import Awaitable, Iterable
from typing import TypeVar

T = TypeVar("T")

async def acall_with_retry(
    fn: Callable[[], Awaitable[T]],
    *,
    max_retries: int,
    ...
) -> T:

def iter_with_retry(
    fn: Callable[[], Iterable[T]],
    *,
    max_retries: int,
    ...
) -> Generator[T, None, None]:

Also applies to: 160-167

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 138 - 145,
Tighten the type hints for acall_with_retry and iter_with_retry: introduce a
TypeVar T and change the fn parameter for acall_with_retry to Callable[[],
Awaitable[T]] with the function returning T, and change iter_with_retry to
accept Callable[[], Iterable[T]] and return Generator[T, None, None]; import
Awaitable and Iterable from collections.abc (and Generator if not already
present) and TypeVar from typing to make the signatures precise and improve
static checking while keeping existing runtime behavior (refer to function names
acall_with_retry and iter_with_retry and update their return annotations
accordingly).

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py (3)

266-333: Consider reducing cognitive complexity.

SonarCloud flagged this function with cognitive complexity of 20 (allowed 15). The nested conditionals within the exception handlers contribute to this. Consider extracting the "giving up" logging block (lines 319-325) into a helper or simplifying the control flow.

💡 Suggested refactor

+def _log_giving_up(
+    logger_instance: logging.Logger,
+    func_name: str,
+    attempt: int,
+    prefix: str,
+) -> None:
+    """Log when giving up after all retry attempts."""
+    if attempt > 0:
+        logger_instance.exception(
+            "Giving up '%s' after %d attempt(s) for %s",
+            func_name,
+            attempt + 1,
+            prefix,
+        )
+

 def retry_with_exponential_backoff(
     ...
 ) -> Callable:
     def decorator(func: Callable) -> Callable:
         `@wraps`(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             for attempt in range(max_retries + 1):
                 try:
                     result = func(*args, **kwargs)
                     if attempt > 0:
                         logger_instance.info(
                             "Successfully completed '%s' after %d retry attempt(s)",
                             func.__name__,
                             attempt,
                         )
                     return result
                 except exceptions as e:
                     delay = _get_retry_delay(...)
                     if delay is None:
-                        if attempt > 0:
-                            logger_instance.exception(
-                                "Giving up '%s' after %d attempt(s) for %s",
-                                func.__name__,
-                                attempt + 1,
-                                prefix,
-                            )
+                        _log_giving_up(logger_instance, func.__name__, attempt, prefix)
                         raise
                     time.sleep(delay)
                 except Exception:
                     raise
         return wrapper
     return decorator

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 266 - 333,
The retry_with_exponential_backoff wrapper has high cognitive complexity due to
nested exception handling and an inline "giving up" logging block; extract that
block into a small helper (e.g., _log_giving_up or _handle_give_up) and call it
from inside the except exceptions as e: branch when _get_retry_delay returns
None, passing logger_instance, prefix, func.__name__, and attempt so the main
wrapper focuses on control flow only, reducing nesting; keep the existing calls
to _get_retry_delay and time.sleep unchanged and ensure the helper performs the
logger.exception call and any formatting before re-raising the exception.

---

145-152: Type hint should reflect awaitable return type.

The fn parameter is awaited on line 158, so the type hint should be Callable[[], Awaitable[object]] rather than Callable[[], object].

💡 Suggested fix

+from collections.abc import Awaitable
+
 async def acall_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Awaitable[object]],
     *,
     max_retries: int,
     retry_predicate: Callable[[Exception], bool],
     description: str = "",
     logger_instance: logging.Logger | None = None,
 ) -> object:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 145 - 152,
Update the type hint for the `fn` parameter in `acall_with_retry` to reflect
that it's awaited: change `Callable[[], object]` to `Callable[[],
Awaitable[object]]`, and add the necessary `Awaitable` import from `typing` (or
`collections.abc`) at the top of the module so the annotation is valid; keep the
function's async signature and return annotation as-is.

---

168-175: Type hint should reflect iterable return type.

The fn result is iterated on line 186, so the type hint should reflect this, e.g., Callable[[], Iterable[T]].

💡 Suggested fix

+from collections.abc import Iterable
+from typing import TypeVar
+
+T = TypeVar("T")
+
 def iter_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Iterable[T]],
     *,
     max_retries: int,
     retry_predicate: Callable[[Exception], bool],
     description: str = "",
     logger_instance: logging.Logger | None = None,
-) -> Generator:
+) -> Generator[T, None, None]:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 168 - 175,
The type hint for iter_with_retry is wrong because fn is expected to return an
iterable; update the signature to use a generic TypeVar (e.g., T) and change
fn's type to Callable[[], Iterable[T]] and the function's return type to
Generator[T, None, None] so callers of iter_with_retry and the loop over fn()
have correct typing; ensure you import or declare typing.TypeVar and Iterable
where needed and update any related annotations in iter_with_retry accordingly.

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py (2)

292-328: Optional: trim one point of cognitive complexity to satisfy SonarCloud (16 → 15).

SonarCloud reports 16 vs. threshold 15. The easiest reduction is extracting the give-up logging into a small helper so wrapper has one less nested branch; alternatively, just raise the threshold if the team prefers. Not blocking.

♻️ Example extraction

+    def _log_giveup(func: Callable, attempt: int) -> None:
+        if attempt > 0:
+            logger_instance.exception(
+                "Giving up '%s' after %d attempt(s) for %s",
+                func.__name__,
+                attempt + 1,
+                prefix,
+            )
+
     def decorator(func: Callable) -> Callable:
         `@wraps`(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
             for attempt in range(max_retries + 1):
                 try:
                     result = func(*args, **kwargs)
                 except exceptions as e:
                     delay = _get_retry_delay(
                         e, attempt, max_retries, retry_predicate, prefix,
                         logger_instance, base_delay, multiplier, 60.0, jitter,
                     )
                     if delay is not None:
                         time.sleep(delay)
                         continue
-                    if attempt > 0:
-                        logger_instance.exception(
-                            "Giving up '%s' after %d attempt(s) for %s",
-                            func.__name__,
-                            attempt + 1,
-                            prefix,
-                        )
+                    _log_giveup(func, attempt)
                     raise

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 292 - 328,
The wrapper function inside decorator has one extra nested branch for the "give
up" logging; extract that logic into a small helper function (e.g., _log_give_up
or _log_and_raise) referenced from decorator/wrapper so wrapper no longer
contains the nested if that calls logger_instance.exception and raise; call the
helper with parameters (logger_instance, func.__name__, attempt, prefix) after
_get_retry_delay returns None and before re-raising the exception, and keep
existing behavior of logging only when attempt > 0 and then re-raising the
exception.

---

145-175: Tighten fn type hints for async and generator variants.

acall_with_retry awaits fn() and iter_with_retry iterates over fn(), but both declare fn: Callable[[], object]. Type checkers will not catch a caller that passes a plain sync function into acall_with_retry, or a non-iterable factory into iter_with_retry. Recommend narrowing:

♻️ Suggested tightening

-from collections.abc import Callable, Generator
+from collections.abc import Awaitable, Callable, Generator, Iterable
@@
 async def acall_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Awaitable[object]],
     *,
@@
 def iter_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Iterable[object]],
     *,

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 145 - 175,
The signatures for acall_with_retry and iter_with_retry should be made generic
and their fn arguments narrowed: introduce a TypeVar T and change
acall_with_retry to async def acall_with_retry(fn: Callable[[], Awaitable[T]],
..., ) -> T so the checker knows fn is awaitable, and change iter_with_retry to
def iter_with_retry(fn: Callable[[], Iterable[T]], ..., ) -> Generator[T, None,
None] (or Callable[[], Iterator[T]] if fn yields an iterator) so the checker
knows fn returns an iterable; update imports (TypeVar, Awaitable,
Iterable/Iterator) and adjust the return type annotations accordingly while
keeping the existing behavior in the bodies of acall_with_retry and
iter_with_retry.

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py (3)

87-97: Minor: debug log counter vs. warning counter use different denominators.

The debug line logs attempt+1/max_retries+1 (total attempts including initial try), while the warning on Line 101-103 logs attempt+1/max_retries (retries only). When reading logs for the same failure, these will show different /N totals and be confusing. Additionally, on the last attempt (attempt == max_retries) the debug line still prints retryable=True even though the function returns None and the caller re-raises — the log reader sees “retryable=True” for a call that was not retried.

Consider either unifying the denominators or appending a "giving up" branch before returning None:

Suggested tweak

-    if not should_retry or attempt >= max_retries:
-        return None
+    if not should_retry or attempt >= max_retries:
+        logger_instance.debug(
+            "Giving up: attempts exhausted or non-retryable (attempt=%d, max_retries=%d, description=%s)",
+            attempt + 1, max_retries, description,
+        )
+        return None

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 87 - 97,
The debug and warning logs in the retry decision flow are inconsistent and can
mislead: update the logging in the function that calls logger_instance.debug
(use symbols attempt, max_retries, should_retry, description) so both messages
use the same denominator (either total attempts = max_retries+1 or retries =
max_retries), and add a conditional "giving up" log path just before returning
None when not should_retry or attempt >= max_retries so the message shows
retryable=False (or explicitly states giving up) for the final non-retry case;
ensure you adjust the debug message's retryable value to reflect the actual
branch taken.

---

176-198: Docstring: clarify that fn must return a fresh iterable on each call.

Each retry re-invokes fn(), which assumes fn is a factory (e.g., lambda: litellm.completion(stream=True, ...)). Passing an already-materialized generator would be silently wrong: after the first failure the exhausted generator would either yield nothing or raise StopIteration/the same error. Worth a one-liner in the docstring so callers don't pass a bare iterator.

Suggested wording

-    """Yield from fn() with retry. Only retries before the first yield.
-
-    Once items have been yielded to the caller a mid-iteration failure is
-    raised immediately — partial output can't be un-yielded.
-    """
+    """Yield from fn() with retry. Only retries before the first yield.
+
+    `fn` must be a factory that returns a *fresh* iterable on each call
+    (e.g., ``lambda: litellm.completion(stream=True, ...)``). A bare
+    generator must not be passed, since retrying would re-consume an
+    already-exhausted iterator.
+
+    Once items have been yielded to the caller a mid-iteration failure is
+    raised immediately — partial output can't be un-yielded.
+    """

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 176 - 198,
Update the function docstring to explicitly state that the parameter fn must be
a zero-argument factory that returns a fresh iterable each time it is called
(i.e., pass a callable like lambda: generator() rather than an
already-materialized iterator), because the code re-invokes fn() on each retry;
reference the use of fn() and the retry loop (for attempt in range(max_retries +
1)) and the has_yielded logic so callers understand why a fresh iterable is
required.

---

145-165: Type hints: fn parameters are incorrectly typed as Callable[[], object] but actual usage differs.

• acall_with_retry awaits the return value (await fn()), so fn should be typed as Callable[[], Awaitable[object]]
• iter_with_retry yields from the return value, so fn should be typed as Callable[[], Iterable[object]]

Current imports at line 10 only include Callable and Generator. Add Awaitable and Iterable to fix type checker compatibility.

Suggested fix

-from collections.abc import Callable, Generator
+from collections.abc import Awaitable, Callable, Generator, Iterable
@@
 async def acall_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Awaitable[object]],
     *,
@@
 def iter_with_retry(
-    fn: Callable[[], object],
+    fn: Callable[[], Iterable[object]],
     *,

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py` around lines 145 - 165,
The type hints for the retry utilities are incorrect: update the imports to
include Awaitable and Iterable, then change acall_with_retry's fn annotation
from Callable[[], object] to Callable[[], Awaitable[object]] and change
iter_with_retry's fn annotation from Callable[[], object] to Callable[[],
Iterable[object]] so the declared callable return types match actual usage;
locate and update the import list (where Callable and Generator are imported)
and the fn parameter annotations in acall_with_retry and iter_with_retry
accordingly.