spsetia-chatbot / app.py
cjzhi98's picture
Update app.py
7898e7d verified
import os
import pickle
from pathlib import Path
import streamlit as st
import streamlit_authenticator as stauth
from streamlit_feedback import streamlit_feedback
import traceback
import uuid
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,
save_feedback
)
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,
}
]
if "chat_history" not in st.session_state:
st.session_state.chat_history = []
if "fbk" not in st.session_state:
st.session_state.fbk = str(uuid.uuid4())
def display_answer():
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
def create_answer(prompt):
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)
st.session_state.chat_history.append(
{
"question": prompt,
"answer": answer,
}
)
except:
st.error("An error occurred. Please try again.")
traceback.print_exc()
display_answer()
if prompt := st.chat_input():
create_answer(prompt)
def fbcb(response):
data = st.session_state.chat_history[-1]
question = data["question"]
answer = data["answer"]
print(question, answer, response)
save_feedback(question, answer, response)
st.success("Thank you for your feedback!")
st.session_state.fbk = str(uuid.uuid4())
streamlit_feedback(
feedback_type="thumbs",
optional_text_label="[Optional]",
align="flex-start",
key=st.session_state.fbk,
on_submit=fbcb,
)
# {'type': 'thumbs', 'score': 'πŸ‘Ž', 'text': 'bad'}
# {'type': 'thumbs', 'score': 'πŸ‘', 'text': 'good'}