Issue fix

app.py

@@ -55,16 +55,38 @@ available_models = ["OpenAI GPT-4o", "LLaMA 3", "Gemini Pro"]
 if "selected_model" not in st.session_state:
     st.session_state["selected_model"] = available_models[0]
 
+previous_model = st.session_state.get("selected_model", available_models[0])
+
 st.session_state["selected_model"] = st.sidebar.selectbox(
     "Select LLM Model",
     available_models,
     index=available_models.index(st.session_state["selected_model"])
 )
 
-# Dynamically initialize LLM based on selection
 OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
-llm = initialize_llm(st.session_state["selected_model"], OPENAI_API_KEY)
-rag_chain = load_rag_chain(st.session_state["selected_model"], OPENAI_API_KEY)
+
+if st.session_state["selected_model"] != previous_model:
+    if st.session_state.messages:
+        col1, col2 = st.columns([3,1])
+        with col1:
+            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()
+    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))
 
 classifier = load_models()