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!