feat: Rate Limiter Per Users

main.py

@@ -1,26 +1,43 @@
-from fastapi import FastAPI, Depends, HTTPException, status, Query
+from fastapi import FastAPI, Depends, HTTPException, status, Query, Request
 from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from jose import JWTError
 from schemas import UserRegister, TokenResponse, RefreshTokenRequest, QueryInput
-from auth import  register_user, get_db, authenticate_user, create_token, verify_token, verify_access_token, Session
+from auth import register_user, get_db, authenticate_user, create_token, verify_token, verify_access_token, Session
 from utils import generate_stream, generate_response, shutdown_event
 from fastapi.security import OAuth2PasswordRequestForm
 from pathlib import Path
 from datetime import timedelta
 import os
 import logging
+from dotenv import load_dotenv
+
+load_dotenv()
+# Import SlowAPI for Rate Limiting
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.middleware import SlowAPIMiddleware
 
 logging.basicConfig(level=logging.INFO)
 
-SECRET_KEY = os.getenv("SECRET_KEY", 'def6nQHONW99pOPyba9DShny6FB1CJJBigZault')
-REFRESH_SECRET_KEY = os.getenv("REFRESH_SECRET_KEY", SECRET_KEY)
+SECRET_KEY = os.getenv("SECRET_KEY")
+REFRESH_SECRET_KEY = os.getenv("REFRESH_SECRET_KEY")
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 30
 REFRESH_TOKEN_EXPIRE_DAYS = 7
 
 app = FastAPI()
 
+# Initialize Rate Limiter
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+
+# Attach Rate Limit Exceeded Handler
+app.add_exception_handler(429, _rate_limit_exceeded_handler)
+
+# Add Middleware for Rate Limiting
+app.add_middleware(SlowAPIMiddleware)
+
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # Entry Endpoint 
@@ -29,14 +46,18 @@ def index() -> FileResponse:
     file_path = Path(__file__).parent / 'static' / 'index.html'
     return FileResponse(path=str(file_path), media_type='text/html')
 
+# Apply Rate Limiting on Register API (Limit: 5 requests per minute)
 @app.post("/register")
-async def register(user: UserRegister, db: Session = Depends(get_db)):
+@limiter.limit("5/minute")
+async def register(request: Request, user: UserRegister, db: Session = Depends(get_db)):
     """Registers a new user."""
     new_user = register_user(user.username, user.password, db)
     return {"message": "User registered successfully", "user": new_user.username}
 
+# Apply Rate Limiting on Login API (Limit: 10 requests per minute)
 @app.post("/login", response_model=TokenResponse)
-async def login(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
+@limiter.limit("10/minute")
+async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
     try:
         user = authenticate_user(form_data.username, form_data.password, db)
 
@@ -70,38 +91,15 @@ async def login(form_data: OAuth2PasswordRequestForm = Depends(), db: Session =
         logging.error(f"Login error: {str(e)}")
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Internal server error")
 
-@app.post("/refresh", response_model=TokenResponse)
-async def refresh(refresh_request: RefreshTokenRequest):
-    try:
-        # Verify the refresh token
-        username = verify_token(refresh_request.refresh_token, REFRESH_SECRET_KEY)
-        
-        # Create new access token
-        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-        access_token = create_token(
-            data={"sub": username},
-            expires_delta=access_token_expires,
-            secret_key=SECRET_KEY
-        )
-        
-        return {
-            "access_token": access_token,
-            "refresh_token": refresh_request.refresh_token,  # Return the same refresh token
-            "token_type": "bearer"
-        }
-        
-    except JWTError:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Could not validate credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    
+# Apply Rate Limiting on Generate API (Limit: 3 requests per 10 seconds)
 @app.post("/generate")
+@limiter.limit("3/10seconds")
 async def generate(
+    request: Request,
     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():
@@ -123,4 +121,4 @@ async def startup_event():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)