Fix Streamlit session state issues with explicit widget keys and robust state management

app.py

@@ -1,4 +1,4 @@
-# Force redeploy trigger - version 2.1
+# Force redeploy trigger - version 2.2
 import streamlit as st
 from utils.config import config
 import requests
@@ -9,28 +9,38 @@ from core.memory import load_user_state, check_redis_health
 # Set page config
 st.set_page_config(page_title="AI Life Coach", page_icon="🧘", layout="centered")
 
-# Initialize session state
-if 'ngrok_url' not in st.session_state:
-    st.session_state.ngrok_url = config.ollama_host
-if 'model_status' not in st.session_state:
-    st.session_state.model_status = "checking"
-if 'available_models' not in st.session_state:
-    st.session_state.available_models = []
-if 'selected_model' not in st.session_state:
-    st.session_state.selected_model = config.local_model_name
+# Comprehensive session state initialization
+session_keys_defaults = {
+    'ngrok_url': config.ollama_host,
+    'model_status': "checking",
+    'available_models': [],
+    'selected_model': config.local_model_name,
+    'selected_model_index': 0,
+    'user_message_input': ""
+}
+
+for key, default_value in session_keys_defaults.items():
+    if key not in st.session_state:
+        st.session_state[key] = default_value
 
 # Sidebar for user selection
 st.sidebar.title("🧘 AI Life Coach")
-user = st.sidebar.selectbox("Select User", ["Rob", "Sarah"])
+user = st.sidebar.selectbox("Select User", ["Rob", "Sarah"], key="user_selector")
 
-# Ngrok URL input in sidebar
+# Ngrok URL input in sidebar - UPDATED VERSION
 st.sidebar.markdown("---")
 st.sidebar.subheader("Ollama Connection")
-ngrok_input = st.sidebar.text_input("Ngrok URL", value=st.session_state.ngrok_url)
-if st.sidebar.button("Update Ngrok URL"):
+ngrok_input = st.sidebar.text_input(
+    "Ngrok URL", 
+    value=st.session_state.ngrok_url,
+    key="ngrok_url_input"  # Explicit key
+)
+
+if st.sidebar.button("Update Ngrok URL", key="update_ngrok_button"):  # Explicit key
     st.session_state.ngrok_url = ngrok_input
     st.session_state.model_status = "checking"
     st.session_state.available_models = []
+    st.session_state.selected_model_index = 0  # Reset model index
     st.sidebar.success("Ngrok URL updated!")
     st.experimental_rerun()
 
@@ -63,22 +73,45 @@ if st.session_state.ngrok_url and st.session_state.model_status != "unreachable"
         # If current selected model not in list, select the first one
         if st.session_state.selected_model not in model_names:
             st.session_state.selected_model = model_names[0]
+            st.session_state.selected_model_index = 0
 
-# Model selector dropdown
+# Model selector dropdown - REPLACED ENTIRE SECTION
 st.sidebar.markdown("---")
 st.sidebar.subheader("Model Selection")
+
+# Initialize selected_model_index in session state if not present
+if 'selected_model_index' not in st.session_state:
+    st.session_state.selected_model_index = 0
+
 if st.session_state.available_models:
+    # Ensure selected_model_index is within bounds
+    if st.session_state.selected_model_index >= len(st.session_state.available_models):
+        st.session_state.selected_model_index = 0
+    
+    # Find index of currently selected model
+    if st.session_state.selected_model in st.session_state.available_models:
+        st.session_state.selected_model_index = st.session_state.available_models.index(st.session_state.selected_model)
+    else:
+        # If current model not available, select first one
+        st.session_state.selected_model_index = 0
+        if st.session_state.available_models:
+            st.session_state.selected_model = st.session_state.available_models[0]
+    
+    # Use index-based selection to avoid widget key issues
     selected_model = st.sidebar.selectbox(
         "Select Model",
         st.session_state.available_models,
-        index=st.session_state.available_models.index(st.session_state.selected_model)
-        if st.session_state.selected_model in st.session_state.available_models
-        else 0
+        index=st.session_state.selected_model_index,
+        key="model_selector"  # Explicit key to prevent conflicts
     )
     st.session_state.selected_model = selected_model
 else:
     st.sidebar.warning("No models available - check Ollama connection")
-    model_input = st.sidebar.text_input("Or enter model name", value=st.session_state.selected_model)
+    model_input = st.sidebar.text_input(
+        "Or enter model name", 
+        value=st.session_state.selected_model,
+        key="manual_model_input"  # Explicit key
+    )
     st.session_state.selected_model = model_input
 
 st.sidebar.markdown("---")
@@ -308,9 +341,18 @@ for msg in conversation:
     content = msg["content"]
     st.markdown(f"**{role}:** {content}")
 
-# Chat input
-user_input = st.text_input("Your message...", key="input")
-if st.button("Send"):
+# Chat input - REPLACED SECTION
+user_input = st.text_input(
+    "Your message...", 
+    key="user_message_input",  # Explicit key
+    placeholder="Type your message here...",
+    value=st.session_state.user_message_input
+)
+
+# Use a separate key for the button to avoid conflicts
+send_button = st.button("Send", key="send_message_button")
+
+if send_button:
     if user_input.strip() == "":
         st.warning("Please enter a message.")
     else:
@@ -334,8 +376,12 @@ if st.button("Send"):
                     st.session_state.selected_model
                 )
                 backend_used = "Ollama"
-            
-            if ai_response:
-                st.markdown(f"**AI Coach ({backend_used}):** {ai_response}")
-            else:
-                st.error(f"Failed to get response from {backend_used}.")
+        
+        if ai_response:
+            st.markdown(f"**AI Coach ({backend_used}):** {ai_response}")
+        else:
+            st.error(f"Failed to get response from {backend_used}.")
+        
+        # Clear input after sending (this helps prevent duplicate sends)
+        st.session_state.user_message_input = ""
+        st.experimental_rerun()