From 126ecbca944f60d6799c9c63e73e364a940af7a7 Mon Sep 17 00:00:00 2001
From: AquibPy <mohdaquib31@gmail.com>
Date: Sun, 16 Jun 2024 00:02:36 +0530
Subject: [PATCH] ADDED: Rate limit in Notegem endpoint

---
 README.md        | 14 ++++++++++++++
 api.py           | 24 +++++++++++++++++++++++-
 requirements.txt |  3 ++-
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index fc44454..1e70732 100644
--- a/README.md
+++ b/README.md
@@ -52,6 +52,20 @@ This project supports MongoDB and Redis for data storage and caching respectivel
 - **Clustering:** Supports clustering for horizontal scaling and distributing data across multiple nodes.
 - **Atomic Operations:** Provides atomic operations on data structures, ensuring consistency and reliability.
 
+# Rate Limiting
+
+This FastAPI application includes rate limiting to control the number of requests that can be made to certain endpoints within a specified time frame. Rate limiting helps prevent abuse of the API and ensures fair usage among consumers.
+
+## Rate Limiting Configuration
+
+Rate limiting is implemented using `slowapi`, which provides middleware for rate limiting based on IP address or other identifiers.
+
+### Configuration Details
+
+- **Limits**: Requests are limited to a certain number per minute.
+- **Identifier**: Rate limiting is applied based on the IP address of the client.
+- **Exceeding Limit**: Clients exceeding the limit receive a 429 HTTP status code with an appropriate message.
+
 ## Endpoints
 
 ### 1. Invoice Extractor
diff --git a/api.py b/api.py
index 80303cd..a75f489 100644
--- a/api.py
+++ b/api.py
@@ -39,6 +39,9 @@
 import tempfile
 import shutil
 from youtube_transcript_api import YouTubeTranscriptApi, TranscriptsDisabled, NoTranscriptFound
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
 
 
 os.environ["LANGCHAIN_TRACING_V2"]="true"
@@ -55,6 +58,16 @@
 app = FastAPI(title="Genify By Mohd Aquib",
               summary="This API contains routes of different Gen AI usecases")
 
+limiter = Limiter(key_func=get_remote_address)
+app.state.limit = limiter
+
+@app.exception_handler(RateLimitExceeded)
+async def rate_limit_exceeded_handler(request:Request,exc: RateLimitExceeded):
+    return JSONResponse(
+        status_code= status.HTTP_429_TOO_MANY_REQUESTS,
+        content= {"response": "Limit exceeded, please try again after 2 minute !!!!!!"}
+    )
+
 templates = Jinja2Templates(directory="templates")
 
 app.allow_dangerous_deserialization = True
@@ -740,7 +753,9 @@ async def medigem(image_file: UploadFile = File(...)):
     return ResponseText(response=remove_substrings(response.text))
 
 @app.post("/NoteGem", description="This API endpoint leverages the Google Gemini AI Model to generate comprehensive notes from YouTube video transcripts")
-def process_video(video_url: str = Form(...)):
+@limiter.limit("5/minute")
+
+async def process_video(request: Request, video_url: str = Form(...)):
     video_id = extract_video_id(video_url)
     if not video_id:
         raise HTTPException(status_code=400, detail="Invalid YouTube URL")
@@ -754,9 +769,16 @@ def process_video(video_url: str = Form(...)):
         raise HTTPException(status_code=500, detail=str(e))
 
     try:
+        cache_key = f"notegem:{video_id}"
+        cached_response = redis.get(cache_key)
+        if cached_response:
+            print("Retrieving response from Redis cache")
+            return ResponseText(response=cached_response.decode("utf-8"))
+        
         model = genai.GenerativeModel(settings.GEMINI_PRO_1_5)
         response = model.generate_content(settings.NOTE_GEN_PROMPT + transcript)
         summary = response.text
+        redis.set(cache_key, summary, ex=60)
         db = MongoDB()
         payload = {
                 "endpoint" : "/NoteGem",
diff --git a/requirements.txt b/requirements.txt
index c1835bb..e306831 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -28,4 +28,5 @@ grpcio
 # crewai
 # crewai-tools
 proto-plus
-langchain-cohere
\ No newline at end of file
+langchain-cohere
+slowapi
\ No newline at end of file