Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import requests
 
-def check_key_gemini_availability(key):
+def check_key_gemini_availability(key, ai_model):
     """
     Checks the availability of a Gemini API key.
     Args:
@@ -19,7 +19,10 @@ def check_key_gemini_availability(key):
             return False, f"Invalid key: 'models' field not found in response from list models endpoint."
 
         # Second, attempt to generate content to further validate the key
-        url_generateContent = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={key}"
+        ai_model_ = ai_model.strip() or "gemini-1.5-flash"
+        if ai_model_ not in result['models']:
+            return False, f"Specified model is not in available 'models'."
+        url_generateContent = f"https://generativelanguage.googleapis.com/v1beta/models/{ai_model_}:generateContent?key={key}"
         # headers = {'Content-Type': 'application/json'}
         data = {"contents": [{"parts": [{"text": "Say \"this is a test\""}]}]}
         rq = requests.post(url_generateContent, json=data) #headers=headers, data=json.dumps(data)
@@ -74,6 +77,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             keys = gr.TextArea(label="API Keys (by lines)")
+            ai_model = ge.Textbox(label="gemini model")
             with gr.Row():
                 clear_button = gr.Button("Clear")
                 submit_button = gr.Button("Submit", variant="primary")
@@ -82,6 +86,6 @@ with gr.Blocks() as demo:
             errs = gr.TextArea(label="Invalid Keys-status code")
 
     clear_button.click(fn=clear_inputs, outputs=[keys])
-    submit_button.click(fn=check_keys, inputs=[keys], outputs=[infos, errs], api_name="check_keys")
+    submit_button.click(fn=check_keys, inputs=[keys,ai_model], outputs=[infos, errs], api_name="check_keys")
 
 demo.launch()