diff --git a/pyproject.toml b/pyproject.toml index bad59039..8f2f531a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ dependencies = [ [project.optional-dependencies] ui = [ "requests==2.32.3", - "streamlit_local_storage==0.0.23", "streamlit==1.38.0", ] app = [ diff --git a/ui/chat.py b/ui/chat.py index a0e2d589..e72d5ab3 100644 --- a/ui/chat.py +++ b/ui/chat.py @@ -52,7 +52,8 @@ else: rag = st.toggle("Activated RAG", value=False, disabled=True) if new_chat: - st.session_state.clear() + st.session_state.pop("messages", None) + st.session_state.pop("sources", None) st.rerun() if "messages" not in st.session_state: diff --git a/ui/config.py b/ui/config.py index e609a85d..b6396104 100644 --- a/ui/config.py +++ b/ui/config.py @@ -3,6 +3,5 @@ BASE_URL = os.getenv("BASE_URL", "http://localhost:8080/v1") EMBEDDINGS_MODEL_TYPE = "text-embeddings-inference" LANGUAGE_MODEL_TYPE = "text-generation" -LOCAL_STORAGE_KEY = "albertApiKey" INTERNET_COLLECTION_ID = "internet" PRIVATE_COLLECTION_TYPE = "private" diff --git a/ui/utils.py b/ui/utils.py index 56c7ecc6..cea8be7e 100644 --- a/ui/utils.py +++ b/ui/utils.py @@ -3,9 +3,8 @@ import requests import streamlit as st -from streamlit_local_storage import LocalStorage -from config import BASE_URL, EMBEDDINGS_MODEL_TYPE, LANGUAGE_MODEL_TYPE, LOCAL_STORAGE_KEY +from config import BASE_URL, EMBEDDINGS_MODEL_TYPE, LANGUAGE_MODEL_TYPE def set_config(): @@ -28,12 +27,11 @@ def header(): st.subheader("Albert playground") # Authentication - local_storage = LocalStorage() - API_KEY = authenticate(local_storage=local_storage) + API_KEY = authenticate() with col2: logout = st.button("Logout") if logout: - local_storage.deleteItem(LOCAL_STORAGE_KEY) + st.session_state.pop("API_KEY") st.rerun() st.markdown("***") @@ -46,15 +44,15 @@ def check_api_key(base_url: str, api_key: str): return response.status_code == 200 -def authenticate(local_storage: LocalStorage): - API_KEY = local_storage.getItem(LOCAL_STORAGE_KEY) +def authenticate(): + API_KEY = st.session_state.get("API_KEY") if API_KEY is None: with st.form(key="my_form"): API_KEY = st.text_input(label="Please enter your API key", type="password") submit = st.form_submit_button(label="Submit") if submit: if check_api_key(base_url=BASE_URL, api_key=API_KEY): - local_storage.setItem(LOCAL_STORAGE_KEY, API_KEY) + st.session_state["API_KEY"] = API_KEY st.toast("Authentication succeed", icon="✅") time.sleep(0.5) st.rerun()