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()