forked from cleveragents/cleveragents-core
aditya
2f69358dcf
All 3 simulations tested with the fixed generate prompt that passes the original DoD to the LLM. Results: - SIM12 (Rate Limiter): 7 files — algorithms.py, store.py, middleware.py, config.py, example_app.py, requirements.txt, __init__.py - SIM14 (JSON Validator): 6 files — validator.py, types.py, errors.py, cli.py, requirements.txt, __init__.py - SIM17 (API Tester): 7 files — runner.py, assertions.py, variables.py, reporter.py, cli.py, requirements.txt, test_suite_example.yaml All Python files pass compile() syntax check. All DoD files present. Files written directly via --output-dir, no manual cleanup needed.
151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
import time
|
|
import math
|
|
from abc import ABC, abstractmethod
|
|
from typing import Tuple, Dict, Any, List
|
|
from .store import store
|
|
|
|
|
|
class RateLimitAlgorithm(ABC):
|
|
"""Abstract base class for rate limiting algorithms."""
|
|
|
|
def __init__(self, limit: int, window: int):
|
|
self.limit = limit
|
|
self.window = window
|
|
|
|
@abstractmethod
|
|
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
|
"""
|
|
Check if request is allowed.
|
|
Returns (allowed, headers_dict).
|
|
"""
|
|
pass
|
|
|
|
def _get_key(self, identifier: str) -> str:
|
|
"""Generate storage key for identifier."""
|
|
return f"{self.__class__.__name__}:{identifier}"
|
|
|
|
|
|
class TokenBucket(RateLimitAlgorithm):
|
|
"""Token bucket rate limiting algorithm."""
|
|
|
|
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
|
key = self._get_key(identifier)
|
|
now = time.time()
|
|
|
|
# Get current bucket state
|
|
bucket_data = store.get(key)
|
|
|
|
if bucket_data is None:
|
|
# Initialize bucket
|
|
bucket_data = {
|
|
'tokens': self.limit,
|
|
'last_refill': now
|
|
}
|
|
else:
|
|
# Refill tokens based on elapsed time
|
|
elapsed = now - bucket_data['last_refill']
|
|
tokens_to_add = elapsed * (self.limit / self.window)
|
|
bucket_data['tokens'] = min(self.limit, bucket_data['tokens'] + tokens_to_add)
|
|
bucket_data['last_refill'] = now
|
|
|
|
# Check if request can be served
|
|
if bucket_data['tokens'] >= 1:
|
|
bucket_data['tokens'] -= 1
|
|
allowed = True
|
|
remaining = int(bucket_data['tokens'])
|
|
else:
|
|
allowed = False
|
|
remaining = 0
|
|
|
|
# Store updated bucket state (TTL = 2 * window for safety)
|
|
store.set(key, bucket_data, ttl=self.window * 2)
|
|
|
|
# Calculate reset time
|
|
if bucket_data['tokens'] < self.limit:
|
|
time_to_full = (self.limit - bucket_data['tokens']) / (self.limit / self.window)
|
|
reset_time = int(now + time_to_full)
|
|
else:
|
|
reset_time = int(now)
|
|
|
|
headers = {
|
|
'X-RateLimit-Limit': str(self.limit),
|
|
'X-RateLimit-Remaining': str(remaining),
|
|
'X-RateLimit-Reset': str(reset_time)
|
|
}
|
|
|
|
return allowed, headers
|
|
|
|
|
|
class SlidingWindow(RateLimitAlgorithm):
|
|
"""Sliding window rate limiting algorithm."""
|
|
|
|
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
|
key = self._get_key(identifier)
|
|
now = time.time()
|
|
|
|
# Get current request timestamps
|
|
timestamps = store.get(key) or []
|
|
|
|
# Remove timestamps outside the window
|
|
window_start = now - self.window
|
|
timestamps = [ts for ts in timestamps if ts > window_start]
|
|
|
|
# Check if request is allowed
|
|
if len(timestamps) < self.limit:
|
|
timestamps.append(now)
|
|
allowed = True
|
|
remaining = self.limit - len(timestamps)
|
|
else:
|
|
allowed = False
|
|
remaining = 0
|
|
|
|
# Store updated timestamps
|
|
store.set(key, timestamps, ttl=self.window + 1)
|
|
|
|
# Calculate reset time (when oldest request in window expires)
|
|
if timestamps:
|
|
reset_time = int(timestamps[0] + self.window)
|
|
else:
|
|
reset_time = int(now + self.window)
|
|
|
|
headers = {
|
|
'X-RateLimit-Limit': str(self.limit),
|
|
'X-RateLimit-Remaining': str(remaining),
|
|
'X-RateLimit-Reset': str(reset_time)
|
|
}
|
|
|
|
return allowed, headers
|
|
|
|
|
|
class FixedWindow(RateLimitAlgorithm):
|
|
"""Fixed window rate limiting algorithm."""
|
|
|
|
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
|
now = time.time()
|
|
|
|
# Calculate current window
|
|
window_start = int(now // self.window) * self.window
|
|
key = f"{self._get_key(identifier)}:{window_start}"
|
|
|
|
# Get current count for this window
|
|
current_count = store.get(key) or 0
|
|
|
|
# Check if request is allowed
|
|
if current_count < self.limit:
|
|
new_count = store.increment(key, ttl=self.window + 1)
|
|
allowed = True
|
|
remaining = max(0, self.limit - new_count)
|
|
else:
|
|
allowed = False
|
|
remaining = 0
|
|
|
|
# Calculate reset time (start of next window)
|
|
reset_time = window_start + self.window
|
|
|
|
headers = {
|
|
'X-RateLimit-Limit': str(self.limit),
|
|
'X-RateLimit-Remaining': str(remaining),
|
|
'X-RateLimit-Reset': str(int(reset_time))
|
|
}
|
|
|
|
return allowed, headers |