Update app.py

app.py

@@ -1,7 +1,8 @@
 from flask import Flask, request, jsonify
-import hmac, hashlib, secrets, time
+import hmac, hashlib, secrets, time, os
 from openai import OpenAI
-import os
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 
 app = Flask(__name__)
 
@@ -10,6 +11,8 @@ SECRET_KEY = os.getenv("SECRET_KEY")
 GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
 endpoint = "https://models.inference.ai.azure.com"
 client = OpenAI(base_url=endpoint,api_key=GITHUB_TOKEN)
+# Initialize Flask-Limiter (without global default limits)
+limiter = Limiter(get_remote_address, app=app)
 
 # Track API statistics
 request_count = 0
@@ -75,22 +78,28 @@ def generate_response(query:str) -> str:
         return response.choices[0].message.content
     except:
         return "API Server is under maintenance. Please Try After Some Time Thank You for using TrueSyncAI Chat API. Have a great day."
-    
+ 
+# Custom Rate Limit Error Handler
+@app.errorhandler(429)
+def ratelimit_exceeded(e):
+    return jsonify({
+        "error": "Rate limit exceeded",
+        "message": "You have reached the maximum limit of 5 requests per minute. Please try again later."
+    }), 429
+   
 # Chat Endpoint
 @app.route("/v1/chat/completions", methods=["POST"])
+@limiter.limit("1 per minute")  # Apply limit only to this function
 def chat():
     global request_count
     data = request.json
-    api_key = data.get("api_key")
-    message = data.get("message", "").strip()
+    api_key, message = data.get("api_key"), data.get("message", "").strip()
     
     if not api_key or not validate_api_key(api_key):
         return jsonify({"error": "Invalid API Key"}), 401
-    
     if not message:
         return jsonify({"error": "Message cannot be empty"}), 400
     
-    # Basic AI response (Can integrate LLMs here)
     response = generate_response(message)
     request_count += 1
     return jsonify({"response": response})