Update api/utils.py

api/utils.py

@@ -3,15 +3,12 @@ import json
 import uuid
 import asyncio
 import random
-import string
 from typing import Any, Dict, Optional
 import os
-from fastapi import HTTPException, Request, Response
-from fastapi.responses import JSONResponse
+from fastapi import HTTPException, Request
 from dotenv import load_dotenv
-
 import httpx
-from api import validate  # Import validate to use getHid
+from api import validate
 from api.config import (
     MODEL_MAPPING,
     get_headers_api_chat,
@@ -44,25 +41,26 @@ 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:
+        # If the IP is new, initialize its counter and 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]
-        # 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
+        # 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}.")
 
 # Helper function to create chat completion data
 def create_chat_completion_data(
@@ -116,7 +114,7 @@ def get_referer_url() -> str:
 
 # Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest, request_obj: Request):
-    referer_url = get_referer_url()  # Get the base URL without model-specific logic
+    referer_url = get_referer_url()
     logger.info(f"Processing streaming response - Model: {request.model} - URL: {referer_url}")
 
     # Get the IP address and check rate limit
@@ -189,7 +187,7 @@ async def process_streaming_response(request: ChatRequest, request_obj: Request)
 
 # Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest, request_obj: Request):
-    referer_url = get_referer_url()  # Get the base URL without model-specific logic
+    referer_url = get_referer_url()
     logger.info(f"Processing non-streaming response - Model: {request.model} - URL: {referer_url}")
 
     # Get the IP address and check rate limit
@@ -202,7 +200,7 @@ async def process_non_streaming_response(request: ChatRequest, request_obj: Requ
 
     headers_api_chat = get_headers_api_chat(referer_url)
     headers_chat = get_headers_chat(referer_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
-    validated_token = validate.getHid()  # Get the validated token from validate.py
+    validated_token = validate.getHid()
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)