Issue with login widget in streamlit

I am trying to create a login system to a stramlit app, but I have some issues with the logic. I have a python file where is the login logic and another one where is the sytem's logic. To be more specific I want to move from login screen to dashboard when click 'login' and from dashboard back to login when click 'logout'

I tried to seperate authedicator.logout() from the login() function of the first python file and call it to the other file inside st.session_state['page'] = 'dashboard' block. Also tried some other things like call run_app() inside st.session_state['page'] = 'dashbpard' block but I don't manage to achieve this. What I want is when I click on 'Log Out' to escape st.session_state['page'] = 'dashboard' block and redirect to st.session_state['page'] = 'login' block.

The first python file with login:

import streamlit as st
import mysql.connector
from mysql.connector import Error
import streamlit_authenticator as stauth
from Streamlit_app import run_app

### def create_connetion() and def fetch_credentials() first ###

def login():
    conn = create_connection()
    credentials = format_credentials(fetch_credentials(conn))

    if credentials:
        authenticator = stauth.Authenticate(
            credentials=credentials,
            cookie_name="auth",
            cookie_key="my_key",
            cookie_expiry_days=3,
        )

        name, auth_status, username = authenticator.login()

        if auth_status is None:
            st.warning("Please enter your credentials")
            return auth_status

        if auth_status is False:
            st.error("Username or Password is incorrect")
            return auth_status

        if auth_status:
            user_info = credentials['usernames'][username]
            first_name = user_info["name"]
            last_name = user_info["last_name"]
            date_of_birth = user_info["date_of_birth"]

            st.sidebar.subheader(f'Welcome {first_name} {last_name}!')

            chat_dir = "your_chat_dir"
            run_app(chat_dir)

            authenticator.logout('Log Out', 'sidebar')
            return auth_status
    else:
        st.error("No users found in the database")

    conn.close()


#st.title('Login System')

#login()

The second python file with the pages:

import streamlit as st
from Streamlit_signup import sign_up
from Streamlit_login import login


if 'page' not in st.session_state:
    st.session_state['page'] = 'login'

if 'auth_status' not in st.session_state:
    st.session_state['auth_status'] = None

if st.session_state['page'] == 'login':
    st.title('AI Chatbot')

    col_1, col_2 = st.columns([14, 2])
    with col_2:
        if st.button(":green[Sign up]"):
            st.session_state['page'] = 'sign_up'
            st.rerun()

    auth_status = login()
    st.write("auth status: " + str(auth_status))

    if auth_status:
        st.write("auth status in dashboard: " + str(auth_status))
        st.session_state['page'] = 'dashboard'
        st.session_state['auth_status'] = auth_status
        st.rerun()

elif st.session_state['page'] == 'sign_up':
    st.title("Sign Up Form")
    sign_up()

    if st.button("Back to Login"):
        st.session_state['page'] = 'login'
        st.session_state['auth_status'] = False
        st.rerun()

elif st.session_state['page'] == 'dashboard':
    st.title("Dashboard")
    st.write(st.session_state['auth_status'])
    login()
    st.session_state['page'] = 'login'

Some pictures to help:

1. Login page

2. Dashboard

3. Remains on dashboard after logout

I think this is a better practice and it works properply.

import streamlit as st
import mysql.connector
from mysql.connector import Error
import streamlit_authenticator as stauth
from Streamlit_app import run_app
from Streamlit_signup import sign_up


def login():
    conn = create_connection()
    credentials = format_credentials(fetch_credentials(conn))

    if credentials:
        authenticator = stauth.Authenticate(
            credentials=credentials,
            cookie_name="auth",
            cookie_key="my_key",
            cookie_expiry_days=3,
        )

        name, auth_status, username = authenticator.login()

        if auth_status is None:
            st.warning("Please enter your credentials")
            return auth_status, authenticator

        if auth_status is False:
            st.error("Username or Password is incorrect")
            return auth_status, authenticator

        if auth_status:
            user_info = credentials['usernames'][username]
            first_name = user_info["name"]
            last_name = user_info["last_name"]
            date_of_birth = user_info["date_of_birth"]

            st.sidebar.subheader(f'Welcome {first_name} {last_name}!')

            chat_dir = "YOUT_CHAT_DIR"
            run_app(chat_dir)

            return auth_status, authenticator
    else:
        st.error("No users found in the database")

    conn.close()


def logout(authenticator):
    authenticator.logout('Log Out', 'sidebar')
    st.session_state['auth_status'] = None


if 'page' not in st.session_state:
    st.session_state['page'] = 'login'

if 'auth_status' not in st.session_state:
    st.session_state['auth_status'] = None

if st.session_state['page'] == 'login':

    auth_status, authenticator = login()

    if auth_status:
        st.session_state['auth_status'] = auth_status
        logout(authenticator)
    else:
        if st.sidebar.button("Sign up"):
            st.session_state['page'] = 'sign_up'
            st.rerun()

elif st.session_state['page'] == 'sign_up':
    sign_up()

    if st.button("Back to Login"):
        st.session_state['page'] = 'login'
        st.rerun()