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