unmas_minebot/streamlit_interface.py

import streamlit as st
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


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"])


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

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