|
|
import streamlit as st |
|
|
import json |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
from bot import Chatbot, Report, CheckerBot, GeneralBot, botstrings |
|
|
from time import sleep |
|
|
|
|
|
|
|
|
def send(message, check=True, add_to_memory=True): |
|
|
store_state() |
|
|
# Check if the message is a tip. |
|
|
if check: |
|
|
tip = checker_bot.check_for_tips(message) |
|
|
else: |
|
|
tip = False |
|
|
|
|
|
if tip: |
|
|
print("☠️ I'TS A TIP ☠️") |
|
|
return False |
|
|
else: |
|
|
with st.chat_message("assistant"): |
|
|
st.markdown(message) |
|
|
# Add assistant response to chat history |
|
|
st.session_state.messages.append({"role": "assistant", "content": message}) |
|
|
if add_to_memory: |
|
|
chatbot.memory.append({"role":"assistant", "content": message}) |
|
|
return True |
|
|
|
|
|
def store_state(): |
|
|
st.session_state.chatbot = chatbot |
|
|
st.session_state.report = report |
|
|
|
|
|
|
|
|
def generate_report(): |
|
|
from jinja2 import Environment, FileSystemLoader |
|
|
|
|
|
# Load the Jinja2 template |
|
|
env = Environment(loader=FileSystemLoader('.')) |
|
|
template = env.get_template('report_template.html') |
|
|
|
|
|
report_html = template.render(report=report) |
|
|
st.session_state.report_html = report_html |
|
|
st.rerun() |
|
|
|
|
|
|
|
|
|
|
|
st.title("UNMAS Bot") |
|
|
|
|
|
|
|
|
if 'report_html' not in st.session_state: |
|
|
st.session_state.report_html = '' |
|
|
|
|
|
|
|
|
# Initialize chat history |
|
|
if "messages" not in st.session_state: |
|
|
st.session_state.messages = [] |
|
|
|
|
|
if 'upload_image' not in st.session_state: |
|
|
st.session_state.upload_image = False |
|
|
|
|
|
# Load report from session state or create a new one |
|
|
if "report" not in st.session_state: |
|
|
report = Report() |
|
|
st.session_state.report = report |
|
|
else: |
|
|
report = st.session_state.report |
|
|
|
|
|
# Load chatbot from session state or create a new one |
|
|
if "chatbot" not in st.session_state: |
|
|
chatbot = Chatbot() |
|
|
st.session_state.chatbot = chatbot |
|
|
else: |
|
|
chatbot = st.session_state.chatbot |
|
|
if report.object is not None: |
|
|
new_system_prompt = botstrings.checker_bot_system_prompt + f' The user seems to have found {report.object}.' |
|
|
chatbot.memory[0] = {"role":"system", "content": new_system_prompt} |
|
|
print(chatbot.memory) |
|
|
|
|
|
|
|
|
# Load checker and general bot from session state or create a new one |
|
|
if "checker_bot" not in st.session_state: |
|
|
checker_bot = CheckerBot() |
|
|
st.session_state.checker_bot = checker_bot |
|
|
else: |
|
|
checker_bot = st.session_state.checker_bot |
|
|
|
|
|
|
|
|
|
|
|
# Display chat messages from history on app rerun |
|
|
for message in st.session_state.messages: |
|
|
with st.chat_message(message["role"]): |
|
|
st.markdown(message["content"]) |
|
|
|
|
|
# Write the report to screen |
|
|
if st.session_state.report_html != '': |
|
|
st.markdown(st.session_state.report_html, unsafe_allow_html=True) |
|
|
st.session_state.report_html = '' |
|
|
|
|
|
|
|
|
if st.session_state.upload_image: |
|
|
user_input = None |
|
|
img_file = st.file_uploader('Upload an image', type=['png', 'jpg', 'jpeg']) |
|
|
user_input = st.chat_input('') |
|
|
if img_file is not None: |
|
|
image = Image.open(img_file) |
|
|
img_array = np.array(image) |
|
|
report.image = img_array |
|
|
send('Thanks!', check=False, add_to_memory=False) |
|
|
question = chatbot.ask_for_info(report, chatbot=chatbot) |
|
|
send(question, check=False) |
|
|
st.session_state.upload_image = False |
|
|
store_state() |
|
|
st.rerun() |
|
|
|
|
|
elif user_input: |
|
|
answer = checker_bot.check_image_answer(user_input, report) |
|
|
if answer == 'no' or answer == 'help': |
|
|
if answer == 'no': |
|
|
send('No problem, we will continue without the image.', check=False, add_to_memory=False) |
|
|
st.session_state.upload_image = False |
|
|
elif answer == 'help': |
|
|
send('We will provide help for sending a picture in WhatsAPP', check=False) #TODO |
|
|
st.session_state.upload_image = False |
|
|
question = chatbot.ask_for_info(report, chatbot=chatbot) |
|
|
send(question, check=False) |
|
|
|
|
|
|
|
|
elif answer == 'yes': |
|
|
send('Great! Please upload the image.', check=False, add_to_memory=False) |
|
|
st.session_state.upload_image = True |
|
|
store_state() |
|
|
st.rerun() |
|
|
|
|
|
else: |
|
|
user_input = st.chat_input('') |
|
|
|
|
|
|
|
|
if user_input and not st.session_state.upload_image: |
|
|
|
|
|
if user_input.lower() == 'report': |
|
|
# Download the report |
|
|
generate_report() |
|
|
|
|
|
st.chat_message("user").markdown(user_input) |
|
|
# Add user message to chat history |
|
|
st.session_state.messages.append({"role": "user", "content": user_input}) |
|
|
chatbot.memory.append({"role":"user", "content": user_input}) |
|
|
|
|
|
# Check the message type |
|
|
print('Check message type') |
|
|
message_type = checker_bot.check_message_type(message=user_input) |
|
|
print('Message type:', message_type) |
|
|
|
|
|
if not chatbot.first_instructions_sent: |
|
|
# Give instructions for how to use the bot |
|
|
chatbot.first_instructions_sent = True |
|
|
chatbot.informations_requested = True |
|
|
send(botstrings.first_instructions, check=False, add_to_memory=False) |
|
|
sleep(1.2) |
|
|
send('So first, tell me what you found?', check=False) |
|
|
|
|
|
else: |
|
|
if 'greeting' in message_type.lower(): |
|
|
# Answer the greeting |
|
|
bot_answer = chatbot.generate(user_input) |
|
|
send(bot_answer['content']) |
|
|
|
|
|
|
|
|
elif 'question' in message_type.lower(): |
|
|
|
|
|
bot_answer = chatbot.generate(user_input) # TODO How to handle questions? |
|
|
|
|
|
# Check if the answer is an answer to the question. |
|
|
answered = checker_bot.check_answer(bot_answer['content'], user_input, chatbot) |
|
|
if not answered: |
|
|
bot_answer = 'I could not understand your question. Please ask again.' |
|
|
send(bot_answer) |
|
|
chatbot.memory.append({"role":"assistant", "content": bot_answer}) |
|
|
|
|
|
elif 'information' in message_type.lower(): |
|
|
|
|
|
if chatbot.informations_requested: |
|
|
answered = checker_bot.check_answer(user_input, chatbot.memory[-1]['content'], chatbot) |
|
|
if answered: |
|
|
# Ask for information and send that message to the report |
|
|
result = checker_bot.check_for_info(user_input, report, report.looking_for) |
|
|
if result: |
|
|
print(result) |
|
|
else: |
|
|
send('I could not understand your message. Please try again.', check=False) |
|
|
question = chatbot.ask_for_info(report, chatbot=chatbot) |
|
|
|
|
|
if question == 'image': |
|
|
question = "Can you upload a picture of what you have found? Make sure not to get close to what you've found when taking the picture. If you can't take a picture, just say 'no'." |
|
|
st.session_state.upload_image = True |
|
|
print(st.session_state.upload_image) |
|
|
send(question, check=False, add_to_memory=False) |
|
|
st.rerun() |
|
|
|
|
|
elif question == 'done': |
|
|
send('Thank you for your help! Below is the report.', check=False) |
|
|
generate_report() |
|
|
|
|
|
else: |
|
|
send(question, check=False) |
|
|
|
|
|
|
|
|
else: |
|
|
send('I could not understand your message. Please try again.', check=False) |
|
|
|
|
|
|
|
|
else: |
|
|
send('I could not understand your message. Please try again.', check=False) |
|
|
|
|
|
store_state()
|
|
|
|
