Update app.py

app.py

@@ -1,112 +1,106 @@
 import streamlit as st
 from google import genai
 
-# Page Config
+# ---------------- Page Config ----------------
 st.set_page_config(page_title="Gemini API Tester", page_icon="⚡")
-
 st.title("⚡ Gemini API Tester")
-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 ---
+# ---------------- Session State ----------------
+state_defaults = {
+    "api_key": None,
+    "api_validated": False,
+    "client": None,
+    "models": [],
+    "selected_model": None,
+    "messages": []
+}
+
+for k, v in state_defaults.items():
+    if k not in st.session_state:
+        st.session_state[k] = v
+
+# ---------------- Sidebar ----------------
 with st.sidebar:
     st.header("🔐 API Configuration")
 
-    # 🔹 API Key Input (only show if not validated)
+    # ---------- NOT VALIDATED ----------
     if not st.session_state.api_validated:
         api_key_input = st.text_input(
-            "Enter your Google API Key",
+            "Enter Google API Key",
             type="password"
         )
 
-        if st.button("Validate API Key"):
+        validate = st.button("Validate API Key")
+
+        if validate:
             if not api_key_input:
-                st.warning("⚠️ Please enter an API key")
+                st.warning("Please enter an API key.")
             else:
                 try:
                     client = genai.Client(api_key=api_key_input)
 
-                    # Fetch Gemini models
                     models = [
-                        m.name
-                        for m in client.models.list()
+                        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")
+                        st.error("API key valid, but no Gemini models available.")
                     else:
-                        # Save to session state
+                        # Persist validated state
                         st.session_state.api_key = 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")
+                        st.success("API key validated successfully")
 
                 except Exception as e:
-                    st.error("❌ Invalid API Key")
+                    st.error("Invalid API key")
                     st.caption(str(e))
 
-    # 🔹 After validation (API key hidden)
+    # ---------- VALIDATED ----------
     else:
-        st.success("🔐 API Key validated")
+        st.success("🔒 API Key validated (hidden)")
 
-        selected_model = st.selectbox(
+        st.session_state.selected_model = st.selectbox(
             "Available Models",
-            st.session_state.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
         )
 
         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 = []
+            for k in ["api_key", "api_validated", "client", "models", "selected_model"]:
+                st.session_state[k] = state_defaults[k]
 
-# Show chat history
+# ---------------- Chat History ----------------
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
 
-# --- Chat Input ---
+# ---------------- Chat Input ----------------
 if prompt := st.chat_input("Ask something..."):
 
     if not st.session_state.api_validated:
-        st.error("❌ Please validate your API key first.")
+        st.error("Please validate an 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)
 
-    # Assistant response
+    # Assistant
     with st.chat_message("assistant"):
         placeholder = st.empty()
         placeholder.markdown("Thinking...")
 
         try:
             response = st.session_state.client.models.generate_content(
-                model=selected_model,
+                model=st.session_state.selected_model,
                 contents=prompt
             )
 
@@ -118,4 +112,4 @@ if prompt := st.chat_input("Ask something..."):
             )
 
         except Exception as e:
-            placeholder.error(f"❌ Error: {e}")
+            placeholder.error(f"Error: {e}")