Add object attribute to Report class and update memory handling in Chatbot

bot.py

@@ -14,6 +14,7 @@ from pprint import pprint
 
 class Report:
     def __init__(self):
+        self.object = None
         self.description = None
         self.incident = None
         self.name = None
@@ -80,7 +81,7 @@ class BaseBot:
     def generate(self, message):
 
         message = re.sub(" +", " ", message)
-        
+
         if self.memory:
             messages = self.memory
             if messages[-1]["content"] != message:
@@ -135,20 +136,42 @@ class Chatbot(BaseBot):
         if looking_for != []:
             report.looking_for = looking_for
             # Ask for the information
-            general_bot.memory = [
-                {"role": "user", "content": 'Formulate a question asking for the following information: ["name", "age"]'},
-                {"role": "assistant", "content": "Can you tell me about the name and age of what you've found?"},
-            ]
             if looking_for == ["image"]:
-                question = 'image'
+                question = "image"
             else:
                 prompt = f"Formulate a question asking for the following information: {looking_for}"
                 general_bot.memory.append({"role": "user", "content": prompt})
                 question = general_bot.generate(prompt)["content"]
+                general_bot.memory = None
+                general_bot.memory = [
+                    {
+                        "role": "user",
+                        "content": 'Formulate a question asking for the following information: ["name", "age"]',
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "Can you tell me about the name and age of what you've found?",
+                    },
+                ]
+                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"Can you tell me about the name and age of the {report.object}?"
+
+                prompt = f"Formulate a question asking for the following information: {looking_for}"
+                general_bot.memory.append({"role": "user", "content": prompt})
+                question = general_bot.generate(prompt)["content"]
+                general_bot.memory = None
             return question
 
         else:
-            return 'done'
+            return "done"
 
 
 class CheckerBot(BaseBot):
@@ -163,19 +186,21 @@ class CheckerBot(BaseBot):
             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?
+        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 # TODO We need to fins a good way to check if the answer is an answer to the question.
+        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:
@@ -210,6 +235,41 @@ class CheckerBot(BaseBot):
         return answered
 
     def check_for_info(self, user_message, report: Report, looking_for: list, n_try=0):
+
+        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 = None
+
+            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:
@@ -267,6 +327,22 @@ class CheckerBot(BaseBot):
 
         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: