api/rpmlimits.py

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}.")