St.Expander resets session state of dataframe when closing and reopening

Hi, I’m fairly new to coding all together and only just started using streamlit to hopefully make a streamlit based app for some automated data analysis for my team. In part of my code I am reading and visualizing an excel file and by using st.aggrid, making it editable for the scientists to change things. Currently, because there are several sheets that first need QC’ing, I am using expander to let the scientist hide the data from each sheet however I noticed that no matter what I try and do to save the edited dataframe, if I close the expansion and then re-expand it, the initial dataframe goes back to starting values and then after the first re-edit takes the saved dataframe back as well.

Is this a bug of expander or my code? If the former, has there been any work around? If the latter, any suggestions or guidance? Below is what I have so far with me visualizing the updated dataframe as well to see what is going on with it.

   with st.expander(df.columns[0]):
        if 'df2_state' not in st.session_state:
            df = st.session_state.df2_state
            grid_return = AgGrid(df, editable=True, reload_data=True)
            df2 = grid_return["data"]
            st.session_state.df2_state = df2
            grid_return = AgGrid(df2, editable=True, reload_data=True)
            AgGrid(df2)
        else:
            grid_return = AgGrid(df, editable=True, reload_data=True)
            df2 = grid_return["data"]
            st.session_state.df2_state = df2
            AgGrid(df2)