Create APIGuard Middleware class

main.py

@@ -3,6 +3,7 @@ from starlette.applications import Starlette
 from starlette.routing import Route
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
+from starlette.middleware.base import BaseHTTPMiddleware
 from gensim.models import KeyedVectors
 """Prompt templates for LLM"""
 from env import LLM_API_KEY
@@ -58,34 +59,36 @@ class QType(Enum):
     STMT = 3
     FILL = 6
 
+class APIGuardMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app):
+        super().__init__(app)
 
+    async def dispatch(self, request, call_next):
 
-async def __api_guard(request, call_next, **kwargs):
-    # Your custom logic here
-    client_url = request.url.path
-    client_ip_addr = request.client.host
-    
-    # IP-based rate limitation
-    async with r.pipeline(transaction=True) as pipeline:
-        try:
-            res = await (pipeline.get(client_ip_addr).execute())[-1]
-        except:
-            res = None
-
-        if res == None:
-            ok = await pipeline.set(client_ip_addr, 25).execute()
-            r.expire(60)
-        elif res > 0:
-            ok = await pipeline.set(client_ip_addr, res-1).execute()
-        else:
-            raise HTTPException(status_code=429, detail="This IP address is rate-limited")
-    
-    response = await call_next(request)
-    
-    # Optionally modify the response object
-    # ...
+        # Get current client url and client IP address
+        client_url = request.url.path
+        client_ip_addr = request.client.host
+
+        # IP-based rate limitation
+        async with r.pipeline(transaction=True) as pipeline:
+            try:
+                res = await (pipeline.get(client_ip_addr).execute())[-1]
+            except:
+                res = None
     
-    return response
+            if res == None:
+                ok = await pipeline.set(client_ip_addr, 25).execute()
+                r.expire(60)
+            elif res > 0:
+                ok = await pipeline.set(client_ip_addr, res-1).execute()
+            else:
+                raise HTTPException(status_code=429, detail="This IP address is rate-limited")
+        
+        # process the request and get the response    
+        response = await call_next(request)
+        return response
+
+
 sys_random = SystemRandom()
 
 # TODO: Change to environment variable in prod.
@@ -651,7 +654,7 @@ middleware = [
         allow_origins=['http://localhost:8100', 'https://text2quiz-three.vercel.app'],
         allow_methods =['*'],
     ),
-    Middleware(__api_guard),
+    Middleware(APIGuardMiddleware),
 ]
 
 app = Starlette(debug=True,routes=[