Coverage for llm_dataset_engine/utils/rate_limiter.py: 100%

1"""

2 Token bucket rate limiter for API calls.

3

4 Implements token bucket algorithm for rate limiting with clean,

5 testable design.

6"""

7

8 import threading

9 import time

10 from typing import Optional

11

12

13 class RateLimiter:

14 """

15 Token bucket rate limiter for controlling API request rates.

16

17 Thread-safe implementation following KISS principle.

18 """

19

20 def __init__(

21 self, requests_per_minute: int, burst_size: Optional[int] = None

22 ):

23 """

24 Initialize rate limiter.

25

26 Args:

27 requests_per_minute: Maximum requests per minute

28 burst_size: Maximum burst size (default: requests_per_minute)

29 """

30 self.rpm = requests_per_minute

31 self.capacity = burst_size or requests_per_minute

32 self.tokens = float(self.capacity)

33 self.last_update = time.time()

34 self.lock = threading.Lock()

35

36 # Calculate refill rate (tokens per second)

37 self.refill_rate = requests_per_minute / 60.0

38

39 def acquire(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:

40 """

41 Acquire tokens for making requests.

42

43 Args:

44 tokens: Number of tokens to acquire

45 timeout: Maximum wait time in seconds (None = wait forever)

46

47 Returns:

48 True if tokens acquired, False if timeout

49

50 Raises:

51 ValueError: If tokens > capacity

52 """

53 if tokens > self.capacity:

54 raise ValueError(

55 f"Requested {tokens} tokens exceeds capacity {self.capacity}"

56 )

57

58 deadline = None if timeout is None else time.time() + timeout

59

60 while True:

61 with self.lock:

62 self._refill()

63

64 if self.tokens >= tokens:

65 self.tokens -= tokens

66 return True

67

68 # Check timeout

69 if deadline is not None and time.time() >= deadline:

70 return False

71

72 # Sleep before retry

73 time.sleep(0.1)

74

75 def _refill(self) -> None:

76 """Refill tokens based on elapsed time."""

77 now = time.time()

78 elapsed = now - self.last_update

79

80 # Add tokens based on refill rate

81 tokens_to_add = elapsed * self.refill_rate

82 self.tokens = min(self.capacity, self.tokens + tokens_to_add)

83

84 self.last_update = now

85

86 @property

87 def available_tokens(self) -> float:

88 """Get current available tokens."""

89 with self.lock:

90 self._refill()

91 return self.tokens

92

93 def reset(self) -> None:

94 """Reset rate limiter to full capacity."""

95 with self.lock:

96 self.tokens = float(self.capacity)

97 self.last_update = time.time()

98