Advanced Streamlit: Session State and Callbacks for Data Labelling Tool

It is very simple to create a web app using Streamlit nowadays, but there are limitations when building something complex. One of them is the lack of statefulness because variables in the code will get reinitialize every time we interact with a widget. The good news is Streamlit now has native support of session state and callbacks function in their 0.84 release. We can now store variables across reruns and define our custom event handler upon interaction with a widget.

Session State

Hacks of session state have existed since October of 2019 but the new st.session_state provides us with an elegant solution to access stateful variables. One of the use cases of session state in my data labelling tool is able to preserve the selected project. The app needs to know what are the available projects and what is the project working on. This can simply be done by:

import streamlit as st

# set variables in session state
st.session_state.projects = [
    'Boring Project', 'Interesting Project'
]
st.session_state.current_project = st.radio(
    'Select a project to work with:',
    st.session_state.projects,
)

The next step is to fetch data of the selected project. Besides the selected project, the app also needs the row index of the data to fetch. For simplicity, we assume there is a function that fetches data from the database as follows:

def fetch_data(project_name: str, row_index: int):
    """ A simple example function to fetch data. """
    dbs = {
        'Boring Project': [...],
        'Interesting Project': [...],
    }
    return dbs[project_name][row_index]  # output can be any form

if 'current_project' in st.session_state:
    row_index = st.session_state.get('row_index', 0)
    data = fetch_data(st.session_state.current_project, row_index)

Once a user clicks the next or previous data, the row_index variable in the session state can be set to another number and new data will be loaded automatically.

# examples to update row index
if st.button('Next', key='next_data'):
    st.session_state.row_index += 1

if st.button('Previous', key='previous_data'):
    st.session_state.row_index -= 1

Callbacks

Callback is the function that gets called when the input widget is triggered. Streamlit widgets that support callback are st.button(), st.radio(), st.text_input(), etc. Let’s look at an example of callback in a submit button to add a new project.

def submit_add_project(project_name: str):
    """ Callback function during adding a new project. """
    # display a warning if the user entered an existing name
    if project_name in st.session_state.projects:
        st.warning(f'The name "{project_name}" is already exists.')
    else:
        st.session_state.projects.append(project_name)

new_project = st.text_input('New project name:',
                            key='input_new_project_name')
st.button('Add project', key='button_add_project',
          on_click=submit_add_project, args=(new_project, ))

Note that the callback function requires project_name argument, hence we need to pass the argument using args. Besides adding a new project, we should also allow the user to delete a project by doing so:

def submit_delete_project(project_name: str):
    """ Callback function during deleting an existing project. """
    st.session_state.projects.remove(project_name)

to_delete = st.selectbox('Project to be deleted:',
                         st.session_state.projects)
st.button('Delete project', key='button_delete_project',
          on_click=submit_delete_project, args=(to_delete, ))

Conclusion

This article shows some examples of session state and callbacks which enable the app to preserve variables across rerun and interaction with widgets. I recommend you to have a try with them for in-depth understanding and check out the streamlit blog post. Lastly, You can find the repo for the data labelling tool in Github.

This is my first medium post and I’m very grateful to have you as my readers. Hope that you enjoyed and learned something from this article. Thank you for your time. Stay safe.