From 9db97231b14266233decc5233fb6e97230d496f4 Mon Sep 17 00:00:00 2001 From: Aarushi Date: Sun, 1 Dec 2024 15:05:39 +0000 Subject: [PATCH 1/4] set up rate limt middleware --- .../autogpt_libs/rate_limit/__init__.py | 0 .../autogpt_libs/rate_limit/config.py | 21 ++++++++++ .../autogpt_libs/rate_limit/limiter.py | 39 +++++++++++++++++++ .../autogpt_libs/rate_limit/middleware.py | 31 +++++++++++++++ 4 files changed, 91 insertions(+) create mode 100644 autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/__init__.py create mode 100644 autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py create mode 100644 autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py create mode 100644 autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/__init__.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py new file mode 100644 index 000000000000..4445dee617a7 --- /dev/null +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py @@ -0,0 +1,21 @@ +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class RateLimitSettings(BaseSettings): + redis_url: str = Field( + default="redis://localhost:6379", + description="Redis URL for rate limiting", + validation_alias="RATE_LIMIT_REDIS_URL" + ) + + requests_per_minute: int = Field( + default=60, + description="Maximum number of requests allowed per minute per API key", + validation_alias="RATE_LIMIT_REQUESTS_PER_MINUTE" + ) + + model_config = SettingsConfigDict(case_sensitive=True, extra="ignore") + + +RATE_LIMIT_SETTINGS = RateLimitSettings() diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py new file mode 100644 index 000000000000..e612e3441ef3 --- /dev/null +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py @@ -0,0 +1,39 @@ +import time +import redis +from typing import Tuple +from .config import RATE_LIMIT_SETTINGS + + +class RateLimiter: + def __init__(self, redis_url: str = RATE_LIMIT_SETTINGS.redis_url, + requests_per_minute: int = RATE_LIMIT_SETTINGS.requests_per_minute): + self.redis = redis.from_url(redis_url) + self.window = 60 + self.max_requests = requests_per_minute + + async def check_rate_limit(self, api_key_id: str) -> Tuple[bool, int, int]: + """ + Check if request is within rate limits. + + Args: + api_key_id: The API key identifier to check + + Returns: + Tuple of (is_allowed, remaining_requests, reset_time) + """ + now = time.time() + window_start = now - self.window + key = f"ratelimit:{api_key_id}:1min" + + pipe = self.redis.pipeline() + pipe.zremrangebyscore(key, 0, window_start) + pipe.zadd(key, {str(now): now}) + pipe.zcount(key, window_start, now) + pipe.expire(key, self.window) + + _, _, request_count, _ = pipe.execute() + + remaining = max(0, self.max_requests - request_count) + reset_time = int(now + self.window) + + return request_count <= self.max_requests, remaining, reset_time diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py new file mode 100644 index 000000000000..8e663a5eff16 --- /dev/null +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py @@ -0,0 +1,31 @@ +from fastapi import Request, HTTPException +from .limiter import RateLimiter + + +async def rate_limit_middleware(request: Request, call_next): + """FastAPI middleware for rate limiting API requests.""" + limiter = RateLimiter() + + if not request.url.path.startswith("/api"): + return await call_next(request) + + api_key = request.headers.get("Authorization") + if not api_key: + return await call_next(request) + + api_key = api_key.replace("Bearer ", "") + + is_allowed, remaining, reset_time = await limiter.check_rate_limit(api_key) + + if not is_allowed: + raise HTTPException( + status_code=429, + detail="Rate limit exceeded. Please try again later." + ) + + response = await call_next(request) + response.headers["X-RateLimit-Limit"] = str(limiter.max_requests) + response.headers["X-RateLimit-Remaining"] = str(remaining) + response.headers["X-RateLimit-Reset"] = str(reset_time) + + return response \ No newline at end of file From 771f48e2f9f962bf1664b8673f4a9e182732c521 Mon Sep 17 00:00:00 2001 From: Aarushi Date: Mon, 2 Dec 2024 12:35:41 +0000 Subject: [PATCH 2/4] adddress review comments --- .../autogpt_libs/rate_limit/config.py | 18 +++++++++++++++--- .../autogpt_libs/rate_limit/limiter.py | 8 +++++--- .../autogpt_libs/rate_limit/middleware.py | 6 ++++-- 3 files changed, 24 insertions(+), 8 deletions(-) diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py index 4445dee617a7..1cfa8e3ab399 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py @@ -3,10 +3,22 @@ class RateLimitSettings(BaseSettings): - redis_url: str = Field( + redis_host: str = Field( default="redis://localhost:6379", - description="Redis URL for rate limiting", - validation_alias="RATE_LIMIT_REDIS_URL" + description="Redis host", + validation_alias="REDIS_HOST" + ) + + redis_port: str = Field( + default="6379", + description="Redis port", + validation_alias="REDIS_PORT" + ) + + redis_password: str = Field( + default="password", + description="Redis password", + validation_alias="REDIS_PASSWORD" ) requests_per_minute: int = Field( diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py index e612e3441ef3..5956d6c1bdc2 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py @@ -1,13 +1,15 @@ import time -import redis +from redis import Redis from typing import Tuple from .config import RATE_LIMIT_SETTINGS class RateLimiter: - def __init__(self, redis_url: str = RATE_LIMIT_SETTINGS.redis_url, + def __init__(self, redis_host: str = RATE_LIMIT_SETTINGS.redis_host, + redis_port: str = RATE_LIMIT_SETTINGS.redis_port, + redis_password: str = RATE_LIMIT_SETTINGS.redis_password, requests_per_minute: int = RATE_LIMIT_SETTINGS.requests_per_minute): - self.redis = redis.from_url(redis_url) + self.redis = Redis(host=redis_host, port=redis_port, password=redis_password, decode_responses=True) self.window = 60 self.max_requests = requests_per_minute diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py index 8e663a5eff16..c244a2911af3 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py @@ -1,8 +1,10 @@ from fastapi import Request, HTTPException +from starlette.middleware.base import RequestResponseEndpoint + from .limiter import RateLimiter -async def rate_limit_middleware(request: Request, call_next): +async def rate_limit_middleware(request: Request, call_next: RequestResponseEndpoint): """FastAPI middleware for rate limiting API requests.""" limiter = RateLimiter() @@ -28,4 +30,4 @@ async def rate_limit_middleware(request: Request, call_next): response.headers["X-RateLimit-Remaining"] = str(remaining) response.headers["X-RateLimit-Reset"] = str(reset_time) - return response \ No newline at end of file + return response From 8a3bbc8b67790410b3931ddc8c0065dc824c3d51 Mon Sep 17 00:00:00 2001 From: Aarushi Date: Mon, 2 Dec 2024 13:14:34 +0000 Subject: [PATCH 3/4] run linter --- .../autogpt_libs/autogpt_libs/rate_limit/limiter.py | 4 +++- .../autogpt_libs/autogpt_libs/rate_limit/middleware.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py index 5956d6c1bdc2..9b06068d5493 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py @@ -1,6 +1,8 @@ import time -from redis import Redis from typing import Tuple + +from redis import Redis + from .config import RATE_LIMIT_SETTINGS diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py index c244a2911af3..ca361285469b 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py @@ -1,4 +1,4 @@ -from fastapi import Request, HTTPException +from fastapi import HTTPException, Request from starlette.middleware.base import RequestResponseEndpoint from .limiter import RateLimiter From 7cb2c66f2e41479a6fb4f5838b3ce8f77e6beffe Mon Sep 17 00:00:00 2001 From: Reinier van der Leer Date: Mon, 2 Dec 2024 14:20:53 +0100 Subject: [PATCH 4/4] format --- .../autogpt_libs/feature_flag/client.py | 2 +- .../autogpt_libs/logging/config.py | 1 - .../autogpt_libs/logging/test_utils.py | 6 +++--- .../autogpt_libs/rate_limit/config.py | 10 ++++------ .../autogpt_libs/rate_limit/limiter.py | 18 +++++++++++++----- .../autogpt_libs/rate_limit/middleware.py | 3 +-- 6 files changed, 22 insertions(+), 18 deletions(-) diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py b/autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py index d9f081e5cb4b..dde516c1d8fa 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py @@ -72,7 +72,7 @@ def feature_flag( """ def decorator( - func: Callable[P, Union[T, Awaitable[T]]] + func: Callable[P, Union[T, Awaitable[T]]], ) -> Callable[P, Union[T, Awaitable[T]]]: @wraps(func) async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/logging/config.py b/autogpt_platform/autogpt_libs/autogpt_libs/logging/config.py index 523f6cf8ec87..10db444247c6 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/logging/config.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/logging/config.py @@ -23,7 +23,6 @@ class LoggingConfig(BaseSettings): - level: str = Field( default="INFO", description="Logging level", diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/logging/test_utils.py b/autogpt_platform/autogpt_libs/autogpt_libs/logging/test_utils.py index b3682d42cf0d..24e20986549c 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/logging/test_utils.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/logging/test_utils.py @@ -24,10 +24,10 @@ ), ("", ""), ("hello", "hello"), - ("hello\x1B[31m world", "hello world"), - ("\x1B[36mHello,\x1B[32m World!", "Hello, World!"), + ("hello\x1b[31m world", "hello world"), + ("\x1b[36mHello,\x1b[32m World!", "Hello, World!"), ( - "\x1B[1m\x1B[31mError:\x1B[0m\x1B[31m file not found", + "\x1b[1m\x1b[31mError:\x1b[0m\x1b[31m file not found", "Error: file not found", ), ], diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py index 1cfa8e3ab399..76c9abaa0729 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/config.py @@ -6,25 +6,23 @@ class RateLimitSettings(BaseSettings): redis_host: str = Field( default="redis://localhost:6379", description="Redis host", - validation_alias="REDIS_HOST" + validation_alias="REDIS_HOST", ) redis_port: str = Field( - default="6379", - description="Redis port", - validation_alias="REDIS_PORT" + default="6379", description="Redis port", validation_alias="REDIS_PORT" ) redis_password: str = Field( default="password", description="Redis password", - validation_alias="REDIS_PASSWORD" + validation_alias="REDIS_PASSWORD", ) requests_per_minute: int = Field( default=60, description="Maximum number of requests allowed per minute per API key", - validation_alias="RATE_LIMIT_REQUESTS_PER_MINUTE" + validation_alias="RATE_LIMIT_REQUESTS_PER_MINUTE", ) model_config = SettingsConfigDict(case_sensitive=True, extra="ignore") diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py index 9b06068d5493..efad05836f4f 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/limiter.py @@ -7,11 +7,19 @@ class RateLimiter: - def __init__(self, redis_host: str = RATE_LIMIT_SETTINGS.redis_host, - redis_port: str = RATE_LIMIT_SETTINGS.redis_port, - redis_password: str = RATE_LIMIT_SETTINGS.redis_password, - requests_per_minute: int = RATE_LIMIT_SETTINGS.requests_per_minute): - self.redis = Redis(host=redis_host, port=redis_port, password=redis_password, decode_responses=True) + def __init__( + self, + redis_host: str = RATE_LIMIT_SETTINGS.redis_host, + redis_port: str = RATE_LIMIT_SETTINGS.redis_port, + redis_password: str = RATE_LIMIT_SETTINGS.redis_password, + requests_per_minute: int = RATE_LIMIT_SETTINGS.requests_per_minute, + ): + self.redis = Redis( + host=redis_host, + port=redis_port, + password=redis_password, + decode_responses=True, + ) self.window = 60 self.max_requests = requests_per_minute diff --git a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py index ca361285469b..496697d8b1e2 100644 --- a/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py +++ b/autogpt_platform/autogpt_libs/autogpt_libs/rate_limit/middleware.py @@ -21,8 +21,7 @@ async def rate_limit_middleware(request: Request, call_next: RequestResponseEndp if not is_allowed: raise HTTPException( - status_code=429, - detail="Rate limit exceeded. Please try again later." + status_code=429, detail="Rate limit exceeded. Please try again later." ) response = await call_next(request)