openrouter-python-sdk-retry-errors/src/openrouter/utils/retries.py

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

import asyncio
import json
import random
import time
from datetime import datetime
from email.utils import parsedate_to_datetime
from typing import Any, List, Optional

import httpx


class BackoffStrategy:
    initial_interval: int
    max_interval: int
    exponent: float
    max_elapsed_time: int

    def __init__(
        self,
        initial_interval: int,
        max_interval: int,
        exponent: float,
        max_elapsed_time: int,
    ):
        self.initial_interval = initial_interval
        self.max_interval = max_interval
        self.exponent = exponent
        self.max_elapsed_time = max_elapsed_time


class RetryConfig:
    strategy: str
    backoff: BackoffStrategy
    retry_connection_errors: bool

    def __init__(
        self, strategy: str, backoff: BackoffStrategy, retry_connection_errors: bool
    ):
        self.strategy = strategy
        self.backoff = backoff
        self.retry_connection_errors = retry_connection_errors


class Retries:
    config: RetryConfig
    status_codes: List[str]

    def __init__(self, config: RetryConfig, status_codes: List[str]):
        self.config = config
        self.status_codes = status_codes


class TemporaryError(Exception):
    response: httpx.Response
    retry_after: Optional[int]

    def __init__(self, response: httpx.Response):
        self.response = response
        self.retry_after = _parse_retry_after_header(response)


class PermanentError(Exception):
    inner: Exception

    def __init__(self, inner: Exception):
        self.inner = inner


def _parse_retry_after_header(response: httpx.Response) -> Optional[int]:
    """Parse Retry-After header from response.

    Returns:
        Retry interval in milliseconds, or None if header is missing or invalid.
    """
    retry_after_header = response.headers.get("retry-after")
    if not retry_after_header:
        return None

    try:
        seconds = float(retry_after_header)
        return round(seconds * 1000)
    except ValueError:
        pass

    try:
        retry_date = parsedate_to_datetime(retry_after_header)
        delta = (retry_date - datetime.now(retry_date.tzinfo)).total_seconds()
        return round(max(0, delta) * 1000)
    except (ValueError, TypeError):
        pass

    return None


TRANSIENT_STATUS_CODES = {408, 429}
TRANSIENT_ERROR_CODES = {408, 429, 500, 502, 503, 504, 524, 529}


def _matches_retry_status_code(response: httpx.Response, status_code: str) -> bool:
    if "X" in status_code.upper():
        code_range = int(status_code[0])
        status_major = response.status_code // 100
        return code_range == status_major

    return response.status_code == int(status_code)


def _error_code_from_json_body(response: httpx.Response) -> Optional[int]:
    try:
        body: Any = response.json()
    except json.JSONDecodeError:
        return None

    if not isinstance(body, dict):
        return None

    error = body.get("error")
    if not isinstance(error, dict):
        return None

    code = error.get("code")
    if isinstance(code, int):
        return code
    if isinstance(code, str) and code.isdigit():
        return int(code)

    return None


def _is_retryable_response(response: httpx.Response, status_codes: List[str]) -> bool:
    if response.status_code in TRANSIENT_STATUS_CODES:
        return True

    for status_code in status_codes:
        if _matches_retry_status_code(response, status_code):
            return True

    if response.status_code != 400:
        return False

    if not response.is_closed:
        response.read()

    inner_code = _error_code_from_json_body(response)
    return inner_code in TRANSIENT_ERROR_CODES


async def _is_retryable_response_async(
    response: httpx.Response, status_codes: List[str]
) -> bool:
    if response.status_code in TRANSIENT_STATUS_CODES:
        return True

    for status_code in status_codes:
        if _matches_retry_status_code(response, status_code):
            return True

    if response.status_code != 400:
        return False

    if not response.is_closed:
        await response.aread()

    inner_code = _error_code_from_json_body(response)
    return inner_code in TRANSIENT_ERROR_CODES


def _get_sleep_interval(
    exception: Exception,
    initial_interval: int,
    max_interval: int,
    exponent: float,
    retries: int,
) -> float:
    """Get sleep interval for retry with exponential backoff.

    Args:
        exception: The exception that triggered the retry.
        initial_interval: Initial retry interval in milliseconds.
        max_interval: Maximum retry interval in milliseconds.
        exponent: Base for exponential backoff calculation.
        retries: Current retry attempt count.

    Returns:
        Sleep interval in seconds.
    """
    if (
        isinstance(exception, TemporaryError)
        and exception.retry_after is not None
        and exception.retry_after > 0
    ):
        return exception.retry_after / 1000

    sleep = (initial_interval / 1000) * exponent**retries + random.uniform(0, 1)
    return min(sleep, max_interval / 1000)


def retry(func, retries: Retries):
    if retries.config.strategy == "backoff":

        def do_request() -> httpx.Response:
            res: httpx.Response
            try:
                res = func()

                if _is_retryable_response(res, retries.status_codes):
                    raise TemporaryError(res)
            except httpx.ConnectError as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except httpx.TimeoutException as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except TemporaryError:
                raise
            except Exception as exception:
                raise PermanentError(exception) from exception

            return res

        return retry_with_backoff(
            do_request,
            retries.config.backoff.initial_interval,
            retries.config.backoff.max_interval,
            retries.config.backoff.exponent,
            retries.config.backoff.max_elapsed_time,
        )

    return func()


async def retry_async(func, retries: Retries):
    if retries.config.strategy == "backoff":

        async def do_request() -> httpx.Response:
            res: httpx.Response
            try:
                res = await func()

                if await _is_retryable_response_async(res, retries.status_codes):
                    raise TemporaryError(res)
            except httpx.ConnectError as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except httpx.TimeoutException as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except TemporaryError:
                raise
            except Exception as exception:
                raise PermanentError(exception) from exception

            return res

        return await retry_with_backoff_async(
            do_request,
            retries.config.backoff.initial_interval,
            retries.config.backoff.max_interval,
            retries.config.backoff.exponent,
            retries.config.backoff.max_elapsed_time,
        )

    return await func()


def retry_with_backoff(
    func,
    initial_interval=500,
    max_interval=60000,
    exponent=1.5,
    max_elapsed_time=3600000,
):
    start = round(time.time() * 1000)
    retries = 0

    while True:
        try:
            return func()
        except PermanentError as exception:
            raise exception.inner
        except Exception as exception:  # pylint: disable=broad-exception-caught
            now = round(time.time() * 1000)
            if now - start > max_elapsed_time:
                if isinstance(exception, TemporaryError):
                    return exception.response

                raise

            sleep = _get_sleep_interval(
                exception, initial_interval, max_interval, exponent, retries
            )
            time.sleep(sleep)
            retries += 1


async def retry_with_backoff_async(
    func,
    initial_interval=500,
    max_interval=60000,
    exponent=1.5,
    max_elapsed_time=3600000,
):
    start = round(time.time() * 1000)
    retries = 0

    while True:
        try:
            return await func()
        except PermanentError as exception:
            raise exception.inner
        except Exception as exception:  # pylint: disable=broad-exception-caught
            now = round(time.time() * 1000)
            if now - start > max_elapsed_time:
                if isinstance(exception, TemporaryError):
                    return exception.response

                raise

            sleep = _get_sleep_interval(
                exception, initial_interval, max_interval, exponent, retries
            )
            await asyncio.sleep(sleep)
            retries += 1