Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import pandas as pd
 import numpy as np
 import plotly.express as px
+from plotly.graph_objects import Figure
 from ydata_profiling import ProfileReport
 from streamlit_pandas_profiling import st_profile_report
 import os
@@ -172,7 +173,7 @@ custom_html = """
 </script>
 """
 
-# Helper Functions (mostly unchanged)
+# Helper Functions (unchanged)
 def update_cleaned_data(df):
     st.session_state.cleaned_data = df
     if 'data_versions' not in st.session_state:
@@ -303,8 +304,8 @@ def main():
     # Hidden Inputs for JS Interaction
     if 'app_mode' not in st.session_state:
         st.session_state.app_mode = "Data Upload"
-    app_mode = st.markdown('<input id="app-mode" type="hidden" value="Data Upload">', unsafe_allow_html=True)
-    chat_trigger = st.markdown('<input id="chat-trigger" type="hidden">', unsafe_allow_html=True)
+    st.markdown(f'<input id="app-mode" type="hidden" value="{st.session_state.app_mode}">', unsafe_allow_html=True)
+    st.markdown('<input id="chat-trigger" type="hidden">', unsafe_allow_html=True)
 
     # Sidebar
     with st.sidebar:
@@ -392,13 +393,15 @@ def main():
                     st.session_state.last_plot = {"type": "Histogram", "x": x_axis, "data": df[[x_axis]].to_json()}
                 st.plotly_chart(fig)
 
-    # Chatbot Logic
-    chat_input = st.session_state.get('chat_input', '')
+    # Chatbot Logic (Moved to Bottom, Fixed Session State Issue)
+    if "chat_trigger" not in st.session_state:
+        st.session_state.chat_trigger = ""
+    chat_input = st.session_state.chat_trigger
     if chat_input:
         st.session_state.chat_history.append({"role": "user", "content": chat_input})
         df = st.session_state.cleaned_data if 'cleaned_data' in st.session_state else pd.DataFrame()
         new_df, result, plot_info = parse_command(chat_input, df)
-        if isinstance(result, px.scatter._chart_types.Scatter) or isinstance(result, px.histogram._chart_types.Histogram):
+        if isinstance(result, Figure):
             st.plotly_chart(result)
             st.session_state.last_plot = plot_info
             st.session_state.vector_store = update_vector_store_with_plot(extract_plot_data(plot_info, new_df), st.session_state.vector_store)
@@ -417,7 +420,7 @@ def main():
           document.dispatchEvent(new CustomEvent('bot_response', {{ detail: {json.dumps(response)} }}));
         </script>
         """, unsafe_allow_html=True)
-    st.session_state.chat_input = st.text_input("Chat with AI", key="chat_input", on_change=lambda: None)
+        st.session_state.chat_trigger = ""  # Reset after processing
 
 if __name__ == "__main__":
     main()