Convert messaget_type to enum

ae/core/agents/high_level_planner_agent.py

@@ -21,13 +21,11 @@ def __init__(self, config_list, user_proxy_agent:ConversableAgent): # type: igno
 
         user_ltm = self.__get_ltm()
         system_message = LLM_PROMPTS["PLANNER_AGENT_PROMPT"]
-        print(f">>> Planner system_message: {system_message}")
+
         if user_ltm: #add the user LTM to the system prompt if it exists
             user_ltm = "\n" + user_ltm
             system_message = Template(system_message).substitute(basic_user_information=user_ltm)
         system_message = system_message + "\n" + f"Today's date is {datetime.now().strftime('%d %B %Y')}"
-        print(f">>> Planner system_message: {system_message}")
-        print ("****************")
         self.agent = autogen.AssistantAgent(
             name="planner_agent",
             system_message=system_message,

ae/core/autogen_wrapper.py

@@ -1,7 +1,6 @@
 import asyncio
 import json
 import os
-import re
 import tempfile
 import traceback
 from string import Template
@@ -20,6 +19,7 @@
 from ae.core.agents.high_level_planner_agent import PlannerAgent  
 from ae.core.prompts import LLM_PROMPTS
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 from ae.utils.autogen_sequential_function_call import UserProxyAgent_SequentialFunctionExecution
 from ae.utils.response_parser import parse_response
 from ae.core.skills.get_url import geturl
@@ -104,16 +104,15 @@ def trigger_nested_chat(manager: autogen.ConversableAgent):
             next_step = content_json.get('next_step', None)
             plan = content_json.get('plan', None)
             if plan is not None:
-                notify_planner_messages(plan, level="plan")
+                notify_planner_messages(plan, message_type=MessageType.PLAN)
+
             if next_step is None: 
-                notify_planner_messages("Received no response, terminating..", level="others") # type: ignore
-                print("Trigger nested chat returned False")
+                notify_planner_messages("Received no response, terminating..", message_type=MessageType.INFO) # type: ignore
                 return False
             else:
-                notify_planner_messages(next_step, level="step") # type: ignore
+                notify_planner_messages(next_step, message_type=MessageType.STEP) # type: ignore
                 return True 
 
-
         def get_url() -> str:
             return asyncio.run(geturl())
 
@@ -135,7 +134,6 @@ def reflection_message(recipient, messages, sender, config): # type: ignore
             content_json = parse_response(last_message)
             next_step = content_json.get('next_step', None)
             if next_step is None: 
-                print ("Message to nested chat returned None")
                 return None
             else:
                 next_step = next_step.strip() +" " + get_url() # type: ignore
@@ -229,10 +227,8 @@ async def __create_user_delegate_agent(self) -> autogen.ConversableAgent:
         """
         def is_planner_termination_message(x: dict[str, str])->bool: # type: ignore
              should_terminate = False
-             print ("Checking Termination for Content:", x)
              function: Any = x.get("function", None)
              if function is not None:
-                 print("Function is not None")
                  return False
 
              content:Any = x.get("content", "") 
@@ -247,10 +243,9 @@ def is_planner_termination_message(x: dict[str, str])->bool: # type: ignore
                     if(_terminate == "yes"):
                         should_terminate = True
                         if final_response:
-                            print("Notifying Final Answer")
-                            notify_planner_messages(final_response, level="answer")
+                            notify_planner_messages(final_response, message_type=MessageType.ANSWER)
                 except json.JSONDecodeError:
-                    print("Error decoding JSON content")
+                    logger.error("Error decoding JSON response {content}. Terminating..")
                     should_terminate = True
 
              return should_terminate # type: ignore

ae/core/playwright_manager.py

@@ -13,7 +13,7 @@
 from ae.utils.js_helper import escape_js_message
 from ae.utils.js_helper import beautify_plan_message
 from ae.utils.logger import logger
-
+from ae.utils.ui_messagetype import MessageType
 # Enusres that playwright does not wait for font loading when taking screenshots. Reference: https://github.com/microsoft/playwright/issues/28995
 os.environ["PW_TEST_SCREENSHOT_NO_FONTS_READY"] = "1"
 
@@ -268,50 +268,47 @@ async def set_user_response_handler(self):
         await context.expose_function('user_response', self.receive_user_response) # type: ignore
 
 
-    async def notify_user(self, message: str, message_type: str = "step"):
+    async def notify_user(self, message: str, message_type: MessageType = MessageType.STEP):
         """
         Notify the user with a message.
 
         Args:
             message (str): The message to notify the user with.
-            message_type (str, optional): The type of message. Defaults to "step", other values are "plan", "action", "answer", "question", "info". 
+            message_type (enum, optional): Values can be 'PLAN', 'QUESTION', 'ANSWER', 'INFO', 'STEP'. Defaults to 'STEP'.
             To Do: Convert to Enum.
         """
 
-        safe_message_type = escape_js_message(message_type)
         if self.ui_manager.overlay_show_details == False:
-            if not (message_type == "plan" or message_type == "question" or message_type == "answer" or message_type == "info"):
+            if not (message_type == MessageType.PLAN or message_type == MessageType.QUESTION or message_type == MessageType.ANSWER or message_type == MessageType.INFO):
                 return
 
         if self.ui_manager.overlay_show_details == True:
-            if not (message_type == "plan" or message_type == "question" or message_type == "answer" or message_type == "info" or message_type == "step"):
+            if not (message_type == MessageType.PLAN  or message_type == MessageType.QUESTION  or message_type == MessageType.ANSWER or message_type == MessageType.INFO or message_type == MessageType.STEP):
                 return
 
-        if message_type == "plan":
+        if message_type == MessageType.PLAN:
             message = beautify_plan_message(message)
             message = "Plan: \n" + message
-        if message_type == "step":
+        if message_type == MessageType.STEP:
             if "confirm" in message.lower():
                 message = "Verify: " + message
             else:
                 message = "Next step: " + message
-        if message_type == "question":
+        if message_type == MessageType.QUESTION:
             message = "Question: " + message
 
-        if message_type == "answer":
+        if message_type == MessageType.ANSWER:
             message = "Response: " + message
         safe_message = escape_js_message(message)
         self.ui_manager.new_system_message(safe_message, message_type)
-
+        safe_message_type = escape_js_message(message_type.value)
         try:
             js_code = f"addSystemMessage({safe_message}, is_awaiting_user_response=false, message_type={safe_message_type});"
             page = await self.get_current_page()
             await page.evaluate(js_code)
-            logger.debug("User notification completed")
         except Exception as e:
-            print(f"Failed to notify user with message \"{message}\". However, most likey this will work itself out after the page loads: {e}")
             logger.debug(f"Failed to notify user with message \"{message}\". However, most likey this will work itself out after the page loads: {e}")
-
+        print("Notify Complete...")
     async def highlight_element(self, selector: str, add_highlight: bool):
         try:
             page: Page = await self.get_current_page()
@@ -355,7 +352,7 @@ async def prompt_user(self, message: str) -> str:
         page = await self.get_current_page()
 
         await self.ui_manager.show_overlay(page)
-        self.log_system_message(message) # add the message to history after the overlay is opened to avoid double adding it. add_system_message below will add it
+        self.log_system_message(message, MessageType.QUESTION) # add the message to history after the overlay is opened to avoid double adding it. add_system_message below will add it
 
         safe_message = escape_js_message(message)
 
@@ -398,7 +395,7 @@ async def take_screenshots(self, name: str, page: Page|None, full_page: bool = T
         try:
             await page.wait_for_load_state(state=load_state, timeout=take_snapshot_timeout) # type: ignore
             await page.screenshot(path=screenshot_path, full_page=full_page, timeout=take_snapshot_timeout, caret="initial", scale="device")
-            print(f"Screen shot saved to: {screenshot_path}")
+            logger.info(f"Screen shot saved to: {screenshot_path}")
         except Exception as e:
             logger.error(f"Failed to take screenshot and save to \"{screenshot_path}\". Error: {e}")
 
@@ -413,14 +410,14 @@ def log_user_message(self, message: str):
         self.ui_manager.new_user_message(message)
 
 
-    def log_system_message(self, message: str):
+    def log_system_message(self, message: str, type: MessageType = MessageType.STEP):
         """
         Log a system message.
 
         Args:
             message (str): The system message to log.
         """
-        self.ui_manager.new_system_message(message)
+        self.ui_manager.new_system_message(message, type)
 
     async def update_processing_state(self, processing_state: str):
         """

ae/core/post_process_responses.py

@@ -4,7 +4,7 @@
 import asyncio
 from ae.core.playwright_manager import PlaywrightManager
 from ae.utils.logger import logger
-
+from ae.utils.ui_messagetype import MessageType
 
 def final_reply_callback_user_proxy(recipient: autogen.ConversableAgent, messages: list[dict[str, Any]], sender: autogen.Agent, config: dict[str, Any]):
     """
@@ -35,8 +35,8 @@ def final_reply_callback_user_proxy(recipient: autogen.ConversableAgent, message
     return False, None
 
 
-def final_reply_callback_planner_agent(message:str, level:str = "plan"): # type: ignore 
+def final_reply_callback_planner_agent(message:str, message_type:MessageType = MessageType.STEP): # type: ignore 
             browser_manager = PlaywrightManager(browser_type='chromium', headless=False)
             loop = asyncio.get_event_loop()
-            loop.run_until_complete(browser_manager.notify_user(message, message_type=level))
+            loop.run_until_complete(browser_manager.notify_user(message, message_type=message_type))
             return False, None  # required to ensure the agent communication flow continues

ae/core/skills/click_using_selector.py

@@ -12,7 +12,7 @@
 from ae.utils.dom_mutation_observer import subscribe
 from ae.utils.dom_mutation_observer import unsubscribe
 from ae.utils.logger import logger
-
+from ae.utils.ui_messagetype import MessageType
 
 async def click(selector: Annotated[str, "The properly formed query selector string to identify the element for the click action (e.g. [mmid='114']). When \"mmid\" attribute is present, use it for the query selector."],
                 wait_before_execution: Annotated[float, "Optional wait time in seconds before executing the click event logic.", float] = 0.0) -> Annotated[str, "A message indicating success or failure of the click."]:
@@ -52,7 +52,7 @@ def detect_dom_changes(changes:str): # type: ignore
     await asyncio.sleep(0.1) # sleep for 100ms to allow the mutation observer to detect changes
     unsubscribe(detect_dom_changes)
     await browser_manager.take_screenshots(f"{function_name}_end", page)
-    await browser_manager.notify_user(result["summary_message"], message_type="action")
+    await browser_manager.notify_user(result["summary_message"], message_type=MessageType.ACTION)
 
     if dom_changes_detected:
         return f"Success: {result['summary_message']}.\n As a consequence of this action, new elements have appeared in view: {dom_changes_detected}. This means that the action to click {selector} is not yet executed and needs further interaction. Get all_fields DOM to complete the interaction."

ae/core/skills/enter_text_and_click.py

@@ -7,6 +7,7 @@
 from ae.core.skills.enter_text_using_selector import do_entertext
 from ae.core.skills.press_key_combination import do_press_key_combination
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 
 async def enter_text_and_click(
@@ -63,10 +64,10 @@ async def enter_text_and_click(
         do_press_key_combination_result = await do_press_key_combination(browser_manager, page, "Enter")
         if do_press_key_combination_result:
             result["detailed_message"] += f" Instead of click, pressed the Enter key successfully on element: \"{click_selector}\"."
-            await browser_manager.notify_user(f"Pressed the Enter key successfully on element: \"{click_selector}\".", message_type="action")
+            await browser_manager.notify_user(f"Pressed the Enter key successfully on element: \"{click_selector}\".", message_type=MessageType.ACTION)
         else:
             result["detailed_message"] += f" Clicking the same element after entering text in it, is of no value. Tried pressing the Enter key on element \"{click_selector}\" instead of click and failed."
-            await browser_manager.notify_user("Failed to press the Enter key on element \"{click_selector}\".", message_type="action")
+            await browser_manager.notify_user("Failed to press the Enter key on element \"{click_selector}\".", message_type=MessageType.ACTION)
     else:
         await browser_manager.highlight_element(click_selector, True)
 

ae/core/skills/enter_text_using_selector.py

@@ -13,6 +13,7 @@
 from ae.utils.dom_mutation_observer import subscribe
 from ae.utils.dom_mutation_observer import unsubscribe
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 
 @dataclass
@@ -126,7 +127,7 @@ def detect_dom_changes(changes:str): # type: ignore
 
     await browser_manager.take_screenshots(f"{function_name}_end", page)
 
-    await browser_manager.notify_user(result["summary_message"], message_type="action")
+    await browser_manager.notify_user(result["summary_message"], message_type=MessageType.ACTION)
     if dom_changes_detected:
         return f"{result['detailed_message']}.\n As a consequence of this action, new elements have appeared in view: {dom_changes_detected}. This means that the action of entering text {text_to_enter} is not yet executed and needs further interaction. Get all_fields DOM to complete the interaction."
     return result["detailed_message"]

ae/core/skills/get_dom_with_content_type.py

@@ -8,6 +8,7 @@
 from ae.utils.dom_helper import wait_for_non_loading_dom_state
 from ae.utils.get_detailed_accessibility_tree import do_get_accessibility_info
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 async def get_dom_with_content_type(
     content_type: Annotated[str, "The type of content to extract: 'text_only': Extracts the innerText of the highest element in the document and responds with text, or 'input_fields': Extracts the text input and button elements in the dom."]
@@ -70,7 +71,7 @@ async def get_dom_with_content_type(
 
     elapsed_time = time.time() - start_time
     logger.info(f"Get DOM Command executed in {elapsed_time} seconds")
-    await browser_manager.notify_user(user_success_message, message_type="action")
+    await browser_manager.notify_user(user_success_message, message_type=MessageType.ACTION)
     return extracted_data # type: ignore
 
 

ae/core/skills/get_url.py

@@ -1,7 +1,5 @@
 from typing import Annotated
 from ae.core.playwright_manager import PlaywrightManager
-from ae.utils.logger import logger
-
 
 async def geturl() -> Annotated[str, "Returns the full URL of the current active web site/page."]:
     """

ae/core/skills/get_user_input.py

@@ -14,12 +14,11 @@ async def get_user_input(questions: Annotated[List[str], "List of questions to a
     Returns:
     - Newline separated list of questions to ask the user
     """
-    print(f"Executing Get User Input Command {questions}")
+
     answers: dict[str, str] = {}
     browser_manager = PlaywrightManager(browser_type='chromium', headless=False)
     if browser_manager.ui_manager:
         for question in questions:
-            print(f"Question: {question}")
             answers[question] = await browser_manager.prompt_user(f"Question: {question}")
     else:
         answers = await answer_questions_over_cli(questions)

ae/core/skills/open_url.py

@@ -2,6 +2,7 @@
 from typing import Annotated
 from ae.core.playwright_manager import PlaywrightManager
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 
 async def openurl(url: Annotated[str, "The URL to navigate to. Value must include the protocol (http:// or https://)."],
@@ -34,7 +35,7 @@ async def openurl(url: Annotated[str, "The URL to navigate to. Value must includ
 
     await browser_manager.take_screenshots(f"{function_name}_end", page)
 
-    await browser_manager.notify_user(f"Opened URL: {url}", message_type="action")
+    await browser_manager.notify_user(f"Opened URL: {url}", message_type=MessageType.ACTION)
         # Get the page title
     title = await page.title()
     url=page.url

ae/core/skills/pdf_text_extractor.py

@@ -7,6 +7,7 @@
 from ae.config import PROJECT_TEMP_PATH
 from ae.core.playwright_manager import PlaywrightManager
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 
 async def extract_text_from_pdf(pdf_url: Annotated[str, "The URL of the PDF file to extract text from."]) -> Annotated[str, "All the text found in the PDF file."]:
@@ -35,7 +36,7 @@ async def extract_text_from_pdf(pdf_url: Annotated[str, "The URL of the PDF file
                     text += page_text + "\n"
         extracted_text = text.strip()
         word_count = len(extracted_text.split())
-        await browser_manager.notify_user(f"Extracted text from the PDF successfully. Found {word_count} words.", message_type="action")
+        await browser_manager.notify_user(f"Extracted text from the PDF successfully. Found {word_count} words.", message_type=MessageType.ACTION)
         return "Text found in the PDF:\n" + extracted_text
     except httpx.HTTPStatusError as e:
         logger.error(f"An error occurred while downloading the PDF from {pdf_url}: {str(e)}")

ae/core/skills/press_key_combination.py

@@ -8,6 +8,7 @@
 from ae.utils.dom_mutation_observer import subscribe  # type: ignore
 from ae.utils.dom_mutation_observer import unsubscribe  # type: ignore
 from ae.utils.logger import logger
+from ae.utils.ui_messagetype import MessageType
 
 
 async def press_key_combination(key_combination: Annotated[str, "The key to press, e.g., Enter, PageDown etc"]) -> str:
@@ -61,7 +62,7 @@ def detect_dom_changes(changes:str): # type: ignore
     if dom_changes_detected:
         return f"Key {key_combination} executed successfully.\n As a consequence of this action, new elements have appeared in view:{dom_changes_detected}. This means that the action is not yet executed and needs further interaction. Get all_fields DOM to complete the interaction."
 
-    await browser_manager.notify_user(f"Key {key_combination} executed successfully", message_type="action")
+    await browser_manager.notify_user(f"Key {key_combination} executed successfully", message_type=MessageType.ACTION)
     return f"Key {key_combination} executed successfully"
 
 

ae/core/system_orchestrator.py

@@ -99,7 +99,7 @@ async def process_command(self, command: str):
         Args:
             command (str): The command to process.
         """
-        print(f"Received command: {command}")
+        logger.info(f"Received command: {command}")
         if command.lower() == 'exit':
             await self.shutdown()
             return