Update app.py

app.py

@@ -4,56 +4,83 @@ from google import genai
 # Page Config
 st.set_page_config(page_title="Gemini API Tester", page_icon="⚡")
 
-# --- Title ---
 st.title("⚡ Gemini API Tester")
-st.markdown("Enter your Google API key to see which Gemini models are available.")
+st.markdown("Enter your Google API key to load available Gemini models.")
+
+# --- Initialize session state ---
+if "api_key" not in st.session_state:
+    st.session_state.api_key = ""
+
+if "api_validated" not in st.session_state:
+    st.session_state.api_validated = False
+
+if "models" not in st.session_state:
+    st.session_state.models = []
+
+if "client" not in st.session_state:
+    st.session_state.client = None
 
 # --- Sidebar ---
 with st.sidebar:
     st.header("🔐 API Configuration")
 
-    api_key = st.text_input(
-        "Enter your Google API Key",
-        type="password"
-    )
-
-    models = []
-    client = None
+    # 🔹 API Key Input (only show if not validated)
+    if not st.session_state.api_validated:
+        api_key_input = st.text_input(
+            "Enter your Google API Key",
+            type="password"
+        )
 
-    if api_key:
-        try:
-            # Initialize client
-            client = genai.Client(api_key=api_key)
-
-            # Fetch available models
-            models = [
-                model.name
-                for model in client.models.list()
-                if "gemini" in model.name.lower()
-            ]
-
-            if models:
-                st.success("✅ API Key is valid")
+        if st.button("Validate API Key"):
+            if not api_key_input:
+                st.warning("⚠️ Please enter an API key")
             else:
-                st.warning("⚠️ API key is valid, but no Gemini models found")
-
-        except Exception as e:
-            st.error("❌ Invalid API Key")
-            st.caption(str(e))
+                try:
+                    client = genai.Client(api_key=api_key_input)
+
+                    # Fetch Gemini models
+                    models = [
+                        m.name
+                        for m in client.models.list()
+                        if "gemini" in m.name.lower()
+                    ]
+
+                    if not models:
+                        st.error("❌ API key valid, but no Gemini models found")
+                    else:
+                        # Save to session state
+                        st.session_state.api_key = api_key_input
+                        st.session_state.client = client
+                        st.session_state.models = models
+                        st.session_state.api_validated = True
+
+                        st.success("✅ API Key validated successfully")
+
+                except Exception as e:
+                    st.error("❌ Invalid API Key")
+                    st.caption(str(e))
+
+    # 🔹 After validation (API key hidden)
+    else:
+        st.success("🔐 API Key validated")
 
-    # Model selection (only show if models exist)
-    selected_model = None
-    if models:
         selected_model = st.selectbox(
             "Available Models",
-            models
+            st.session_state.models
         )
 
+        if st.button("🔄 Change API Key"):
+            st.session_state.api_key = ""
+            st.session_state.api_validated = False
+            st.session_state.models = []
+            st.session_state.client = None
+            st.experimental_rerun()
+
 # --- Chat State ---
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat history
+# Show chat history
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
@@ -61,12 +88,14 @@ for msg in st.session_state.messages:
 # --- Chat Input ---
 if prompt := st.chat_input("Ask something..."):
 
-    if not client or not selected_model:
-        st.error("❌ Please enter a valid API key and select a model first.")
+    if not st.session_state.api_validated:
+        st.error("❌ Please validate your API key first.")
         st.stop()
 
     # User message
-    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.messages.append(
+        {"role": "user", "content": prompt}
+    )
     with st.chat_message("user"):
         st.markdown(prompt)
 
@@ -76,7 +105,7 @@ if prompt := st.chat_input("Ask something..."):
         placeholder.markdown("Thinking...")
 
         try:
-            response = client.models.generate_content(
+            response = st.session_state.client.models.generate_content(
                 model=selected_model,
                 contents=prompt
             )