Spaces:
Sleeping
Sleeping
| import os | |
| import pickle | |
| from pathlib import Path | |
| import streamlit as st | |
| import streamlit_authenticator as stauth | |
| import traceback | |
| from dataset_loader import load_dataset | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| if not os.path.isdir("spsetia_chatbot_data"): | |
| load_dataset() | |
| from spsetia_chatbot_data.utils import ( | |
| run_conversation, | |
| record_db, | |
| hashed_passwords, | |
| greeting | |
| ) | |
| st.set_page_config(page_title="S P Setia Chatbot", page_icon=":robot:", layout="wide") | |
| # --- USER AUTHENTICATION --- | |
| names = ["S P Setia Admin"] | |
| usernames = ["spsetia_admin"] | |
| authenticator = stauth.Authenticate( | |
| names, | |
| usernames, | |
| hashed_passwords, | |
| "spsetia_chatbot", | |
| "eqiddnksadjw", | |
| cookie_expiry_days=30, | |
| ) | |
| name, authentication_status, username = authenticator.login("Login", "main") | |
| if authentication_status == False: | |
| st.error("Username/password is incorrect") | |
| if authentication_status == None: | |
| st.warning("Please enter your username and password") | |
| if authentication_status: | |
| st.sidebar.title(f"S P Setia Chatbot") | |
| st.sidebar.header(f"Welcome {name}") | |
| authenticator.logout("Logout", "sidebar") | |
| if st.sidebar.button("Refresh Chat"): | |
| st.session_state.messages = [ | |
| {"role": "assistant", "content": greeting} | |
| ] | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [ | |
| {"role": "assistant", "content": greeting} | |
| ] | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.write(message["content"]) | |
| if prompt := st.chat_input(): | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| with st.chat_message("user"): | |
| st.write(prompt) | |
| if st.session_state.messages[-1]["role"] != "assistant": | |
| try: | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| response = run_conversation( | |
| st.session_state.messages, prompt, username | |
| ) | |
| qa_id = response[4] | |
| prompt_tokens = response[5] | |
| completion_tokens = response[6] | |
| if response[3]: | |
| if isinstance(response[0], list): | |
| answer = "" | |
| for item in response[0]: | |
| answer += f"""{item[0]}\n""" | |
| st.text(item[0]) | |
| for image in item[1]: | |
| if image: | |
| st.image(image) | |
| reference = f""" \n\nReference \nfolder: {response[2]["folder"]} \nfile: {response[2]["file"]}""" | |
| st.text(reference) | |
| answer += reference | |
| record_db( | |
| content=answer, | |
| qa_id=qa_id, | |
| msg_mode=1, | |
| recipient_id=username, | |
| prompt_tokens=prompt_tokens, | |
| completion_tokens=completion_tokens, | |
| ) | |
| else: | |
| answer = response[0] | |
| reference = f""" \n\nReference \nfolder: {response[2]["folder"]} \nfile: {response[2]["file"]}""" | |
| answer += reference | |
| st.text(answer) | |
| record_db( | |
| content=answer, | |
| qa_id=qa_id, | |
| msg_mode=1, | |
| recipient_id=username, | |
| prompt_tokens=prompt_tokens, | |
| completion_tokens=completion_tokens, | |
| ) | |
| else: | |
| answer = response[0] | |
| st.text(answer) | |
| if images := response[1]: | |
| if isinstance(images, list): | |
| for image in images: | |
| st.image(image) | |
| record_db( | |
| content=answer, | |
| qa_id=qa_id, | |
| msg_mode=1, | |
| recipient_id=username, | |
| prompt_tokens=prompt_tokens, | |
| completion_tokens=completion_tokens, | |
| ) | |
| message = {"role": "assistant", "content": answer} | |
| st.session_state.messages.append(message) | |
| except: | |
| st.error("An error occurred. Please try again.") | |
| traceback.print_exc() | |