unmas_minebot/bot.py

import re
import json
import os
import openai
from dotenv import load_dotenv
import requests
from bot_strings import botstrings

# Call the load_dotenv function
load_dotenv()

from pprint import pprint


class Report:
    def __init__(self):
        self.object = None
        self.description = None
        self.incident = None
        self.name = None
        self.role = None
        self.location = None
        self.image = None
        self.size = None
        self.shape = None
        self.color = None
        self.organization = None
        self.type = None
        self.marked = None
        self.blockage = None
        self.area_description = None
        self.contact_person_phone_number = None
        self.contact_person_name = None
        self.at_the_location: bool = None
        self.coordinates = None
        self.woreda = None

        self.info_groups = [
            ["description"],
            ["color", "size", "shape"],
            ["at_the_location"],
            #['coordinates'],
            ["image"],
            ["woreda"], 
            ["location"],
            ["area_description"],
            ["incident"],
            ["blockage", "marked"],
        ]

        self.info_explicit = [
            "name",
            "contact_person_name",
            "contact_person_phone_number",
        ]

        self.info_professional = ["organization", "role", "type"]

        self.descriptions = {
            "description": "A general description of the object.",
            "at_the_location": "If the person reporting is still where the object was found. (True or False)",
            "incident": "Has there been an incident related to the object?",
            "name": "The name of the person reporting.",
            "role": "The role or occupation of the person reporting.",
            "woreda": "The woreda where the object was found.",
            "location": "The closest city or village to where the object was found.",  # TODO Add more details
            "coordinates": "The coordinates of the person reporting.",
            "image": "A picture of the object.",
            "size": "The size of the object.",
            "shape": "The shape of the object.",
            "color": "The color of the object.",
            "organization": "The organization of the person reporting.",
            "marked": "Is the object marked?",
            "blockage": "Is the object blocking anything?",
            "area_description": "A description of the area where the object is located.",
            "contact_person_phone_number": "The phone number of the contact person.",
            "contact_person_name": "The name of the contact person.",
            "type": "The type of the object.",
        }

        self.looking_for: list = self.info_groups[0]


class BaseBot:
    def __init__(self) -> None:
        # Fetch the OpenAI key from the environment variables
        self.OPEN_AI_KEY = os.getenv("OPEN_AI_KEY")
        # # Set the OpenAI key
        # openai.api_key = self.OPEN_AI_KEY
        self.url = "https://api.openai.com/v1/chat/completions"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.OPEN_AI_KEY}",
        }

        self.memory: list[dict] = None

    def generate(self, message):

        message = re.sub(" +", " ", message)

        if self.memory:
            messages = self.memory
            if messages[-1]["content"] != message:
                messages.append({"role": "user", "content": message})
        else:
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": message},
            ]

        data = {"model": "gpt-3.5-turbo", "messages": messages}

        # Print the last message in yellow
        last_message = messages[-1]["content"]
        print("\n\033[94m" + last_message + "\033[0m\n")

        response = requests.post(self.url, headers=self.headers, json=data).json()
        pprint(response)
        answer = response["choices"][0]["message"]

        print("\033[95m" + answer["content"] + "\033[0m")

        return answer


class Chatbot(BaseBot):
    def __init__(self) -> None:
        system_prompt = botstrings.chatbot_system_prompt
        super().__init__()  # Inherit from BaseBot
        # TODO Check if there is an older conversation with the same number.
        # Set the system prompt and the first user message.
        self.first_instructions_sent = False
        self.informations_requested = False
        self.memory = [{"role": "system", "content": system_prompt}]

    def ask_for_info(self, report: Report, chatbot: "Chatbot"):
        # Ask for information
        looking_for = []
        group_found = False
        for group in report.info_groups:
            if not group_found:
                for info in group:
                    print("#", info, getattr(report, info))
                    if getattr(report, info) is None:
                        group_found = True
                        looking_for.append(info)

            # Check so that the set of info is not already asked for without an answer
            if looking_for == report.looking_for:
                for info in looking_for:
                    setattr(report, info, "Unanswered")
                looking_for = []
                group_found = False

        if looking_for != []:
            report.looking_for = looking_for
            # Ask for the information
            if looking_for == ["image"]:
                question = "image"
            else:
                description_dict = {}
                for i in looking_for:
                    description_dict[i] = report.descriptions[i]

                # Give the bot examles of how to ask for the information
                general_bot.memory = [
                    {
                        "role": "user",
                        "content": """I'm looking for the information in this disctionary:
                        {'description': 'If the person reporting is still where the object was found?'}
                        Formulate a question asking for the information in the dictionary.""",
                    },
                    {
                        "role": "assistant",
                        "content": f"Are you still at the place where object was found?",
                    },
                ]

                # If we know what the object is, we can ask about it.
                if report.object is not None:
                    general_bot.memory.insert(
                        0,
                        {
                            "role": "system",
                            "content": botstrings.general_bot_system_prompt
                            + f" The object found might be {report.object}.",
                        },
                    )
                    
                    general_bot.memory[2]["content"] = f"Are you still at the place where the {report.object} was found?"

                object_string = ''
                if report.object is not None:
                    object_string = f" The object found is {report.object}."
                prompt = f"""I'm looking for the information about an object. {object_string}. The information needed is described in this disctionary:
                        {description_dict} 
                        Formulate a question asking for the information in the dictionary."""
                
                question = general_bot.generate(prompt)["content"]
                general_bot.memory = []

            return question

        else:
            return "done"


