Update main.py

main.py

@@ -1,40 +1,17 @@
-import os
-import re
-import random
-import string
-import uuid
-import json
-import logging
-import asyncio
-import time
-from collections import defaultdict
-from typing import List, Dict, Any, Optional, Union, AsyncGenerator
-
-from aiohttp import ClientSession, ClientResponseError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    handlers=[logging.StreamHandler()]
-)
-logger = logging.getLogger(__name__)
-
-# Load environment variables
-API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
-RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
-
-if not API_KEYS or API_KEYS == ['']:
-    logger.error("No API keys found. Please set the API_KEYS environment variable.")
-    raise Exception("API_KEYS environment variable not set.")
-
 # Simple in-memory rate limiter based solely on IP addresses
-rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+async def rate_limiter_per_ip(request: Request):
+    client_ip = request.client.host
+    current_time = time.time()
+
+    # Initialize or update the count and timestamp
+    if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
+        rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
+    else:
+        if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
+            logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
+            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address')
+        rate_limit_store[client_ip]["count"] += 1
 
-# Define cleanup interval and window
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
@@ -383,7 +360,7 @@ class ChatRequest(BaseModel):
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
 
-@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
+@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip: Request)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
     # Redact user messages only for logging purposes