Update app.py

app.py

@@ -257,57 +257,54 @@ async def generate_image():
 
 genai.configure(api_key="AIzaSyBPIdkEyVTDZnmXrBi4ykf0sOfkbOvxAzo")
 
-# Set up the model
-generation_config = {
-    "temperature": 0.9,
-    "top_p": 1,
-    "top_k": 1,
-    "max_output_tokens": 2048,
-}
-
-safety_settings = [
-    {
-        "category": "HARM_CATEGORY_HARASSMENT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_HATE_SPEECH",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-    {
-        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-    },
-]
-
-model = genai.GenerativeModel(
-    model_name="gemini-1.0-pro-001",
-    generation_config=generation_config,
-    safety_settings=safety_settings
-)
+DEFAULT_MODELS = ["gemini-1.0-pro", "gemini-1.0-pro-001"]
 
 @app.route('/gemini', methods=['GET'])
 def gemini():
     prompt = request.args.get('prompt')
+    model_name = request.args.get('model')
 
     if not prompt:
         return jsonify({'error': 'Prompt parameter is required'}), 400
 
+    if model_name and model_name not in DEFAULT_MODELS:
+        return jsonify({'error': f'Model {model_name} not found'}), 400
+
     try:
+        # Use the specified model or default model
+        selected_model = model_name if model_name in DEFAULT_MODELS else DEFAULT_MODELS[0]
+
+        # Set up the selected model
+        generation_config = {
+            "temperature": 1,
+            "top_p": 1,
+            "top_k": 1,
+            "max_output_tokens": 2048,
+        }
+        safety_settings = [
+            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        ]
+
+        model = genai.GenerativeModel(
+            model_name=selected_model,
+            generation_config=generation_config,
+            safety_settings=safety_settings
+        )
+
+        # Start the conversation and generate response
         convo = model.start_chat(history=[])
         convo.send_message(prompt)
         response = convo.last.text
 
-        return jsonify({'response': response})
+        return jsonify({'status': 'true', 'response': response})
 
     except Exception as e:
         error_message = str(e)
         app.logger.error("Failed to generate content: %s", error_message)
-        return jsonify({'error': 'Failed to generate content. Error: {}'.format(error_message)}), 500
+        return jsonify({'status': 'false', 'error': 'Failed to generate content.'.}), 500
 
     
 if __name__ == "__main__":