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 from pprint import pprint 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') return template.render(report=report) def report_button(): if st.button("Generate report"): report = st.session_state.report report_html = generate_report(report) st.markdown(report_html, unsafe_allow_html=True) st.title("UNMAS Bot") # 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 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 # 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 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"]) print('Upload image', st.session_state.upload_image) if st.session_state.upload_image: user_input = None img_file = st.file_uploader('Upload an image', type=['png', 'jpg', 'jpeg']) 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() # user_input = st.chat_input('') # if user_input: # st.chat_message("user").markdown(user_input) # chatbot.memory.append({"role":"user", "content": user_input}) # st.session_state.messages.append({"role": "user", "content": user_input}) else: user_input = st.chat_input('') if user_input and not st.session_state.upload_image: if user_input.lower() == 'show report': # Download the report report_button() 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?" 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! I will now process the information and get back to you. Below is the report.', check=False) r = report.__dict__.copy() del r['descriptions'] del r['image'] st.write(json.dumps(r, indent=4)) 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()