From 59fadec3a06b25c436aabdec063ee46f8d7dd259 Mon Sep 17 00:00:00 2001 From: Lasse Studion Date: Thu, 22 Feb 2024 09:51:29 +0300 Subject: [PATCH] Refactor chatbot and report generation --- streamlit_interface.py | 83 ++++++++++++++++++++++++++---------------- 1 file changed, 52 insertions(+), 31 deletions(-) diff --git a/streamlit_interface.py b/streamlit_interface.py index a065bbb..466d274 100644 --- a/streamlit_interface.py +++ b/streamlit_interface.py @@ -4,7 +4,8 @@ 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. @@ -37,18 +38,18 @@ def generate_report(): env = Environment(loader=FileSystemLoader('.')) template = env.get_template('report_template.html') - return template.render(report=report) + report_html = template.render(report=report) + st.session_state.report_html = report_html + st.rerun() -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") -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: @@ -57,19 +58,24 @@ if "messages" not in st.session_state: 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 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: @@ -79,16 +85,22 @@ 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) +# 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) @@ -99,13 +111,25 @@ if st.session_state.upload_image: 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}) + + 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('') @@ -113,9 +137,9 @@ else: if user_input and not st.session_state.upload_image: - if user_input.lower() == 'show report': + if user_input.lower() == 'report': # Download the report - report_button() + generate_report() st.chat_message("user").markdown(user_input) # Add user message to chat history @@ -167,18 +191,15 @@ if user_input and not st.session_state.upload_image: question = chatbot.ask_for_info(report, chatbot=chatbot) if question == 'image': - question = "Can you upload a picture of what you have found?" + 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! 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)) + send('Thank you for your help! Below is the report.', check=False) + generate_report() else: send(question, check=False)