Issue fix

app.py

@@ -65,28 +65,36 @@ st.session_state["selected_model"] = st.sidebar.selectbox(
 
 OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
 
+if "model_change_confirmed" not in st.session_state:
+    st.session_state.model_change_confirmed = False
+
 if st.session_state["selected_model"] != previous_model:
     if st.session_state.messages:
-        col1, col2 = st.columns([3,1])
-        with col1:
+        st.session_state.model_change_confirmed = False  # Reset confirmation state
+        with st.sidebar:
             st.warning("Changing models will clear current conversation.")
-        with col2:
-            if st.button("Confirm Change"):
-                st.session_state.messages = []
-                st.session_state.current_image = None
-                st.session_state.llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
-                st.session_state.rag_chain = load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY)
-                st.rerun()
-            if st.button("Cancel"):
-                st.session_state["selected_model"] = previous_model
-                st.rerun()
+            col1, col2 = st.columns(2)
+            with col1:
+                if st.button("Confirm Change", key="confirm_model_change"):
+                    st.session_state.messages = []
+                    st.session_state.current_image = None
+                    st.session_state.model_change_confirmed = True
+                    st.rerun()
+            with col2:
+                if st.button("Cancel", key="cancel_model_change"):
+                    st.session_state["selected_model"] = previous_model
+                    st.rerun()
     else:
-        # No messages yet, just switch
-        st.session_state.llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
-        st.session_state.rag_chain = load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY)
-
-llm = st.session_state.get("llm", initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY))
-rag_chain = st.session_state.get("rag_chain", load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY))
+        st.session_state.model_change_confirmed = True
+
+if "model_change_confirmed" not in st.session_state or st.session_state.model_change_confirmed:
+    llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
+    rag_chain = load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY)
+    st.session_state.llm = llm
+    st.session_state.rag_chain = rag_chain
+else:
+    llm = st.session_state.get("llm", initialize_llm(previous_model, OPENAI_API_KEY))
+    rag_chain = st.session_state.get("rag_chain", load_rag_chain(previous_model, OPENAI_API_KEY))
 
 classifier = load_models()