Update app.py

app.py

@@ -1,23 +1,26 @@
 # streamlit_app.py
 import streamlit as st
-from pages import page_1, page_2  # Adjusted import statement
+from pages import page_1, page_2
+
+def go_to_main():
+    st.session_state.page = 'main'
+
+def go_to_page_1():
+    st.session_state.page = 'page_1'
+
+def go_to_page_2():
+    st.session_state.page = 'page_2'
 
-# Initialize session state
 if 'page' not in st.session_state:
     st.session_state.page = 'main'
 
-# Main Page
 if st.session_state.page == 'main':
     st.markdown("# Main Page 🏠")
-    if st.button('Go to Page 1'):
-        st.session_state.page = 'page_1'
-    elif st.button('Go to Page 2'):
-        st.session_state.page = 'page_2'
+    st.button('Go to Page 1', on_click=go_to_page_1)
+    st.button('Go to Page 2', on_click=go_to_page_2)
 
-# Page 1
 elif st.session_state.page == 'page_1':
-    page_1.run()
+    page_1.run(go_to_main, go_to_page_2)
 
-# Page 2
 elif st.session_state.page == 'page_2':
-    page_2.run()
+    page_2.run(go_to_main, go_to_page_1)