Update api/utils.py

api/utils.py

@@ -1,13 +1,12 @@
-from datetime import datetime, timedelta
+from datetime import datetime
 import json
 import uuid
 import asyncio
 import random
 from typing import Any, Dict, Optional
-import os
+import httpx
 from fastapi import HTTPException, Request
 from dotenv import load_dotenv
-import httpx
 from api import validate
 from api.config import (
     MODEL_MAPPING,
@@ -20,49 +19,12 @@ from api.config import (
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
+from api.rpmlimits import check_rate_limit, get_client_ip  # Import rate limit functions
 
 # 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 = {}
-
-# 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:
-        # 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]
-        # 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(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
 ) -> Dict[str, Any]:
@@ -81,13 +43,11 @@ def create_chat_completion_data(
         "usage": None,
     }
 
-# Function to convert message to dictionary format, ensuring base64 data and optional model prefix
 def message_to_dict(message, model_prefix: Optional[str] = None):
     content = message.content if isinstance(message.content, str) else message.content[0]["text"]
     if model_prefix:
         content = f"{model_prefix} {content}"
     if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
-        # Ensure base64 images are always included for all models
         return {
             "role": message.role,
             "content": content,
@@ -99,25 +59,19 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
         }
     return {"role": message.role, "content": content}
 
-# Function to strip model prefix from content if present
 def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
-    """Remove the model prefix from the response content if present."""
     if model_prefix and content.startswith(model_prefix):
         logger.debug(f"Stripping prefix '{model_prefix}' from content.")
         return content[len(model_prefix):].strip()
     return content
 
-# Simplified function to get the base referer URL
 def get_referer_url() -> str:
-    """Return the base URL for the referer without model-specific logic."""
     return BASE_URL
 
-# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest, request_obj: Request):
     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
     client_ip = get_client_ip(request_obj)
     check_rate_limit(client_ip)
 
@@ -126,7 +80,7 @@ async def process_streaming_response(request: ChatRequest, request_obj: Request)
     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
+    validated_token = validate.getHid()
     logger.info(f"Retrieved validated token for IP {client_ip}: {validated_token}")
 
     if request.model == 'o1-preview':
@@ -185,12 +139,10 @@ async def process_streaming_response(request: ChatRequest, request_obj: Request)
             logger.error(f"Error occurred during request (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, request_obj: Request):
     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
     client_ip = get_client_ip(request_obj)
     check_rate_limit(client_ip)