Sharing the same widget across several pages of a multi-page app, using session_state

🎈 Using Streamlit
,
1

Hello, I am currently working on my multi-page app, which involves the OpenAI API. In order to manage the API connection, I created a class called APIKeyManager:

APIKeyManager 
import os

import openai
import streamlit as st
from dotenv import load_dotenv


load_dotenv(".env")


class APIKeyManager:
    def __init__(self):
        self.local_api_key = st.session_state.get(
            "local_api_key", os.getenv("LOCAL_OPENAI_API_KEY")
        )
        self.use_local_key = st.session_state.get("use_local_key", False)
        self.user_api_key = st.session_state.get("user_api_key", "")

    def display(self):
        with st.sidebar:
            st.title("OpenAI API Manager")

            self.use_local_key = st.checkbox(
                label="Default API key",
                help="Use my own API key, if you don't have any.",
                on_change=self.check_api_key,
                value=st.session_state.get("use_local_key", False),
                key="use_local_key",
                kwargs={"type": "local"},
            )

            with st.form("api_form"):
                self.user_api_key = st.text_input(
                    label="Enter your API key:",
                    value=self.user_api_key,
                    placeholder="sk-...",
                    type="password",
                    autocomplete="",
                    disabled=self.use_local_key,
                )
                st.form_submit_button(
                    label="Submit",
                    use_container_width=True,
                    disabled=self.use_local_key,
                    on_click=self.check_api_key,
                    kwargs={"type": "human"},
                )

        if st.session_state.get("valid_api_key"):
            st.sidebar.success("Successfully authenticated", icon="🔐")
        else:
            st.sidebar.error("Please add your OpenAI API key to continue")
            st.sidebar.info(
                "Obtain your key from: https://platform.openai.com/account/api-keys"
            )

    def check_api_key(self, type: str):
        api_key = None

        if type == "local" and st.session_state.get("use_local_key"):
            api_key = self.local_api_key
        elif type == "human":
            api_key = self.user_api_key

        try:
            openai.api_key = api_key
            _ = openai.Model.list()
            st.toast("Authentication successful!", icon="✅")
            st.session_state.valid_api_key = True
            self.store_api_key(api_key)
        except openai.error.AuthenticationError:
            st.toast("Authentication error", icon="🚫")
            st.session_state.valid_api_key = False
            self.delete_api_key()

    def store_api_key(self, api_key):
        st.session_state.OPENAI_API_KEY = api_key
        os.environ["OPENAI_API_KEY"] = api_key

    def delete_api_key(self):
        st.session_state.OPENAI_API_KEY = None
        os.environ.pop("OPENAI_API_KEY", None)

Basically, it is a class with a .display() method, in which widgets (a checkbox, and a text_input inside a form) are initialized in the sidebar. It uses streamlit.session_state for variable storing.

Since my app is multi-page, I need this API key manager in the sidebar of every page of my app. I used session_state in order to persist the state of the APIKeyManager.

Home.py 
import streamlit as st
from api_manager import APIKeyManager

st.set_page_config(page_title="Home", layout="wide")


def initialize_api_key_manager():
    if "api_key_manager" not in st.session_state:
        st.session_state.clear()
        st.session_state.api_key_manager = APIKeyManager()


initialize_api_key_manager()


def main():
    st.session_state.api_key_manager.display()

    st.title("Home")
    # rest of the code ...


if __name__ == "__main__":
    main()
Page1.py 
import streamlit as st

st.set_page_config(page_title="Home", layout="wide")


def main():
    if api_key_manager := st.session_state.get("api_key_manager"):
        api_key_manager.display()

    st.title("Page 1")
    # rest of the code ...


if __name__ == "__main__":
    main()

The other pages follow the same structure as Page1.

I could say this works, but I can see that when switching pages, the widget in the sidebar flickers. It very quickly disappears, and then appears again.

Do you have any idea if there could be a better solution? Thanks in advance!

2

This topic was automatically closed 180 days after the last reply. New replies are no longer allowed.