Update app.py

app.py

@@ -5,8 +5,9 @@ from google import genai
 st.set_page_config(page_title="Gemini API Tester", page_icon="⚡")
 st.title("⚡ Gemini API Tester")
 
-# ---------------- Session State ----------------
-state_defaults = {
+# ---------------- Session State Defaults ----------------
+defaults = {
+    "api_key_input": "",
     "api_key": None,
     "api_validated": False,
     "client": None,
@@ -15,7 +16,7 @@ state_defaults = {
     "messages": []
 }
 
-for k, v in state_defaults.items():
+for k, v in defaults.items():
     if k not in st.session_state:
         st.session_state[k] = v
 
@@ -23,21 +24,22 @@ for k, v in state_defaults.items():
 with st.sidebar:
     st.header("🔐 API Configuration")
 
-    # ---------- NOT VALIDATED ----------
+    # -------- BEFORE VALIDATION --------
     if not st.session_state.api_validated:
-        api_key_input = st.text_input(
+        st.text_input(
             "Enter Google API Key",
-            type="password"
+            type="password",
+            key="api_key_input"
         )
 
-        validate = st.button("Validate API Key")
-
-        if validate:
-            if not api_key_input:
+        if st.button("Validate API Key"):
+            if not st.session_state.api_key_input:
                 st.warning("Please enter an API key.")
             else:
                 try:
-                    client = genai.Client(api_key=api_key_input)
+                    client = genai.Client(
+                        api_key=st.session_state.api_key_input
+                    )
 
                     models = [
                         m.name for m in client.models.list()
@@ -47,34 +49,35 @@ with st.sidebar:
                     if not models:
                         st.error("API key valid, but no Gemini models available.")
                     else:
-                        # Persist validated state
-                        st.session_state.api_key = api_key_input
+                        # ✅ Persist validated state
+                        st.session_state.api_key = st.session_state.api_key_input
                         st.session_state.client = client
                         st.session_state.models = models
                         st.session_state.selected_model = models[0]
                         st.session_state.api_validated = True
 
-                        st.success("API key validated successfully")
+                        # 🔒 Clear visible input immediately
+                        st.session_state.api_key_input = ""
+
+                        st.success("API key validated")
 
                 except Exception as e:
                     st.error("Invalid API key")
                     st.caption(str(e))
 
-    # ---------- VALIDATED ----------
+    # -------- AFTER VALIDATION --------
     else:
         st.success("🔒 API Key validated (hidden)")
 
-        st.session_state.selected_model = st.selectbox(
+        st.selectbox(
             "Available Models",
             st.session_state.models,
-            index=st.session_state.models.index(
-                st.session_state.selected_model
-            ) if st.session_state.selected_model in st.session_state.models else 0
+            key="selected_model"
         )
 
         if st.button("🔄 Change API Key"):
-            for k in ["api_key", "api_validated", "client", "models", "selected_model"]:
-                st.session_state[k] = state_defaults[k]
+            for k in defaults:
+                st.session_state[k] = defaults[k]
 
 # ---------------- Chat History ----------------
 for msg in st.session_state.messages:
@@ -89,11 +92,13 @@ if prompt := st.chat_input("Ask something..."):
         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)
 
-    # Assistant
+    # Assistant response
     with st.chat_message("assistant"):
         placeholder = st.empty()
         placeholder.markdown("Thinking...")