class CheckerBot(BaseBot):
    def __init__(self) -> None:
        self.system_prompt = botstrings.checker_bot_system_prompt
        super().__init__()  # Inherit from BaseBot
        self.history = []

    def check_message_type(self, message):
        prompt = f"""What kind of message is this?
            '''{message}'''
            Don't care about the content, only the type of message. \
            Answer with any of the following: ["greeting", "question", "information", "statement"]\
            """

        result = self.generate(prompt)["content"]
        if "greeting" in result.lower():
            message_type = "greeting"
        elif "question" in result.lower():
            message_type = "question"
        elif (
            "information" in result.lower() or "statement" in result.lower()
        ):  # TODO should "statement" be here?
            message_type = "information"

        return message_type

    def check_answer(self, answer, question, chatbot: Chatbot):
        return True  #! We need to finf a good way to check if the answer is an answer to the question.
        question = question.replace("\n", " ")
        answer = answer.replace("\n", " ")
        if question == botstrings.first_instructions:
            question = "What have you found?"

        prompt = f'''
        Have a look at this conversation:\n
        """
        {chat2string(chatbot.memory[-4:]).strip()}
        """\n
        Is the last message a resonable answer to the last question ({question})?
        Answer ONLY with any of the following: "True", "False", "Unclear"
        '''.strip()
        result = self.generate(prompt)["content"]

        if "unclear" in result.lower():
            prompt = f"""
            A user is having a conversation with an assistant. This is the conversation:\n'''{chatbot.memory}'''\
            Is the last message ('''{answer}''') an answer to the question ('{question}')?
            Answer ONLY with any of the following: "True", "False", "Unclear"\
            """
            self.memory.append({"role": "user", "content": prompt})

            result = self.generate(prompt)

        if "true" in result.lower():
            answered = True

        elif "false" in result.lower() or "unclear" in result.lower():
            answered = False

        return answered

    def check_for_info(self, user_message, report: Report, question: str, looking_for: list, n_try=0) -> dict:

        if report.object is None:

            prompt = f"""
                    A user has sent a message:\n'''{user_message}'''
                    If the user is describing what he/she has found, describe the object with a single word or a very short phrase. Else answer "null".\
                    """
            messages = [
                {
                    "role": "user",
                    "content": prompt.replace(user_message, "I think it's a bomb"),
                },
                {"role": "assistant", "content": "a bomb"},
                {
                    "role": "user",
                    "content": prompt.replace(user_message, "I don'w know"),
                },
                {"role": "assistant", "content": "null"},
                {
                    "role": "user",
                    "content": prompt.replace(
                        user_message, "something that looks weird"
                    ),
                },
                {"role": "assistant", "content": "a weird thing"},
                {"role": "user", "content": prompt},
            ]
            self.memory = messages
            result = self.generate(user_message)["content"]
            self.memory = []

            if result not in ["null", "None", "", "Unknown"]:
                print("Object:", result)
                report.object = result

        info_dict = {}
        for info in looking_for:
            if getattr(report, info) is None:
                info_dict[info] = report.descriptions[info]

        # If the bot asks for a description, it should also check for color, size and shape.
        if looking_for == ["description"]:
            for i in ["color", "size", "shape"]:
                info_dict[i] = report.descriptions[i]

        prompt = f""""
    A user was asked '''{question}''' 
    This is the answer from the user: '''{user_message}'''\n
    This is a dictionary describing what info I want:\n\n{info_dict} \n\n\
    Take a look at the message along with the question and create a dictionary with the information that is requested.\
    There might not be information available in the mesasge for all fields. \
    If you can't find information for a certain field, fill that with a python null value ("None" or "null").
    Do not include any explanations, only provide a  RFC8259 compliant JSON response  following this format without deviation.\
        """

        n_try += 1
        result = self.generate(prompt)

        try:
            json_content = json.loads(
                result["content"]
            )  # Parse the string back into a dictionary
            for key, value in json_content.items():
                if value not in ["null", "None", "", "Unknown", None]:
                    setattr(report, key, value)
                else:  # Don't ask for these again.
                    if (
                        key in ["shape", "size", "color", "marked", "blockage"]
                        and key in looking_for
                    ):
                        setattr(report, key, "Unknown")

        except:
            try:
                result_content = result_content[
                    result_content.find("{") : result_content.rfind("}") + 1
                ]

                json_content = json.loads(
                    result["content"]
                )  # Parse the string back into a dictionary
                for key, value in json_content.items():
                    if value not in ["null", "None", "", "Unknown"]:
                        setattr(report, key, value)

            except:
                if n_try > 3:
                    return False
                self.check_for_info(user_message, report, looking_for = looking_for, question=question, n_try=n_try)

        return json_content

    def check_for_tips(self, bot_message):
        # Check if the message is a tip
        prompt = f"""
        An assistant wants to send this message to a user: '''{bot_message}'''\
        Is the message a tip or recommendation on how to handle a suspicious object?
        Answer ONLY with any of the following: "True", "False"\
        """
        result = self.generate(prompt)["content"]
        print(result)
        if "true" in result.lower():
            is_tip = True
        elif "false" in result.lower():
            is_tip = False

        return is_tip

    def check_image_answer(self, message: str, report: Report):
        prompt = f"""
            user has been asked to upload an image and answered with a message:
            {message}
            What is the meaning of the message? Answer ONLY with one of ["no, I can't", "yes, I can", "I don't know how to send it"]
            """

        if "no" in self.generate(prompt)["content"].lower():
            answer = "no"
            report.image = "Cannot send image"
        elif "yes" in self.generate(prompt)["content"].lower():
            answer = "yes"
        else:
            answer = "help"
        return answer


