Merge pull request #107 from ant-xuexiao/feat/implement-rate-limit-and-throw-429-if-exceeded

feat: impl rate limit and throw 429 if too many requests

Author: xingwanying

server/auth/get_user_info.py

@@ -0,0 +1,26 @@
+from uilts.env import get_env_variable
+import httpx
+import secrets
+
+
+AUTH0_DOMAIN = get_env_variable("AUTH0_DOMAIN")
+
+async def getUserInfoByToken(token):
+    userinfo_url = f"https://{AUTH0_DOMAIN}/userinfo"
+    
+    headers = {"authorization": f"Bearer {token}"}
+    async with httpx.AsyncClient() as client:
+        user_info_response = await client.get(userinfo_url, headers=headers)
+        if user_info_response.status_code == 200:
+            user_info = user_info_response.json()
+            data = {
+                "id": user_info["sub"],
+                "nickname": user_info.get("nickname"),
+                "name": user_info.get("name"),
+                "picture": user_info.get("picture"),
+                "sub": user_info["sub"],
+                "sid": secrets.token_urlsafe(32)
+            }
+            return data
+        else :
+            return {}

server/routers/auth.py

@@ -1,10 +1,11 @@
 from fastapi import APIRouter,Cookie, Request, HTTPException, status, Response
-from uilts.env import get_env_variable
-from fastapi_auth0 import Auth0
+
 from fastapi.responses import RedirectResponse
 import httpx
+
 from db.supabase.client import get_client
-import secrets
+from auth.get_user_info import getUserInfoByToken
+from uilts.env import get_env_variable
 
 AUTH0_DOMAIN = get_env_variable("AUTH0_DOMAIN")
 
@@ -19,34 +20,12 @@
 WEB_URL =  get_env_variable("WEB_URL")
 
 
-auth = Auth0(domain=AUTH0_DOMAIN, api_audience=API_AUDIENCE, scopes={'read': 'get list'})
-
 router = APIRouter(
     prefix="/api/auth",
     tags=["auth"],
     responses={404: {"description": "Not found"}},
 )
 
-async def getUserInfoByToken(token):
-    userinfo_url = f"https://{AUTH0_DOMAIN}/userinfo"
-    
-
-    headers = {"authorization": f"Bearer {token}"}
-    async with httpx.AsyncClient() as client:
-        user_info_response = await client.get(userinfo_url, headers=headers)
-        if user_info_response.status_code == 200:
-            user_info = user_info_response.json()
-            data = {
-                "id": user_info["sub"],
-                "nickname": user_info.get("nickname"),
-                "name": user_info.get("name"),
-                "picture": user_info.get("picture"),
-                "sub": user_info["sub"],
-                "sid": secrets.token_urlsafe(32)
-            }
-            return data
-        else :
-            return {}
 
 async def getTokenByCode(code):
     token_url = f"https://{AUTH0_DOMAIN}/oauth/token"

server/routers/chat.py

@@ -1,7 +1,8 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
 from fastapi.responses import StreamingResponse
 from data_class import ChatData
 from agent import qa_chat, bot_builder
+from verify.rate_limit import verify_rate_limit
 
 
 router = APIRouter(
@@ -10,12 +11,12 @@
     responses={404: {"description": "Not found"}},
 )
 
-@router.post("/qa", response_class=StreamingResponse)
+@router.post("/qa", response_class=StreamingResponse, dependencies=[Depends(verify_rate_limit)])
 def run_qa_chat(input_data: ChatData):
     result = qa_chat.agent_chat(input_data)
     return StreamingResponse(result, media_type="text/event-stream")
 
-@router.post("/builder", response_class=StreamingResponse)
+@router.post("/builder", response_class=StreamingResponse, dependencies=[Depends(verify_rate_limit)])
 def run_bot_builder(input_data: ChatData):
     result = bot_builder.agent_chat(input_data)
     return StreamingResponse(result, media_type="text/event-stream")

server/routers/health_checker.py

@@ -1,4 +1,6 @@
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends
+
+from verify.rate_limit import verify_rate_limit
 
 router = APIRouter(
     prefix="/api",
@@ -8,4 +10,8 @@
 
 @router.get("/health_checker")
 def health_checker():
+    return { "Hello": "World" }
+
+@router.get("/login_checker", dependencies=[Depends(verify_rate_limit)])
+def login_checker():
     return { "Hello": "World" }

server/routers/rag.py

@@ -1,6 +1,7 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
 from rag import retrieval
 from data_class import S3Config
+from verify.rate_limit import verify_rate_limit
 
 router = APIRouter(
     prefix="/api",
@@ -9,12 +10,12 @@
 )
 
 
-@router.post("/rag/add_knowledge")
+@router.post("/rag/add_knowledge", dependencies=[Depends(verify_rate_limit)])
 def add_knowledge(config: S3Config):
     data=retrieval.add_knowledge(config)
     return data
 
-@router.post("/rag/search_knowledge")
+@router.post("/rag/search_knowledge", dependencies=[Depends(verify_rate_limit)])
 def search_knowledge(query: str):
     data=retrieval.search_knowledge(query)
     return data

server/verify/rate_limit.py

@@ -0,0 +1,45 @@
+
+from fastapi import Cookie, HTTPException
+from datetime import datetime, timedelta
+
+from auth.get_user_info import getUserInfoByToken
+from db.supabase.client import get_client
+
+RATE_LIMIT_REQUESTS = 100
+RATE_LIMIT_DURATION = timedelta(minutes=1)
+
+async def verify_rate_limit(petercat: str = Cookie(None)):
+    if not petercat:
+        raise HTTPException(status_code=403, detail="Must Login")
+    user = await getUserInfoByToken(petercat)
+    user_id = user['id']
+
+    supabase = get_client()
+    table = supabase.table("user_token_usage")
+    rows = table.select('id, user_id, last_request, request_count').eq('user_id', user_id).execute()
+
+    now = datetime.now().isoformat()
+    user_usage = rows.data[0] if len(rows.data) > 0 else { "user_id": user_id, 'request_count': 0, 'last_request': now }
+
+    # Calculate the time elapsed since the last request
+    elapsed_time = datetime.now() - datetime.fromisoformat(user_usage["last_request"])
+    
+    if elapsed_time > RATE_LIMIT_DURATION:
+        # If the elapsed time is greater than the rate limit duration, reset the count
+        user_usage['request_count'] = 1
+    else:
+        if user_usage['request_count'] >= RATE_LIMIT_REQUESTS:
+            # If the request count exceeds the rate limit, return a JSON response with an error message
+            raise HTTPException(
+                status_code=429, 
+                detail="Rate Limit Exceeded, Try It Later",
+                headers={"Retry-After": "60"}
+            )
+
+        user_usage['request_count'] = int(user_usage['request_count']) + 1
+    
+    user_usage['last_request'] = datetime.now().isoformat()
+
+    table.upsert(user_usage).execute()
+
+    return user