|
|
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): |
|
|
""" |
|
|
Sends a message to the chat interface. |
|
|
|
|
|
Parameters: |
|
|
- message (str): The message to be sent. |
|
|
- check (bool): Whether to check the message for tips. Default is True. |
|
|
- add_to_memory (bool): Whether to add the message to the chatbot's memory. Default is True. |
|
|
|
|
|
Returns: |
|
|
- bool: True if the message was sent successfully, False if it was a tip. |
|
|
""" |
|
|
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(): |
|
|
""" |
|
|
Stores the current state of the chatbot and report in the session state. |
|
|
""" |
|
|
st.session_state.chatbot = chatbot |
|
|
st.session_state.report = report |
|
|
|
|
|
|
|
|
def generate_report(): |
|
|
""" |
|
|
Generates a report using a Jinja2 template and stores it in the session state. |
|
|
""" |
|
|
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") |
|
|
|
|
|
|
|
|
# Check if 'report_html', 'messages', 'upload_image',are already in the session state, if not, initialize them |
|
|
if 'report_html' not in st.session_state: |
|
|
st.session_state.report_html = '' |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
# Check if 'report' and the bots are already in the session state, if not, initialize them |
|
|
if "report" not in st.session_state: |
|
|
report = Report() |
|
|
st.session_state.report = report |
|
|
else: |
|
|
report = st.session_state.report |
|
|
|
|
|
if "chatbot" not in st.session_state: |
|
|
chatbot = Chatbot() |
|
|
st.session_state.chatbot = chatbot |
|
|
else: |
|
|
# If 'chatbot' is already in the session state, retrieve it |
|
|
chatbot = st.session_state.chatbot |
|
|
# If the report object is not None, update the chatbot's memory with a new system prompt |
|
|
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) |
|
|
|
|
|
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"]) |
|
|
|
|
|
# If 'report_html' in the session state is not an empty string, write the report to the screen and reset 'report_html' to an empty string |
|
|
if st.session_state.report_html != '': |
|
|
st.markdown(st.session_state.report_html, unsafe_allow_html=True) |
|
|
st.session_state.report_html = '' |
|
|
|
|
|
|
|
|
# Check if 'upload_image' is set to True in the session state |
|
|
if st.session_state.upload_image: |
|
|
import base64 |
|
|
from io import BytesIO |
|
|
# Initialize user_input to None |
|
|
user_input = None |
|
|
# Create a file uploader widget that accepts png, jpg, and jpeg files |
|
|
img_file = st.file_uploader('Upload an image', type=['png', 'jpg', 'jpeg']) |
|
|
# Create a chat input widget and store the user's input in user_input |
|
|
user_input = st.chat_input('') |
|
|
|
|
|
# If an image file was uploaded |
|
|
if img_file is not None: |
|
|
# Open the image file and convert it to a numpy array and store the image array in the report object |
|
|
image = Image.open(img_file) |
|
|
img_array = np.array(image) |
|
|
|
|
|
# Convert the image to Base64 |
|
|
buffered = BytesIO() |
|
|
|
|
|
# Convert image to RGB if it's RGBA |
|
|
if image.mode == 'RGBA': |
|
|
image = image.convert('RGB') |
|
|
|
|
|
image.save(buffered, format="JPEG") |
|
|
img_str = base64.b64encode(buffered.getvalue()).decode() |
|
|
|
|
|
report.image = img_str |
|
|
send('Thanks!', check=False, add_to_memory=False) |
|
|
# Ask the chatbot for more information and store the question |
|
|
question = chatbot.ask_for_info(report, chatbot=chatbot) |
|
|
send(question, check=False) |
|
|
# Set 'upload_image' to False in the session state |
|
|
st.session_state.upload_image = False |
|
|
store_state() |
|
|
st.rerun() |
|
|
|
|
|
# If the user entered something in the chat input widget |
|
|
elif user_input: |
|
|
# Check the user's answer to the image question |
|
|
answer = checker_bot.check_image_answer(user_input, report) |
|
|
if answer == 'no' or answer == 'help': |
|
|
# If the user answered 'no' |
|
|
if answer == 'no': |
|
|
# Send a message to the user and set 'upload_image' to False in the session state |
|
|
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 a help message to the user and set 'upload_image' to False in the session state |
|
|
send('We will provide help for sending a picture in WhatsAPP', check=False) #TODO |
|
|
st.session_state.upload_image = False |
|
|
# Ask the chatbot for more information and store the question |
|
|
question = chatbot.ask_for_info(report, chatbot=chatbot) |
|
|
send(question, check=False) |
|
|
|
|
|
# If the user answered 'yes' |
|
|
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) |
|
|
question = '' |
|
|
|
|
|
# Make the last bot message the question |
|
|
for i in chatbot.memory: |
|
|
if i['role'] == 'assistant': |
|
|
question = i['content'] |
|
|
if answered: |
|
|
# Ask for information and send that message to the report |
|
|
result = checker_bot.check_for_info(user_input, report, looking_for=report.looking_for, question=question) |
|
|
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()
|
|
|
|
