Refactor chatbot and report generation

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)