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