hanzoai_agent/sdk/python/hanzo_agents/rate_limiter.py

import asyncio
import hashlib
import os
import random
import time
from typing import Any, Optional
from hanzo_agents.logger import log_debug


class RateLimitError(Exception):
    """Custom exception for rate limit errors"""

    def __init__(self, message: str, retry_after: Optional[float] = None):
        super().__init__(message)
        self.retry_after = retry_after


class StatelessRateLimiter:
    """
    Stateless rate limiter with adaptive exponential backoff.

    Designed to work across hundreds of containers without coordination.
    Uses container-specific jitter to naturally distribute load.
    """

    def __init__(
        self,
        max_retries: int = 20,
        base_delay: float = 1.0,
        max_delay: float = 300.0,
        jitter_factor: float = 0.25,
        circuit_breaker_threshold: int = 10,
        circuit_breaker_timeout: int = 300,
    ):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.jitter_factor = jitter_factor
        self.circuit_breaker_threshold = circuit_breaker_threshold
        self.circuit_breaker_timeout = circuit_breaker_timeout

        # Container-specific seed for consistent but distributed jitter
        self._container_seed = self._get_container_seed()

        # Circuit breaker state (per-instance)
        self._consecutive_failures = 0
        self._circuit_open_time = None

    def _get_container_seed(self) -> int:
        """Generate a container-specific seed for consistent jitter distribution"""
        # Use hostname, process ID, and other container-specific identifiers
        identifier = f"{os.getenv('HOSTNAME', 'localhost')}-{os.getpid()}"
        return int(hashlib.md5(identifier.encode()).hexdigest()[:8], 16)

    def _is_rate_limit_error(self, error: Exception) -> bool:
        """
        Universal rate limit error detection for any LLM provider.

        Args:
            error: Exception to check

        Returns:
            bool: True if this is a rate limit error
        """
        # Check for specific LLM rate limit error
        if hasattr(error, "__class__") and "RateLimitError" in str(error.__class__):
            return True

        # Check HTTP status codes
        if hasattr(error, "response"):
            if hasattr(error.response, "status_code"):
                if error.response.status_code in [429, 503]:
                    return True

        # Check for HTTP status in error attributes
        if hasattr(error, "status_code"):
            if error.status_code in [429, 503]:
                return True

        # Check error message for rate limit keywords
        error_message = str(error).lower()
        rate_limit_keywords = [
            "rate limit",
            "rate-limit",
            "rate_limit",
            "too many requests",
            "quota exceeded",
            "temporarily rate-limited",
            "rate limited",
            "requests per",
            "rpm exceeded",
            "tpm exceeded",
            "usage limit",
            "throttled",
            "throttling",
        ]

        return any(keyword in error_message for keyword in rate_limit_keywords)

    def _extract_retry_after(self, error: Exception) -> Optional[float]:
        """
        Extract retry-after value from error if available.

        Args:
            error: Exception that may contain retry-after information

        Returns:
            Optional[float]: Retry-after seconds if found
        """
        # Check for Retry-After header in HTTP response
        if hasattr(error, "response") and hasattr(error.response, "headers"):
            retry_after = error.response.headers.get("Retry-After")
            if retry_after:
                try:
                    return float(retry_after)
                except ValueError:
                    pass

        # Check for retry_after in error attributes
        if hasattr(error, "retry_after"):
            try:
                return float(error.retry_after)
            except (ValueError, TypeError):
                pass

        return None

    def _calculate_backoff_delay(
        self, attempt: int, retry_after: Optional[float] = None
    ) -> float:
        """
        Calculate backoff delay with exponential backoff and jitter.

        Args:
            attempt: Current attempt number (0-based)
            retry_after: Server-suggested retry delay

        Returns:
            float: Delay in seconds
        """
        # Use server-suggested delay if available and reasonable
        if retry_after and retry_after <= self.max_delay:
            base_delay = retry_after
        else:
            # Exponential backoff: base_delay * (2 ^ attempt)
            base_delay = min(self.base_delay * (2**attempt), self.max_delay)

        # Add container-specific jitter to distribute load
        # Use container seed to ensure consistent but distributed jitter
        random.seed(self._container_seed + attempt)
        jitter_range = base_delay * self.jitter_factor
        jitter = random.uniform(-jitter_range, jitter_range)

        # Ensure minimum delay and apply jitter
        delay = max(0.1, base_delay + jitter)

        log_debug(
            f"Rate limit backoff: attempt={attempt}, base_delay={base_delay:.2f}s, jitter={jitter:.2f}s, total_delay={delay:.2f}s"
        )

        return delay

    def _check_circuit_breaker(self) -> bool:
        """
        Check if circuit breaker is open.

        Returns:
            bool: True if circuit is open (should not retry)
        """
        if self._circuit_open_time is None:
            return False

        # Check if circuit breaker timeout has passed
        if time.time() - self._circuit_open_time > self.circuit_breaker_timeout:
            log_debug("Circuit breaker timeout passed, attempting to close circuit")
            self._circuit_open_time = None
            self._consecutive_failures = 0
            return False

        return True

    def _update_circuit_breaker(self, success: bool):
        """
        Update circuit breaker state based on operation result.

        Args:
            success: Whether the operation succeeded
        """
        if success:
            # Reset on success
            self._consecutive_failures = 0
            if self._circuit_open_time:
                log_debug("Circuit breaker closed after successful request")
                self._circuit_open_time = None
        else:
            # Increment failures
            self._consecutive_failures += 1

            # Open circuit if threshold reached
            if (
                self._consecutive_failures >= self.circuit_breaker_threshold
                and self._circuit_open_time is None
            ):
                self._circuit_open_time = time.time()
                log_debug(
                    f"Circuit breaker opened after {self._consecutive_failures} consecutive failures"
                )

    async def execute_with_retry(self, func, *args, **kwargs) -> Any:
        """
        Execute a function with rate limit retry logic.

        Args:
            func: Async function to execute
            *args: Positional arguments for func
            **kwargs: Keyword arguments for func

        Returns:
            Any: Result of successful function execution

        Raises:
            RateLimitError: If max retries exceeded or circuit breaker is open
            Exception: Original exception if not rate limit related
        """
        # Check circuit breaker
        if self._check_circuit_breaker():
            raise RateLimitError(
                f"Circuit breaker is open. Too many consecutive rate limit failures. "
                f"Will retry after {self.circuit_breaker_timeout} seconds."
            )

        last_error = None

        for attempt in range(self.max_retries + 1):
            try:
                # Execute the function
                result = await func(*args, **kwargs)

                # Success - update circuit breaker and return
                self._update_circuit_breaker(success=True)

                if attempt > 0:
                    log_debug(f"Rate limit retry succeeded on attempt {attempt + 1}")

                return result

            except Exception as error:
                last_error = error

                # Check if this is a rate limit error
                if not self._is_rate_limit_error(error):
                    # Not a rate limit error - re-raise immediately
                    raise error

                # Update circuit breaker for rate limit failure
                self._update_circuit_breaker(success=False)

                # Check if we've exceeded max retries
                if attempt >= self.max_retries:
                    log_debug(f"Rate limit max retries ({self.max_retries}) exceeded")
                    break

                # Extract retry-after if available
                retry_after = self._extract_retry_after(error)

                # Calculate backoff delay
                delay = self._calculate_backoff_delay(attempt, retry_after)

                log_debug(
                    f"Rate limit detected on attempt {attempt + 1}, retrying in {delay:.2f}s. Error: {str(error)[:100]}"
                )

                # Wait before retry
                await asyncio.sleep(delay)

        # All retries exhausted
        raise RateLimitError(
            f"Rate limit retries exhausted after {self.max_retries} attempts. "
            f"Last error: {str(last_error)}"
        )