Coverage for llm_dataset_engine/utils/retry_handler.py: 96%

1""" 

2Retry handling with exponential backoff. 

3 

4Provides robust retry logic for transient failures following clean code 

5principles with single responsibility. 

6""" 

7 

8import time 

9from typing import Callable, Optional, Type, TypeVar 

10 

11from tenacity import ( 

12 Retrying, 

13 retry_if_exception_type, 

14 stop_after_attempt, 

15 wait_exponential, 

16) 

17 

18T = TypeVar("T") 

19 

20 

21class RetryableError(Exception): 

22 """Base class for errors that should be retried.""" 

23 

24 pass 

25 

26 

27class RateLimitError(RetryableError): 

28 """Rate limit exceeded error.""" 

29 

30 pass 

31 

32 

33class NetworkError(RetryableError): 

34 """Network-related error.""" 

35 

36 pass 

37 

38 

39class RetryHandler: 

40 """ 

41 Handles retry logic with exponential backoff. 

42  

43 Follows SOLID principles with single responsibility for retry logic. 

44 """ 

45 

46 def __init__( 

47 self, 

48 max_attempts: int = 3, 

49 initial_delay: float = 1.0, 

50 max_delay: float = 60.0, 

51 exponential_base: int = 2, 

52 retryable_exceptions: Optional[tuple[Type[Exception], ...]] = None, 

53 ): 

54 """ 

55 Initialize retry handler. 

56 

57 Args: 

58 max_attempts: Maximum retry attempts 

59 initial_delay: Initial delay in seconds 

60 max_delay: Maximum delay in seconds 

61 exponential_base: Base for exponential backoff 

62 retryable_exceptions: Exception types to retry 

63 """ 

64 self.max_attempts = max_attempts 

65 self.initial_delay = initial_delay 

66 self.max_delay = max_delay 

67 self.exponential_base = exponential_base 

68 

69 if retryable_exceptions is None: 

70 self.retryable_exceptions = ( 

71 RetryableError, 

72 RateLimitError, 

73 NetworkError, 

74 ) 

75 else: 

76 self.retryable_exceptions = retryable_exceptions 

77 

78 def execute(self, func: Callable[[], T]) -> T: 

79 """ 

80 Execute function with retry logic. 

81 

82 Args: 

83 func: Function to execute 

84 

85 Returns: 

86 Result from function 

87 

88 Raises: 

89 Exception: If all retries exhausted 

90 """ 

91 retryer = Retrying( 

92 stop=stop_after_attempt(self.max_attempts), 

93 wait=wait_exponential( 

94 multiplier=self.initial_delay, 

95 max=self.max_delay, 

96 exp_base=self.exponential_base, 

97 ), 

98 retry=retry_if_exception_type(self.retryable_exceptions), 

99 reraise=True, 

100 ) 

101 

102 return retryer(func) 

103 

104 def calculate_delay(self, attempt: int) -> float: 

105 """ 

106 Calculate delay for given attempt number. 

107 

108 Args: 

109 attempt: Attempt number (1-based) 

110 

111 Returns: 

112 Delay in seconds 

113 """ 

114 delay = self.initial_delay * (self.exponential_base ** (attempt - 1)) 

115 return min(delay, self.max_delay) 

116