class GeneralBot(BaseBot):
    def __init__(self) -> None:
        self.system_prompt = botstrings.general_bot_system_prompt
        super().__init__()
        self.memory = []


def send(message, check=True):

    # Check if the message is a tip.
    if check:
        tip = checker_bot.check_for_tips(message)
    else:
        tip = False

    if tip:
        return False
    else:
        print(message)  # TODO Send the message to the user
        return True


def chat2string(chat: list):
    chat_string = ""
    for message in chat:
        chat_string += f"{message['role']}: {message['content']}\n"
    return chat_string


if __name__ == "__main__":

    #! For testing
    user_input = "I have found a strange looking thing on my lawn."

    # Initiate the chatbot.
    chatbot = Chatbot(
        first_message=user_input,
        system_prompt="You are an assistant chatting with a user.",
    )  # TODO ? Add phone number, etc.
    checker_bot = CheckerBot()
    general_bot = GeneralBot()
    report = Report()

    while True:
        # Check the message type
        message_type = checker_bot.check_message_type(message=user_input)
        print("Message type:", message_type)

        if "greeting" in message_type.lower():
            # Answer the greeting
            bot_answer = chatbot.generate()
            if not chatbot.first_instructions_sent:
                # Give instructions for how to use the bot
                chatbot.instructions_sent = True
                chatbot.informations_requested = True
                chatbot.memory.append(
                    {"role": "system", "content": botstrings.first_instructions}
                )
            else:
                chatbot.memory.append(bot_answer)

        elif "question" in message_type.lower():
            if not chatbot.first_instructions_sent:
                # Give instructions for how to use the bot
                chatbot.first_instructions_sent = True
                chatbot.memory.append(
                    {"role": "system", "content": botstrings.first_instructions}
                )

            bot_answer = chatbot.generate()  # 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": "system", "content": bot_answer})

        elif "information" in message_type.lower():
            if not chatbot.first_instructions_sent:
                # Give instructions for how to use the bot
                chatbot.first_instructions_sent = True
                chatbot.informations_requested = True
                chatbot.memory.append(
                    {"role": "system", "content": botstrings.first_instructions}
                )
                send(botstrings.first_instructions, check=False)
            else:
                if chatbot.informations_requested:
                    question = ''
                    for i in chatbot.memory:
                        if i["role"] == "assistant":
                            question = i["content"]

                    print('QUESTION:', question)
                    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, looking_for=report.looking_for, question=question
                        )
                        if result:
                            pprint(result)
                        else:
                            send(
                                "I could not understand your message. Please try again.",
                                check=False,
                            )
                        looking_for = chatbot.ask_for_info(report)

                    else:
                        send(
                            "I could not understand your message. Please try again.",
                            check=False,
                        )

        else:
            print("Unknown message type")

        user_input = input(">>> ")

        if os.path.isfile(user_input):

            report.image = user_input
        else:
            print("User input is not a valid image file.")

general_bot = GeneralBot()