Update api/utils.py

api/utils.py

@@ -1,13 +1,15 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 import json
 import uuid
 import asyncio
 import random
 import string
 from typing import Any, Dict, Optional
+import os
+from fastapi import HTTPException, Request
+from dotenv import load_dotenv
 
 import httpx
-from fastapi import HTTPException
 from api import validate  # Import validate to use getHid
 from api.config import (
     MODEL_MAPPING,
@@ -22,13 +24,51 @@ from api.config import (
 from api.models import ChatRequest
 from api.logger import setup_logger
 
+# Initialize environment variables and logger
+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 = {}
+
 # Helper function to create a random alphanumeric chat ID
 def generate_chat_id(length: int = 7) -> str:
     characters = string.ascii_letters + string.digits
     return ''.join(random.choices(characters, k=length))
 
+# Function to get the IP address of the requester
+def get_client_ip(request: Request) -> str:
+    """Retrieve the IP address of the client making the request."""
+    return request.client.host
+
+# Function to limit requests per IP per minute
+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:
+        request_counts[ip] = {"count": 1, "timestamp": current_time}
+        logger.info(f"New IP {ip} added to request counts.")
+    else:
+        ip_data = request_counts[ip]
+        # Reset the count if the timestamp is more than a minute old
+        if current_time - ip_data["timestamp"] > timedelta(minutes=1):
+            request_counts[ip] = {"count": 1, "timestamp": current_time}
+            logger.info(f"Request count reset for IP {ip}.")
+        else:
+            # Increment the count and check if it exceeds the limit
+            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=f"Rate limit exceeded. Maximum {REQUEST_LIMIT_PER_MINUTE} requests per minute allowed."
+                )
+
 # Helper function to create chat completion data
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
@@ -82,19 +122,22 @@ def get_referer_url(chat_id: str, model: str) -> str:
     return BASE_URL
 
 # Process streaming response with headers from config.py
-async def process_streaming_response(request: ChatRequest):
+async def process_streaming_response(request: ChatRequest, request_obj: Request):
     chat_id = generate_chat_id()
     referer_url = get_referer_url(chat_id, request.model)
     logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
 
+    # Get the IP address and check rate limit
+    client_ip = get_client_ip(request_obj)
+    check_rate_limit(client_ip)
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
 
     headers_api_chat = get_headers_api_chat(referer_url)
     validated_token = validate.getHid()  # Get the validated token from validate.py
-    logger.info(f"Retrieved validated token: {validated_token}")
-
+    logger.info(f"Retrieved validated token for IP {client_ip}: {validated_token}")
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(1, 60)
@@ -147,18 +190,22 @@ async def process_streaming_response(request: ChatRequest):
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}")
+            logger.error(f"HTTP error occurred for Chat ID {chat_id} (IP: {client_ip}): {e}")
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
+            logger.error(f"Error occurred during request for Chat ID {chat_id} (IP: {client_ip}): {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
 # Process non-streaming response with headers from config.py
-async def process_non_streaming_response(request: ChatRequest):
+async def process_non_streaming_response(request: ChatRequest, request_obj: Request):
     chat_id = generate_chat_id()
     referer_url = get_referer_url(chat_id, request.model)
     logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
 
+    # Get the IP address and check rate limit
+    client_ip = get_client_ip(request_obj)
+    check_rate_limit(client_ip)
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
@@ -206,10 +253,10 @@ async def process_non_streaming_response(request: ChatRequest):
                 async for chunk in response.aiter_text():
                     full_response += chunk
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}")
+            logger.error(f"HTTP error occurred for Chat ID {chat_id} (IP: {client_ip}): {e}")
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
+            logger.error(f"Error occurred during request for Chat ID {chat_id} (IP: {client_ip}): {e}")
             raise HTTPException(status_code=500, detail=str(e))
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]