Update api/rpmlimits.py

api/rpmlimits.py

@@ -1,44 +1,15 @@
-from datetime import datetime, timedelta
-from fastapi import HTTPException, Request
-from api.logger import setup_logger
-import os
-from dotenv import load_dotenv
-
-load_dotenv()
-logger = setup_logger(__name__)
-
-# Set request limit per minute from environment variable
-REQUEST_LIMIT_PER_MINUTE = int(os.getenv("REQUEST_LIMIT_PER_MINUTE", "10"))
-
-# Dictionary to track IP addresses and request counts
-request_counts = {}
-
-def get_client_ip(request: Request) -> str:
-    """Retrieve the IP address of the client making the request."""
-    return request.client.host
-
-def check_rate_limit(ip: str):
-    """Check if the IP has exceeded the request limit per minute."""
-    current_time = datetime.now()
-    
-    if ip not in request_counts:
-        # New IP, initialize with a count and set the timestamp
-        request_counts[ip] = {"count": 1, "timestamp": current_time}
-        logger.info(f"New IP {ip} added to request counts.")
-    else:
-        ip_data = request_counts[ip]
-        # Check if the timestamp is more than a minute old
-        if current_time - ip_data["timestamp"] < timedelta(minutes=1):
-            # If within the same minute, increment the count
-            ip_data["count"] += 1
-            logger.info(f"IP {ip} made request number {ip_data['count']}.")
-            if ip_data["count"] > REQUEST_LIMIT_PER_MINUTE:
-                logger.warning(f"Rate limit exceeded for IP {ip}.")
-                raise HTTPException(
-                    status_code=429,
-                    detail={"error": {"message": "Rate limit exceeded. Please wait and try again.", "type": "rate_limit"}},
-                )
-        else:
-            # If more than a minute has passed, reset the count and timestamp
-            request_counts[ip] = {"count": 1, "timestamp": current_time}
-            logger.info(f"Request count reset for IP {ip}.")
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from fastapi import FastAPI, Request
+from starlette.responses import JSONResponse
+
+# Initialize the limiter
+limiter = Limiter(key_func=get_remote_address)
+
+# Exception handler for rate limit exceeded
+async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
+    return JSONResponse(
+        status_code=429,
+        content={"error": {"message": "Rate limit exceeded. Please wait and try again.", "type": "rate_limit"}}
+    )