Update api/routes.py

api/routes.py

@@ -1,16 +1,33 @@
 import json
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import StreamingResponse
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from slowapi import Limiter
 from api.auth import verify_app_secret
 from api.config import ALLOWED_MODELS
 from api.models import ChatRequest
 from api.utils import process_non_streaming_response, process_streaming_response
 from api.logger import setup_logger
 
+# Set up logging
 logger = setup_logger(__name__)
 
+# Initialize the router
 router = APIRouter()
 
+# Set up limiter for rate limiting
+limiter = Limiter(key_func=get_remote_address)
+
+# Rate limit exception handler
+@router.exception_handler(RateLimitExceeded)
+async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
+    return Response(
+        status_code=429,
+        content=json.dumps({"error": {"message": "Rate limit exceeded. Please wait and try again.", "type": "rate_limit"}}),
+        media_type="application/json"
+    )
+
 @router.options("/v1/chat/completions")
 @router.options("/api/v1/chat/completions")
 async def chat_completions_options():
@@ -30,8 +47,9 @@ async def list_models():
 
 @router.post("/v1/chat/completions")
 @router.post("/api/v1/chat/completions")
+@limiter.limit("10/minute")  # Limit to 10 requests per minute per IP
 async def chat_completions(
-    request: Request,  # Add Request to capture the HTTP request object
+    request: Request,
     app_secret: str = Depends(verify_app_secret)
 ):
     logger.info("Entering chat_completions route")
@@ -51,7 +69,7 @@ async def chat_completions(
     if chat_request.stream:
         logger.info("Streaming response")
         return StreamingResponse(
-            process_streaming_response(chat_request, request),  # Pass request as request_obj
+            process_streaming_response(chat_request, request),
             media_type="text/event-stream"
         )
     else: