Spaces:
Running
Running
import os | |
import time | |
import streamlit as st | |
from backend.constants.streamlit_keys import DATA_INITIALIZE_NOT_STATED, DATA_INITIALIZE_COMPLETED, \ | |
DATA_INITIALIZE_STARTED | |
from backend.constants.variables import DATA_INITIALIZE_STATUS, JUMP_QUERY_ASK, CHAINS_RETRIEVERS_MAPPING, \ | |
TABLE_EMBEDDINGS_MAPPING, RETRIEVER_TOOLS, USER_NAME, GLOBAL_CONFIG, update_global_config | |
from backend.construct.build_all import build_chains_and_retrievers, load_embedding_models, update_retriever_tools | |
from backend.types.global_config import GlobalConfig | |
from logger import logger | |
from ui.chat_page import chat_page | |
from ui.home import render_home | |
from ui.retrievers import render_retrievers | |
# warnings.filterwarnings("ignore", category=UserWarning) | |
def prepare_environment(): | |
os.environ['TOKENIZERS_PARALLELISM'] = 'true' | |
os.environ["LANGCHAIN_TRACING_V2"] = "false" | |
# os.environ["LANGCHAIN_API_KEY"] = "" | |
os.environ["OPENAI_API_BASE"] = st.secrets['OPENAI_API_BASE'] | |
os.environ["OPENAI_API_KEY"] = st.secrets['OPENAI_API_KEY'] | |
os.environ["AUTH0_CLIENT_ID"] = st.secrets['AUTH0_CLIENT_ID'] | |
os.environ["AUTH0_DOMAIN"] = st.secrets['AUTH0_DOMAIN'] | |
update_global_config(GlobalConfig( | |
openai_api_base=st.secrets['OPENAI_API_BASE'], | |
openai_api_key=st.secrets['OPENAI_API_KEY'], | |
auth0_client_id=st.secrets['AUTH0_CLIENT_ID'], | |
auth0_domain=st.secrets['AUTH0_DOMAIN'], | |
myscale_user=st.secrets['MYSCALE_USER'], | |
myscale_password=st.secrets['MYSCALE_PASSWORD'], | |
myscale_host=st.secrets['MYSCALE_HOST'], | |
myscale_port=st.secrets['MYSCALE_PORT'], | |
query_model="gpt-3.5-turbo-0125", | |
chat_model="gpt-3.5-turbo-0125", | |
untrusted_api=st.secrets['UNSTRUCTURED_API'], | |
myscale_enable_https=st.secrets.get('MYSCALE_ENABLE_HTTPS', True), | |
)) | |
# when refresh browser, all session keys will be cleaned. | |
def initialize_session_state(): | |
if DATA_INITIALIZE_STATUS not in st.session_state: | |
st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_NOT_STATED | |
logger.info(f"Initialize session state key: {DATA_INITIALIZE_STATUS}") | |
if JUMP_QUERY_ASK not in st.session_state: | |
st.session_state[JUMP_QUERY_ASK] = False | |
logger.info(f"Initialize session state key: {JUMP_QUERY_ASK}") | |
def initialize_chat_data(): | |
if st.session_state[DATA_INITIALIZE_STATUS] != DATA_INITIALIZE_COMPLETED: | |
start_time = time.time() | |
st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_STARTED | |
st.session_state[TABLE_EMBEDDINGS_MAPPING] = load_embedding_models() | |
st.session_state[CHAINS_RETRIEVERS_MAPPING] = build_chains_and_retrievers() | |
st.session_state[RETRIEVER_TOOLS] = update_retriever_tools() | |
# mark data initialization finished. | |
st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_COMPLETED | |
end_time = time.time() | |
logger.info(f"ChatData initialized finished in {round(end_time - start_time, 3)} seconds, " | |
f"session state keys: {list(st.session_state.keys())}") | |
st.set_page_config( | |
page_title="ChatData", | |
page_icon="https://myscale.com/favicon.ico", | |
initial_sidebar_state="expanded", | |
layout="wide", | |
) | |
prepare_environment() | |
initialize_session_state() | |
initialize_chat_data() | |
if USER_NAME in st.session_state: | |
chat_page() | |
else: | |
if st.session_state[JUMP_QUERY_ASK]: | |
render_retrievers() | |
else: | |
render_home() | |