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