File size: 1,615 Bytes
3b993c4 24d33b9 3b993c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import asyncio
import math
import time
from vsp.shared import logger_factory
logger = logger_factory.get_logger(__name__)
class BedrockRateLimiter:
"""
A rate limiter for AWS Bedrock API calls.
This class implements a token bucket algorithm to manage API request rates.
"""
def __init__(self, rate: int, per: float = 60.0):
"""
Initialize the BedrockRateLimiter.
Args:
rate (int): The number of requests allowed per time period.
per (float): The time period in seconds. Defaults to 60.0.
"""
self.rate = rate
self.per = per
self.allowance = rate
self.last_check = time.time()
async def acquire(self) -> None:
"""
Acquire permission to make an API call, respecting rate limits.
This method implements a token bucket algorithm. If the rate limit is exceeded,
it will pause execution for an appropriate amount of time.
Raises:
No specific exceptions are raised, but the method may cause the execution
to sleep for up to 30 seconds if the rate limit is exceeded.
"""
now = time.time()
time_passed = now - self.last_check
self.last_check = now
self.allowance += math.ceil(time_passed * (self.rate / self.per))
if self.allowance > self.rate:
self.allowance = self.rate
if self.allowance < 1:
logger.info("Rate limit exceeded", sleep_time=30)
await asyncio.sleep(30)
self.allowance = 0
else:
self.allowance -= 1
|