handle the logout functionality correctly

Signed-off-by: lochanpaudel <lochanpaudel10@gmail.com>

Author: nahcol10

src/mvt/app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from menu import menu
+from menu import menu, handle_logout
 from database import create_connection, create_table, create_prompts_table, get_user, insert_user
 from homepage import gethomepage
 import os
@@ -162,5 +162,7 @@ def show_credentials_error():
         create_prompts_table(conn)
         conn.close()
 
-# Show menu
+# Show menu and logout button
+if "user_type" in st.session_state and st.session_state.user_type is not None:
+    handle_logout()
 menu()

src/mvt/menu.py

@@ -3,16 +3,12 @@
 def handle_logout():
     """Handle logout action and state cleanup"""
     if st.sidebar.button("Log out"):
-        st.session_state['user_type'] = None
-        st.session_state['username'] = None
-        st.session_state.pop('first_login', None)
+        for key in list(st.session_state.keys()):
+            del st.session_state[key]
         st.logout()
         st.switch_page("app.py")
 
 def authenticated_menu():
-    # Show logout button at the top of sidebar
-    handle_logout()
-    
     # Show a navigation menu for authenticated users
     st.sidebar.page_link("pages/chatbot.py", label="AIFAQ ChatBot")
     if st.session_state.user_type in ["admin"]: