Update app.py

app.py

@@ -1,4 +1,6 @@
 from flask import Flask, request, jsonify, send_from_directory
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 from PIL import Image
 from io import BytesIO
 from prodiapy import Prodia
@@ -13,8 +15,15 @@ import string
 import random 
 
 
+
 app = Flask(__name__)
 
+limiter = Limiter(
+    app,
+    key_func=get_remote_address,
+    default_limits=["30 per minute"]
+)
+
 # Define the path to static files
 static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
 
@@ -244,6 +253,47 @@ async def generate_image():
     except Exception as e:
         print('Error:', str(e))
         return jsonify({"status": "error", "error": "Internal Server Error"}), 500
+
+
+MODELS = {
+    'gemini-1.0-pro-001': 'gemini-1.0-pro-001',
+    'gemini-1.0-pro': 'gemini-1.0-pro',
+    'gemini-pro': 'gemini-pro',
+    'chat-bison-001': 'chat-bison-001'
+}
+
+@app.route('/gemini', methods=['GET'])
+@limiter.limit("30 per minute")
+def gemini():
+    prompt = request.args.get('prompt')
+    model = request.args.get('model', 'gemini-pro')
+
+    if not prompt:
+        return jsonify({'error': 'Prompt parameter is required'}), 400
+
+    if model not in MODELS:
+        return jsonify({'error': 'Invalid model'}), 400
+
+    url = "https://generativelanguage.googleapis.com/v1beta/models/{}/{}".format(MODELS[model], "generateContent" if model != 'chat-bison-001' else "generateMessage")
+    data = {
+        "contents": [
+            {
+                "parts": [
+                    {
+                        "text": prompt
+                    }
+                ]
+            }
+        ]
+    }
+
+    response = requests.post(url + "?key=AIzaSyCKc8gUTrb7KjhAwY6NwATaMWIDr5etKFk", json=data)
+
+    if response.status_code == 200:
+        result = response.json()
+        return jsonify(result)
+    else:
+        return jsonify({'error': 'Failed to generate content'}), 500
     
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)