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.
94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
import functools
|
|
from typing import Union, Optional
|
|
from flask import request, jsonify, Response
|
|
from .algorithms import RateLimitAlgorithm, TokenBucket, SlidingWindow, FixedWindow
|
|
from .config import RateLimitConfig, LimitType
|
|
|
|
|
|
def rate_limit(
|
|
config: RateLimitConfig,
|
|
algorithm: str = "fixed_window"
|
|
) -> callable:
|
|
"""
|
|
Flask decorator for rate limiting routes.
|
|
|
|
Args:
|
|
config: RateLimitConfig instance
|
|
algorithm: Algorithm type ("token_bucket", "sliding_window", "fixed_window")
|
|
"""
|
|
|
|
# Create algorithm instance
|
|
algorithm_map = {
|
|
"token_bucket": TokenBucket,
|
|
"sliding_window": SlidingWindow,
|
|
"fixed_window": FixedWindow
|
|
}
|
|
|
|
if algorithm not in algorithm_map:
|
|
raise ValueError(f"Unknown algorithm: {algorithm}")
|
|
|
|
rate_limiter = algorithm_map[algorithm](config.limit, config.window)
|
|
|
|
def decorator(f):
|
|
@functools.wraps(f)
|
|
def wrapper(*args, **kwargs):
|
|
# Get identifier based on limit type
|
|
identifier = _get_identifier(config)
|
|
|
|
if identifier is None:
|
|
return jsonify({
|
|
"error": "Missing API key" if config.limit_type == LimitType.API_KEY else "Unable to identify client"
|
|
}), 400
|
|
|
|
# Check rate limit
|
|
allowed, headers = rate_limiter.is_allowed(identifier)
|
|
|
|
if not allowed:
|
|
response = jsonify({"error": config.error_message})
|
|
response.status_code = config.error_code
|
|
|
|
# Add rate limit headers
|
|
for header_name, header_value in headers.items():
|
|
response.headers[header_name] = header_value
|
|
|
|
return response
|
|
|
|
# Call original function
|
|
result = f(*args, **kwargs)
|
|
|
|
# Add rate limit headers to successful response
|
|
if isinstance(result, Response):
|
|
response = result
|
|
else:
|
|
# Handle tuple returns (response, status_code, headers)
|
|
if isinstance(result, tuple):
|
|
if len(result) == 2:
|
|
response = jsonify(result[0])
|
|
response.status_code = result[1]
|
|
elif len(result) == 3:
|
|
response = jsonify(result[0])
|
|
response.status_code = result[1]
|
|
response.headers.update(result[2])
|
|
else:
|
|
response = jsonify(result)
|
|
else:
|
|
response = jsonify(result) if not isinstance(result, Response) else result
|
|
|
|
# Add rate limit headers
|
|
for header_name, header_value in headers.items():
|
|
response.headers[header_name] = header_value
|
|
|
|
return response
|
|
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
def _get_identifier(config: RateLimitConfig) -> Optional[str]:
|
|
"""Get client identifier based on configuration."""
|
|
if config.limit_type == LimitType.IP:
|
|
return request.environ.get('REMOTE_ADDR') or request.remote_addr
|
|
elif config.limit_type == LimitType.API_KEY:
|
|
return request.headers.get(config.api_key_header)
|
|
else:
|
|
return None |