feat: Rate Limiter

main.py

@@ -12,7 +12,9 @@ import os
 import logging
 from dotenv import load_dotenv
 
+# Load environment variables
 load_dotenv()
+
 # Import SlowAPI for Rate Limiting
 from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
@@ -20,21 +22,37 @@ from slowapi.middleware import SlowAPIMiddleware
 
 logging.basicConfig(level=logging.INFO)
 
-SECRET_KEY = os.getenv("SECRET_KEY")
-REFRESH_SECRET_KEY = os.getenv("REFRESH_SECRET_KEY")
+# Security Keys with Default Fallbacks
+SECRET_KEY = os.getenv("SECRET_KEY", "default_fallback_256_bit_key")
+REFRESH_SECRET_KEY = os.getenv("REFRESH_SECRET_KEY", SECRET_KEY)
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 30
 REFRESH_TOKEN_EXPIRE_DAYS = 7
 
 app = FastAPI()
 
+# Custom Key Function for Rate Limiting (Handles Proxies)
+def custom_key_func(request: Request):
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0]  # Get real client IP if behind proxy
+    return get_remote_address(request)
+
 # Initialize Rate Limiter
-limiter = Limiter(key_func=get_remote_address)
+limiter = Limiter(key_func=custom_key_func)
 app.state.limiter = limiter
 
 # Attach Rate Limit Exceeded Handler
 app.add_exception_handler(429, _rate_limit_exceeded_handler)
 
+# Custom Rate Limit Response
+@app.exception_handler(429)
+async def rate_limit_exceeded_handler(request: Request, exc):
+    return JSONResponse(
+        status_code=429,
+        content={"error": "Rate limit exceeded. Please try again later."}
+    )
+
 # Add Middleware for Rate Limiting
 app.add_middleware(SlowAPIMiddleware)
 
@@ -99,7 +117,6 @@ async def generate(
     query_input: QueryInput,
     username: str = Depends(verify_access_token),
     stream: bool = Query(False, description="Enable streaming response"),
-    
 ):
     """Handles both streaming and non-streaming responses, with shutdown detection."""
     if shutdown_event.is_set():