Update main.py

main.py

@@ -14,10 +14,51 @@ import json
 from typing import Optional
 import datetime
 from usage_tracker import UsageTracker
+from fastapi.middleware.base import BaseHTTPMiddleware
+from collections import defaultdict
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, requests_per_second: int = 2):
+        super().__init__(app)
+        self.requests_per_second = requests_per_second
+        self.last_request_time = defaultdict(float)
+        self.tokens = defaultdict(lambda: requests_per_second)
+        self.last_update = defaultdict(float)
+        
+    async def dispatch(self, request: Request, call_next):
+        client_ip = request.client.host
+        current_time = time.time()
+        
+        # Update tokens
+        time_passed = current_time - self.last_update[client_ip]
+        self.last_update[client_ip] = current_time
+        self.tokens[client_ip] = min(
+            self.requests_per_second,
+            self.tokens[client_ip] + time_passed * self.requests_per_second
+        )
+        
+        # Check if request can be processed
+        if self.tokens[client_ip] < 1:
+            return JSONResponse(
+                status_code=429,
+                content={
+                    "detail": "Too many requests. Please try again later.",
+                    "retry_after": round((1 - self.tokens[client_ip]) / self.requests_per_second)
+                }
+            )
+            
+        # Consume a token
+        self.tokens[client_ip] -= 1
+        
+        # Process the request
+        response = await call_next(request)
+        return response
+
 usage_tracker = UsageTracker()
 load_dotenv() #idk why this shi
 
 app = FastAPI()
+app.add_middleware(RateLimitMiddleware, requests_per_second=2)
 
 # Get API keys and secret endpoint from environment variables
 api_keys_str = os.getenv('API_KEYS') #deprecated -_-