chore: import upstream snapshot with attribution

gui_agents/s2/WAA_setup.md

@@ -0,0 +1,136 @@
+# Introduction
+
+This is the WindowsAgentArena (WAA) setup with Agent S2 (and beyond). Why do we need a setup guide? Despite the thorough [README.md](https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file "https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file"), we have to include our code into their repository _and_ fix up a number of setup issues from the WAA environment. Sadly, this isn’t the most straightforward.
+
+# Initial WAA Setup
+
+The initial WAA setup is straightforward. Follow the [README.md](https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file "https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file") on their repository. After you’ve finished this, try running `run-local.sh`. This will start up an experiment with their default `Navi` agent. At this point, the environment is _sufficient to run evaluation_, but it’s incomplete and thus the evaluation won’t be exactly correct due to environment issues.
+
+![](./images/waa_setup/fig1.png)
+
+Figure 1: Bash script chain of execution.
+
+While we’re at it, look to understand the following things:
+
+-   the entire README.md (especially the [Bring Your Own Agent guide](https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file#-byoa-bring-your-own-agent "https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file#-byoa-bring-your-own-agent"))
+    
+-   the _long_ chain of bash scripts that start the run (Figure 1)
+    
+-   the `run.py` to see how the agent/environment are instantiated and used together
+    
+-   the folder structure of the repository and the purpose of each folder
+    
+
+# Fixing Setup Issues
+
+By now, your WAA environment should be set up to run locally. There are two major problems:
+
+-   setup issues
+    
+-   the VM persists across examples (it won’t reset after every example is completed which may make evaluation unfair)
+    
+
+Let’s tackle the first one: setup issues.
+
+### Office Apps Aren’t Installed
+
+The first issue I ran into was the office apps aren’t installed. Why is that? Turns out all apps installed in the VM during the initial setup stage install via the links from this [file](https://github.com/microsoft/WindowsAgentArena/blob/main/src/win-arena-container/vm/setup/tools_config.json "https://github.com/microsoft/WindowsAgentArena/blob/main/src/win-arena-container/vm/setup/tools_config.json") (`tools_config.json`). At the time of writing this, only the office links do not work. Try out all the links to make sure they work. If the links do not lead to a download (and some error occurs), then that app was not installed in the VM. What do we do? Two options:
+
+-   redo the entire initial setup stage (time consuming; ~**4** hours for me and even then, it would just not work a lot of the times; ideally, WAA is setup on Linux as I’ve had no issues so far with it)
+    
+-   Enter the VM and install the apps manually (easier and faster)
+    
+
+We’ll do the second approach.
+
+You can access the VM via `https://localhost:8006`. You can turn the VM on by `run-local.sh`. There’s probably a better/faster way to do it, but this doesn’t take too much time anyways (~**1-2** mins). After the VM has started, enter the VM (the agent may be trying to take actions, but you can either just override the action in `run.py` with `import time; time.sleep(10000)` [here](https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/src/win-arena-container/client/lib_run_single.py#L58 "https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/src/win-arena-container/client/lib_run_single.py#L58") or fight the agent for control of the VM!).
+
+Inside the VM, navigate to their [download page](https://www.libreoffice.org/download/download-libreoffice/ "https://www.libreoffice.org/download/download-libreoffice/") and download the latest LibreOffice version. After it’s downloaded, complete the setup wizard and make sure to delete the downloaded `*.msi` file in the VM. Finally, test the download by opening up LibreOffice Writer and Calc.
+
+### Google Chrome Pop-ups
+
+In Google Chrome, there a couple unexpected pop-ups.
+
+![](./images/waa_setup/fig2.png)
+
+Figure 2: Pop-ups on Chrome.
+
+Close all these pop-ups and [make Google Chrome your default web browser](https://support.google.com/chrome/answer/95417?hl=en&co=GENIE.Platform%3DDesktop#zippy=%2Cmac%2Cwindows "https://support.google.com/chrome/answer/95417?hl=en&co=GENIE.Platform%3DDesktop#zippy=%2Cmac%2Cwindows").
+
+### VSCode Pop-ups
+
+This isn’t as important, but there are a couple initial pop-ups in VSCode that you can close.
+
+### Note: `set_cell_values`
+
+_Important if you’re using_ `set_cell_values`
+
+Agent S2 uses a special grounding function called `set_cell_values` that takes advantage of the `soffice` CLI and `unotools` [Python library](https://pypi.org/project/unotools/ "https://pypi.org/project/unotools/"). TL; DR, this function lets the agent set the cell values for a given spreadsheet and sheet.
+
+For this function to work on WAA, the set up is a bit messy…
+
+1.  Connect into the VM
+    
+2.  Open up a terminal and run `python --version`, you should see you’re using the GIMP Python which is `2.x`. This won’t let you use the `soffice` CLI or `import uno` in Python code.
+    
+3.  In the `Desktop` directory within a terminal, do `pip freeze > requirements.txt` to save all the PYPI libraries from the GIMP Python to a `requirements.txt`.
+    
+4.  Configuring Python path to LibreOffice’s Python
+    
+    1.  In the File Explorer, locate the `python.exe` file from LibreOffice. You can do this with `where python`. Copy this path.
+        
+    2.  In the Search bar in the bottom task bar inside the VM, search for “environment variables”.
+        
+    3.  Click on “Environment Variables” and click on “Path” under “System variables”. Paste the copied path from step (a) into there and ensure this path is _above_ the GIMP Python path so it takes precedence.
+        
+    4.  Reopen a terminal and run `soffice` to ensure it is now working. Create a temporary python file and ensure `import uno` works.
+        
+5.  LibreOffice’s Python should be `3.10` or above. However, it does not come with pip. To install pip, download this [file](https://bootstrap.pypa.io/get-pip.py "https://bootstrap.pypa.io/get-pip.py") and execute `python get-pip.py` to install it. Ensure the `python` here is LibreOffice’s Python. Next, install `pip install -r requirements.txt` using the `requirements.txt` from step 3. This is to ensure LibreOffice’s Python has all the dependencies needed for evaluation (pyautogui, etc).
+    
+6.  Clean up all installer files. Then, inside the [WAA repository code](https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/src/win-arena-container/client/desktop_env/controllers/python.py#L193 "https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/src/win-arena-container/client/desktop_env/controllers/python.py#L193"), change this line
+    
+
+`command_list = ["python", "-c", self.pkgs_prefix.format(command=command)]`
+
+to:
+
+`command_list = ["absolute/path/to/libreoffice/python", "-c", self.pkgs_prefix.format(command=command)]`
+
+This ensures that the subprocess running in the flask server inside the VM will use that specific Python version.
+
+### Double Checking…
+
+Double check all apps can be used and no unexpected pop-ups or issues are in the way. Any apps you open make sure to close them upon finishing your clean-up. Make sure any installation files you have in `Downloads` are deleted (and removed from Recycle Bin) to keep the environment clean. At the end, this is our **golden image**. You may want to save a copy of this VM somewhere safe so that you can always copy it back into the WAA repository to be reused (refer to [this](https://github.com/microsoft/WindowsAgentArena/tree/main?tab=readme-ov-file#additional-notes "https://github.com/microsoft/WindowsAgentArena/tree/main?tab=readme-ov-file#additional-notes")).
+
+# Set up Agent S2 with WAA Locally
+
+Take the time to understand the [Agent-S repository](https://github.com/simular-ai/Agent-S "https://github.com/simular-ai/Agent-S").
+
+1.  Instead of following the [README.md](https://github.com/simular-ai/Agent-S/blob/main/README.md "https://github.com/simular-ai/Agent-S/blob/main/README.md") for Agent S2, you need to clone the repository then `pip install -r requirements.txt`
+    
+2.  Move the s2 folder to the [mm_agents](https://github.com/microsoft/WindowsAgentArena/tree/main/src/win-arena-container/client/mm_agents "https://github.com/microsoft/WindowsAgentArena/tree/main/src/win-arena-container/client/mm_agents") folder in WAA. Follow the [Bring Your Own Agent guide](https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file#-byoa-bring-your-own-agent "https://github.com/microsoft/WindowsAgentArena?tab=readme-ov-file#-byoa-bring-your-own-agent").
+    
+    1.  You will need to move the `agent_s.py` file out to the `s2` folder and update all the relevant import statements
+        
+3.  Make the necessary changes in `run.py` and `lib_run_single.py` to accommodate Agent S2 (replace the Navi Agent with Agent S2).
+    
+4.  Test it by running the experiments! Don’t forget when you do `run-local.sh`, now you need to specify Agent S2 instead of the navi agent `agent="agent_s"`.
+    
+5.  You may have some import errors and these libraries need to be installed inside the `winarena` container (I think). You can just add the pip install commands to the bash script where the error stems from (hacky).
+    
+
+#### Perplexica
+
+There may be a Perplexica issue. The Perplexica URL must be configured so that the agent in the `winarena` Docker container can communicate with `localhost:3001` which is the forwarded port from the Perplexica container. On Mac/Windows this can be fixed by changing the `PERPLEXICA_URL` to `http://host.docker.internal:3001/api/search` . On Linux, I just disabled it… I haven’t tried, but you can add `--add-host=host.docker.internal:host-gateway` as a flag to the docker command [here](https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/scripts/run.sh#L223 "https://github.com/microsoft/WindowsAgentArena/blob/6d39ed88c545a0d40a7a02e39b928e278df7332b/scripts/run.sh#L223") (run.sh). This may let you use `http://host.docker.internal:3001/api/search` as the `PERPLEXICA_URL`
+
+# Agent S2 with WAA on Azure
+
+1.  Ensure you have:
+    
+    1.  a **clean copy** of the golden image
+        
+    2.  the correct Azure subscription (so you’re not using your own payment method)
+        
+2.  Follow the Azure deployment in the [README.md](https://github.com/microsoft/WindowsAgentArena/blob/main/README.md "https://github.com/microsoft/WindowsAgentArena/blob/main/README.md").
+    
+3.  Test it! If this works, then we have a resettable golden image and WAA can be ran in parallel, making evaluation much _much_ faster! Good luck!

gui_agents/s2/agents/agent_s.py

@@ -0,0 +1,425 @@
+import json
+import logging
+import os
+import platform
+from typing import Dict, List, Optional, Tuple
+
+from gui_agents.s2.agents.grounding import ACI
+from gui_agents.s2.agents.worker import Worker
+from gui_agents.s2.agents.manager import Manager
+from gui_agents.s2.utils.common_utils import Node
+from gui_agents.utils import download_kb_data
+from gui_agents.s2.core.engine import (
+    OpenAIEmbeddingEngine,
+    GeminiEmbeddingEngine,
+    AzureOpenAIEmbeddingEngine,
+)
+
+logger = logging.getLogger("desktopenv.agent")
+
+
+class UIAgent:
+    """Base class for UI automation agents"""
+
+    def __init__(
+        self,
+        engine_params: Dict,
+        grounding_agent: ACI,
+        platform: str = platform.system().lower(),
+        action_space: str = "pyautogui",
+        observation_type: str = "a11y_tree",
+        search_engine: str = "perplexica",
+    ):
+        """Initialize UIAgent
+
+        Args:
+            engine_params: Configuration parameters for the LLM engine
+            grounding_agent: Instance of ACI class for UI interaction
+            platform: Operating system platform (macos, linux, windows)
+            action_space: Type of action space to use (pyautogui, aci)
+            observation_type: Type of observations to use (a11y_tree, mixed)
+            engine: Search engine to use (perplexica, LLM)
+        """
+        self.engine_params = engine_params
+        self.grounding_agent = grounding_agent
+        self.platform = platform
+        self.action_space = action_space
+        self.observation_type = observation_type
+        self.engine = search_engine
+
+    def reset(self) -> None:
+        """Reset agent state"""
+        pass
+
+    def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
+        """Generate next action prediction
+
+        Args:
+            instruction: Natural language instruction
+            observation: Current UI state observation
+
+        Returns:
+            Tuple containing agent info dictionary and list of actions
+        """
+        pass
+
+    def update_narrative_memory(self, trajectory: str) -> None:
+        """Update narrative memory with task trajectory
+
+        Args:
+            trajectory: String containing task execution trajectory
+        """
+        pass
+
+    def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str:
+        """Update episodic memory with subtask trajectory
+
+        Args:
+            meta_data: Metadata about current subtask execution
+            subtask_trajectory: String containing subtask execution trajectory
+
+        Returns:
+            Updated subtask trajectory
+        """
+        pass
+
+
+class AgentS2(UIAgent):
+    """Agent that uses hierarchical planning and directed acyclic graph modeling for UI automation"""
+
+    def __init__(
+        self,
+        engine_params: Dict,
+        grounding_agent: ACI,
+        platform: str = platform.system().lower(),
+        action_space: str = "pyautogui",
+        observation_type: str = "mixed",
+        search_engine: Optional[str] = None,
+        memory_root_path: str = os.getcwd(),
+        use_default_kb: bool = False,
+        memory_folder_name: str = "kb_s2",
+        kb_release_tag: str = "v0.2.2",
+        embedding_engine_type: str = "openai",
+        embedding_engine_params: Dict = {},
+    ):
+        """Initialize AgentS2
+
+        Args:
+            engine_params: Configuration parameters for the LLM engine
+            grounding_agent: Instance of ACI class for UI interaction
+            platform: Operating system platform (darwin, linux, windows)
+            action_space: Type of action space to use (pyautogui, other)
+            observation_type: Type of observations to use (a11y_tree, screenshot, mixed)
+            search_engine: Search engine to use (LLM, perplexica)
+            use_default_kb: True to use the default OpenAI kb.
+            memory_root_path: Path to memory directory. Defaults to current working directory.
+            memory_folder_name: Name of memory folder. Defaults to "kb_s2".
+            kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
+            embedding_engine_type: Embedding engine to use for knowledge base. Defaults to "openai". Supports "openai" and "gemini".
+            embedding_engine_params: Parameters for embedding engine. Defaults to {}.
+        """
+        super().__init__(
+            engine_params,
+            grounding_agent,
+            platform,
+            action_space,
+            observation_type,
+            search_engine,
+        )
+
+        self.memory_root_path = memory_root_path
+        self.memory_folder_name = memory_folder_name
+        self.kb_release_tag = kb_release_tag
+
+        # Initialize agent's knowledge base on user's current working directory.
+        self.local_kb_path = os.path.join(
+            self.memory_root_path, self.memory_folder_name
+        )
+
+        if use_default_kb:
+            if not os.path.exists(os.path.join(self.local_kb_path, self.platform)):
+                print("Downloading Agent S2's default knowledge base...")
+                download_kb_data(
+                    version="s2",
+                    release_tag=kb_release_tag,
+                    download_dir=self.local_kb_path,
+                    platform=self.platform,
+                )
+                print(
+                    f"Successfully completed download of knowledge base for version s2, tag {self.kb_release_tag}, platform {self.platform}."
+                )
+            else:
+                print(
+                    f"Path local_kb_path {self.local_kb_path} already exists. Skipping download."
+                )
+                print(
+                    f"If you'd like to re-download the initial knowledge base, please delete the existing knowledge base at {self.local_kb_path}."
+                )
+                print(
+                    "Note, the knowledge is continually updated during inference. Deleting the knowledge base will wipe out all experience gained since the last knowledge base download."
+                )
+
+        if embedding_engine_type == "openai":
+            self.embedding_engine = OpenAIEmbeddingEngine(**embedding_engine_params)
+        elif embedding_engine_type == "gemini":
+            self.embedding_engine = GeminiEmbeddingEngine(**embedding_engine_params)
+        elif embedding_engine_type == "azure":
+            self.embedding_engine = AzureOpenAIEmbeddingEngine(
+                **embedding_engine_params
+            )
+
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset agent state and initialize components"""
+        # Initialize core components
+        self.planner = Manager(
+            engine_params=self.engine_params,
+            grounding_agent=self.grounding_agent,
+            local_kb_path=self.local_kb_path,
+            embedding_engine=self.embedding_engine,
+            search_engine=self.engine,
+            platform=self.platform,
+        )
+        self.executor = Worker(
+            engine_params=self.engine_params,
+            grounding_agent=self.grounding_agent,
+            local_kb_path=self.local_kb_path,
+            embedding_engine=self.embedding_engine,
+            platform=self.platform,
+        )
+
+        # Reset state variables
+        self.requires_replan: bool = True
+        self.needs_next_subtask: bool = True
+        self.step_count: int = 0
+        self.turn_count: int = 0
+        self.failure_subtask: Optional[Node] = None
+        self.should_send_action: bool = False
+        self.completed_tasks: List[Node] = []
+        self.current_subtask: Optional[Node] = None
+        self.subtasks: List[Node] = []
+        self.search_query: str = ""
+        self.subtask_status: str = "Start"
+
+    def reset_executor_state(self) -> None:
+        """Reset executor and step counter"""
+        self.executor.reset()
+        self.step_count = 0
+
+    def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
+        # Initialize the three info dictionaries
+        planner_info = {}
+        executor_info = {}
+        evaluator_info = {
+            "obs_evaluator_response": "",
+            "num_input_tokens_evaluator": 0,
+            "num_output_tokens_evaluator": 0,
+            "evaluator_cost": 0.0,
+        }
+        actions = []
+
+        # If the DONE response by the executor is for a subtask, then the agent should continue with the next subtask without sending the action to the environment
+        while not self.should_send_action:
+            self.subtask_status = "In"
+            # If replan is true, generate a new plan. True at start, after a failed plan, or after subtask completion
+            if self.requires_replan:
+                logger.info("(RE)PLANNING...")
+                planner_info, self.subtasks = self.planner.get_action_queue(
+                    instruction=instruction,
+                    observation=observation,
+                    failed_subtask=self.failure_subtask,
+                    completed_subtasks_list=self.completed_tasks,
+                    remaining_subtasks_list=self.subtasks,
+                )
+
+                self.requires_replan = False
+                if "search_query" in planner_info:
+                    self.search_query = planner_info["search_query"]
+                else:
+                    self.search_query = ""
+
+            # use the exectuor to complete the topmost subtask
+            if self.needs_next_subtask:
+                logger.info("GETTING NEXT SUBTASK...")
+
+                # this can be empty if the DAG planner deems that all subtasks are completed
+                if len(self.subtasks) <= 0:
+                    self.requires_replan = True
+                    self.needs_next_subtask = True
+                    self.failure_subtask = None
+                    self.completed_tasks.append(self.current_subtask)
+
+                    # reset executor state
+                    self.reset_executor_state()
+                    self.should_send_action = True
+                    self.subtask_status = "Done"
+                    executor_info = {
+                        "executor_plan": "agent.done()",
+                        "plan_code": "agent.done()",
+                        "reflection": "agent.done()",
+                    }
+                    actions = ["DONE"]
+                    break
+
+                self.current_subtask = self.subtasks.pop(0)
+                logger.info(f"NEXT SUBTASK: {self.current_subtask}")
+                self.needs_next_subtask = False
+                self.subtask_status = "Start"
+
+            # get the next action from the executor
+            executor_info, actions = self.executor.generate_next_action(
+                instruction=instruction,
+                search_query=self.search_query,
+                subtask=self.current_subtask.name,
+                subtask_info=self.current_subtask.info,
+                future_tasks=self.subtasks,
+                done_task=self.completed_tasks,
+                obs=observation,
+            )
+
+            self.step_count += 1
+
+            # set the should_send_action flag to True if the executor returns an action
+            self.should_send_action = True
+
+            # replan on failure
+            if "FAIL" in actions:
+                self.requires_replan = True
+                self.needs_next_subtask = True
+
+                # assign the failed subtask
+                self.failure_subtask = self.current_subtask
+
+                # reset the step count, executor, and evaluator
+                self.reset_executor_state()
+
+                # if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
+                if self.subtasks:
+                    self.should_send_action = False
+
+            # replan on subtask completion
+            elif "DONE" in actions:
+                self.requires_replan = True
+                self.needs_next_subtask = True
+                self.failure_subtask = None
+                self.completed_tasks.append(self.current_subtask)
+
+                # reset the step count, executor, and evaluator
+                self.reset_executor_state()
+
+                # if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
+                if self.subtasks:
+                    self.should_send_action = False
+                self.subtask_status = "Done"
+
+            self.turn_count += 1
+
+        # reset the should_send_action flag for next iteration
+        self.should_send_action = False
+
+        # concatenate the three info dictionaries
+        info = {
+            **{
+                k: v
+                for d in [planner_info or {}, executor_info or {}, evaluator_info or {}]
+                for k, v in d.items()
+            }
+        }
+        info.update(
+            {
+                "subtask": self.current_subtask.name,
+                "subtask_info": self.current_subtask.info,
+                "subtask_status": self.subtask_status,
+            }
+        )
+
+        return info, actions
+
+    def update_narrative_memory(self, trajectory: str) -> None:
+        """Update narrative memory from task trajectory
+
+        Args:
+            trajectory: String containing task execution trajectory
+        """
+        try:
+            reflection_path = os.path.join(
+                self.local_kb_path, self.platform, "narrative_memory.json"
+            )
+            try:
+                reflections = json.load(open(reflection_path))
+            except:
+                reflections = {}
+
+            if self.search_query not in reflections:
+                reflection = self.planner.summarize_narrative(trajectory)
+                reflections[self.search_query] = reflection
+
+            with open(reflection_path, "w") as f:
+                json.dump(reflections, f, indent=2)
+
+        except Exception as e:
+            logger.error(f"Failed to update narrative memory: {e}")
+
+    def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str:
+        """Update episodic memory from subtask trajectory
+
+        Args:
+            meta_data: Metadata about current subtask execution
+            subtask_trajectory: String containing subtask execution trajectory
+
+        Returns:
+            Updated subtask trajectory
+        """
+        subtask = meta_data["subtask"]
+        subtask_info = meta_data["subtask_info"]
+        subtask_status = meta_data["subtask_status"]
+        # Handle subtask trajectory
+        if subtask_status == "Start" or subtask_status == "Done":
+            # If it's a new subtask start, finalize the previous subtask trajectory if it exists
+            if subtask_trajectory:
+                subtask_trajectory += "\nSubtask Completed.\n"
+                subtask_key = subtask_trajectory.split(
+                    "\n----------------------\n\nPlan:\n"
+                )[0]
+                try:
+                    subtask_path = os.path.join(
+                        self.local_kb_path, self.platform, "episodic_memory.json"
+                    )
+                    kb = json.load(open(subtask_path))
+                except:
+                    kb = {}
+                if subtask_key not in kb.keys():
+                    subtask_summarization = self.planner.summarize_episode(
+                        subtask_trajectory
+                    )
+                    kb[subtask_key] = subtask_summarization
+                else:
+                    subtask_summarization = kb[subtask_key]
+                logger.info("subtask_key: %s", subtask_key)
+                logger.info("subtask_summarization: %s", subtask_summarization)
+                with open(subtask_path, "w") as fout:
+                    json.dump(kb, fout, indent=2)
+                # Reset for the next subtask
+                subtask_trajectory = ""
+            # Start a new subtask trajectory
+            subtask_trajectory = (
+                "Task:\n"
+                + self.search_query
+                + "\n\nSubtask: "
+                + subtask
+                + "\nSubtask Instruction: "
+                + subtask_info
+                + "\n----------------------\n\nPlan:\n"
+                + meta_data["executor_plan"]
+                + "\n"
+            )
+        elif subtask_status == "In":
+            # Continue appending to the current subtask trajectory if it's still ongoing
+            subtask_trajectory += (
+                "\n----------------------\n\nPlan:\n"
+                + meta_data["executor_plan"]
+                + "\n"
+            )
+
+        return subtask_trajectory

gui_agents/s2/agents/grounding.py

@@ -0,0 +1,600 @@
+import ast
+import re
+from collections import defaultdict
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pytesseract
+from PIL import Image
+from pytesseract import Output
+
+from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
+from gui_agents.s2.core.mllm import LMMAgent
+from gui_agents.s2.utils.common_utils import (
+    call_llm_safe,
+    parse_single_code_from_string,
+)
+
+
+class ACI:
+    def __init__(self):
+        self.notes: List[str] = []
+
+
+# Agent action decorator
+def agent_action(func):
+    func.is_agent_action = True
+    return func
+
+
+UBUNTU_APP_SETUP = f"""import subprocess;
+import difflib;
+import pyautogui;
+pyautogui.press('escape');
+time.sleep(0.5);
+output = subprocess.check_output(['wmctrl', '-lx']);
+output = output.decode('utf-8').splitlines();
+window_titles = [line.split(None, 4)[2] for line in output];
+closest_matches = difflib.get_close_matches('APP_NAME', window_titles, n=1, cutoff=0.1);
+if closest_matches:
+    closest_match = closest_matches[0];
+    for line in output:
+        if closest_match in line:
+            window_id = line.split()[0]
+            break;
+subprocess.run(['wmctrl', '-ia', window_id])
+subprocess.run(['wmctrl', '-ir', window_id, '-b', 'add,maximized_vert,maximized_horz'])
+"""
+
+
+SET_CELL_VALUES_CMD = """import uno
+import subprocess
+
+def identify_document_type(component):
+    if component.supportsService("com.sun.star.sheet.SpreadsheetDocument"):
+        return "Calc"
+
+    if component.supportsService("com.sun.star.text.TextDocument"):
+        return "Writer"
+
+    if component.supportsService("com.sun.star.sheet.PresentationDocument"):
+        return "Impress"
+
+    return None
+
+def cell_ref_to_indices(cell_ref):
+    column_letters = ''.join(filter(str.isalpha, cell_ref))
+    row_number = ''.join(filter(str.isdigit, cell_ref))
+
+    col = sum((ord(char.upper()) - ord('A') + 1) * (26**idx) for idx, char in enumerate(reversed(column_letters))) - 1
+    row = int(row_number) - 1
+    return col, row
+
+def set_cell_values(new_cell_values: dict[str, str], app_name: str = "Untitled 1", sheet_name: str = "Sheet1"):
+    new_cell_values_idx = {{}}
+    for k, v in new_cell_values.items():
+        try:
+            col, row = cell_ref_to_indices(k)
+        except:
+            col = row = None
+
+        if col is not None and row is not None:
+            new_cell_values_idx[(col, row)] = v
+
+    # Clean up previous TCP connections.
+    subprocess.run(
+        'echo \"password\" | sudo -S ss --kill --tcp state TIME-WAIT sport = :2002',
+        shell=True,
+        check=True,
+        text=True,
+        capture_output=True
+    )
+
+    # Dynamically allow soffice to listen on port 2002.
+    subprocess.run(
+        [
+            "soffice",
+            "--accept=socket,host=localhost,port=2002;urp;StarOffice.Service"
+        ]
+    )
+
+    local_context = uno.getComponentContext()
+    resolver = local_context.ServiceManager.createInstanceWithContext(
+        "com.sun.star.bridge.UnoUrlResolver", local_context
+    )
+    context = resolver.resolve(
+        f"uno:socket,host=localhost,port=2002;urp;StarOffice.ComponentContext"
+    )
+    desktop = context.ServiceManager.createInstanceWithContext(
+        "com.sun.star.frame.Desktop", context
+    )
+
+    # Collect all LibreOffice-related opened windows.
+    documents = []
+    for i, component in enumerate(desktop.Components):
+        title = component.Title
+        doc_type = identify_document_type(component)
+        documents.append((i, component, title, doc_type))
+
+    # Find the LibreOffice Calc app and the sheet of interest.
+    spreadsheet = [doc for doc in documents if doc[3] == "Calc"]
+    selected_spreadsheet = [doc for doc in spreadsheet if doc[2] == app_name]
+    if spreadsheet:
+        try:
+            if selected_spreadsheet:
+                spreadsheet = selected_spreadsheet[0][1]
+            else:
+                spreadsheet = spreadsheet[0][1]
+
+            sheet = spreadsheet.Sheets.getByName(sheet_name)
+        except:
+            raise ValueError(f"Could not find sheet {{sheet_name}} in {{app_name}}.")
+
+        for (col, row), value in new_cell_values_idx.items():
+            cell = sheet.getCellByPosition(col, row)
+
+            # Set the cell value.
+            if isinstance(value, (int, float)):
+                cell.Value = value
+            elif isinstance(value, str):
+                if value.startswith("="):
+                    cell.Formula = value
+                else:
+                    cell.String = value
+            elif isinstance(value, bool):
+                cell.Value = 1 if value else 0
+            elif value is None:
+                cell.clearContents(0)
+            else:
+                raise ValueError(f"Unsupported cell value type: {{type(value)}}")
+
+    else:
+        raise ValueError(f"Could not find LibreOffice Calc app corresponding to {{app_name}}.")
+
+set_cell_values(new_cell_values={cell_values}, app_name="{app_name}", sheet_name="{sheet_name}")        
+"""
+
+
+# ACI primitives are parameterized by description, and coordinate generation uses a pretrained grounding model
+class OSWorldACI(ACI):
+    def __init__(
+        self,
+        platform: str,
+        engine_params_for_generation: Dict,
+        engine_params_for_grounding: Dict,
+        width: int = 1920,
+        height: int = 1080,
+    ):
+        self.platform = (
+            platform  # Dictates how the switch_applications agent action works.
+        )
+
+        # Configure scaling
+        self.width = width
+        self.height = height
+
+        # Maintain state for save_to_knowledge
+        self.notes = []
+
+        # Coordinates used during ACI execution
+        self.coords1 = None
+        self.coords2 = None
+
+        # Configure the visual grounding model responsible for coordinate generation
+        self.grounding_model = LMMAgent(engine_params_for_grounding)
+        self.engine_params_for_grounding = engine_params_for_grounding
+
+        # Configure text grounding agent
+        self.text_span_agent = LMMAgent(
+            engine_params=engine_params_for_generation,
+            system_prompt=PROCEDURAL_MEMORY.PHRASE_TO_WORD_COORDS_PROMPT,
+        )
+
+    # Given the state and worker's referring expression, use the grounding model to generate (x,y)
+    def generate_coords(self, ref_expr: str, obs: Dict) -> List[int]:
+
+        # Reset the grounding model state
+        self.grounding_model.reset()
+
+        # Configure the context, UI-TARS demo does not use system prompt
+        prompt = f"Query:{ref_expr}\nOutput only the coordinate of one point in your response.\n"
+        self.grounding_model.add_message(
+            text_content=prompt, image_content=obs["screenshot"], put_text_last=True
+        )
+
+        # Generate and parse coordinates
+        response = call_llm_safe(self.grounding_model)
+        print("RAW GROUNDING MODEL RESPONSE:", response)
+        numericals = re.findall(r"\d+", response)
+        assert len(numericals) >= 2
+        return [int(numericals[0]), int(numericals[1])]
+
+    # Calls pytesseract to generate word level bounding boxes for text grounding
+    def get_ocr_elements(self, b64_image_data: str) -> Tuple[str, List]:
+        image = Image.open(BytesIO(b64_image_data))
+        image_data = pytesseract.image_to_data(image, output_type=Output.DICT)
+
+        # Clean text by removing leading and trailing spaces and non-alphabetical characters, but keeping punctuation
+        for i, word in enumerate(image_data["text"]):
+            image_data["text"][i] = re.sub(
+                r"^[^a-zA-Z\s.,!?;:\-\+]+|[^a-zA-Z\s.,!?;:\-\+]+$", "", word
+            )
+
+        ocr_elements = []
+        ocr_table = "Text Table:\nWord id\tText\n"
+        # Obtain the <id, text, group number, word number> for each valid element
+        grouping_map = defaultdict(list)
+        ocr_id = 0
+        for i in range(len(image_data["text"])):
+            block_num = image_data["block_num"][i]
+            if image_data["text"][i]:
+                grouping_map[block_num].append(image_data["text"][i])
+                ocr_table += f"{ocr_id}\t{image_data['text'][i]}\n"
+                ocr_elements.append(
+                    {
+                        "id": ocr_id,
+                        "text": image_data["text"][i],
+                        "group_num": block_num,
+                        "word_num": len(grouping_map[block_num]),
+                        "left": image_data["left"][i],
+                        "top": image_data["top"][i],
+                        "width": image_data["width"][i],
+                        "height": image_data["height"][i],
+                    }
+                )
+                ocr_id += 1
+
+        return ocr_table, ocr_elements
+
+    # Given the state and worker's text phrase, generate the coords of the first/last word in the phrase
+    def generate_text_coords(
+        self, phrase: str, obs: Dict, alignment: str = ""
+    ) -> List[int]:
+
+        ocr_table, ocr_elements = self.get_ocr_elements(obs["screenshot"])
+
+        alignment_prompt = ""
+        if alignment == "start":
+            alignment_prompt = "**Important**: Output the word id of the FIRST word in the provided phrase.\n"
+        elif alignment == "end":
+            alignment_prompt = "**Important**: Output the word id of the LAST word in the provided phrase.\n"
+
+        # Load LLM prompt
+        self.text_span_agent.reset()
+        self.text_span_agent.add_message(
+            alignment_prompt + "Phrase: " + phrase + "\n" + ocr_table, role="user"
+        )
+        self.text_span_agent.add_message(
+            "Screenshot:\n", image_content=obs["screenshot"], role="user"
+        )
+
+        # Obtain the target element
+        response = call_llm_safe(self.text_span_agent)
+        print("TEXT SPAN AGENT RESPONSE:", response)
+        numericals = re.findall(r"\d+", response)
+        if len(numericals) > 0:
+            text_id = int(numericals[-1])
+        else:
+            text_id = 0
+        elem = ocr_elements[text_id]
+
+        # Compute the element coordinates
+        if alignment == "start":
+            coords = [elem["left"], elem["top"] + (elem["height"] // 2)]
+        elif alignment == "end":
+            coords = [elem["left"] + elem["width"], elem["top"] + (elem["height"] // 2)]
+        else:
+            coords = [
+                elem["left"] + (elem["width"] // 2),
+                elem["top"] + (elem["height"] // 2),
+            ]
+        return coords
+
+    # Takes a description based action and assigns the coordinates for any coordinate based action
+    # Raises an error if function can't be parsed
+    def assign_coordinates(self, plan: str, obs: Dict):
+
+        # Reset coords from previous action generation
+        self.coords1, self.coords2 = None, None
+
+        try:
+            # Extract the function name and args
+            action = parse_single_code_from_string(plan.split("Grounded Action")[-1])
+            function_name = re.match(r"(\w+\.\w+)\(", action).group(1)
+            args = self.parse_function_args(action)
+        except Exception as e:
+            raise RuntimeError(f"Error in parsing grounded action: {e}") from e
+
+        # arg0 is a description
+        if (
+            function_name in ["agent.click", "agent.type", "agent.scroll"]
+            and len(args) >= 1
+            and args[0] != None
+        ):
+            self.coords1 = self.generate_coords(args[0], obs)
+        # arg0 and arg1 are descriptions
+        elif function_name == "agent.drag_and_drop" and len(args) >= 2:
+            self.coords1 = self.generate_coords(args[0], obs)
+            self.coords2 = self.generate_coords(args[1], obs)
+        # arg0 and arg1 are text phrases
+        elif function_name == "agent.highlight_text_span" and len(args) >= 2:
+            self.coords1 = self.generate_text_coords(args[0], obs, alignment="start")
+            self.coords2 = self.generate_text_coords(args[1], obs, alignment="end")
+
+    # Resize from grounding model dim into OSWorld dim (1920 * 1080)
+    def resize_coordinates(self, coordinates: List[int]) -> List[int]:
+        # User explicitly passes the grounding model dimensions
+        if {"grounding_width", "grounding_height"}.issubset(
+            self.engine_params_for_grounding
+        ):
+            grounding_width = self.engine_params_for_grounding["grounding_width"]
+            grounding_height = self.engine_params_for_grounding["grounding_height"]
+        # Default to (1000, 1000), which is UI-TARS resizing
+        else:
+            grounding_width = 1000
+            grounding_height = 1000
+
+        return [
+            round(coordinates[0] * self.width / grounding_width),
+            round(coordinates[1] * self.height / grounding_height),
+        ]
+
+    # Given a generated ACI function, returns a list of argument values, where descriptions are at the front of the list
+    def parse_function_args(self, function: str) -> List[str]:
+        tree = ast.parse(function)
+        call_node = tree.body[0].value
+
+        def safe_eval(node):
+            if isinstance(
+                node, ast.Constant
+            ):  # Handles literals like numbers, strings, etc.
+                return node.value
+            else:
+                return ast.unparse(node)  # Return as a string if not a literal
+
+        positional_args = [safe_eval(arg) for arg in call_node.args]
+        keyword_args = {kw.arg: safe_eval(kw.value) for kw in call_node.keywords}
+
+        res = []
+
+        for key, val in keyword_args.items():
+            if "description" in key:
+                res.append(val)
+
+        for arg in positional_args:
+            res.append(arg)
+
+        return res
+
+    @agent_action
+    def click(
+        self,
+        element_description: str,
+        num_clicks: int = 1,
+        button_type: str = "left",
+        hold_keys: List = [],
+    ):
+        """Click on the element
+        Args:
+            element_description:str, a detailed descriptions of which element to click on. This description should be at least a full sentence.
+            num_clicks:int, number of times to click the element
+            button_type:str, which mouse button to press can be "left", "middle", or "right"
+            hold_keys:List, list of keys to hold while clicking
+        """
+        x, y = self.resize_coordinates(self.coords1)
+        command = "import pyautogui; "
+
+        # TODO: specified duration?
+        for k in hold_keys:
+            command += f"pyautogui.keyDown({repr(k)}); "
+        command += f"""import pyautogui; pyautogui.click({x}, {y}, clicks={num_clicks}, button={repr(button_type)}); """
+        for k in hold_keys:
+            command += f"pyautogui.keyUp({repr(k)}); "
+        # Return pyautoguicode to click on the element
+        return command
+
+    @agent_action
+    def switch_applications(self, app_code):
+        """Switch to a different application that is already open
+        Args:
+            app_code:str the code name of the application to switch to from the provided list of open applications
+        """
+        if self.platform == "darwin":
+            return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
+        elif self.platform == "linux":
+            return UBUNTU_APP_SETUP.replace("APP_NAME", app_code)
+        elif self.platform == "windows":
+            return f"import pyautogui; import time; pyautogui.hotkey('win', 'd', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
+
+    @agent_action
+    def open(self, app_or_filename: str):
+        """Open any application or file with name app_or_filename. Use this action to open applications or files on the desktop, do not open manually.
+        Args:
+            app_or_filename:str, the name of the application or filename to open
+        """
+        return f"import pyautogui; pyautogui.hotkey('win'); time.sleep(0.5); pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); pyautogui.hotkey('enter'); time.sleep(0.5)"
+
+    @agent_action
+    def type(
+        self,
+        element_description: Optional[str] = None,
+        text: str = "",
+        overwrite: bool = False,
+        enter: bool = False,
+    ):
+        """Type text into a specific element
+        Args:
+            element_description:str, a detailed description of which element to enter text in. This description should be at least a full sentence.
+            text:str, the text to type
+            overwrite:bool, Assign it to True if the text should overwrite the existing text, otherwise assign it to False. Using this argument clears all text in an element.
+            enter:bool, Assign it to True if the enter key should be pressed after typing the text, otherwise assign it to False.
+        """
+
+        if self.coords1 is not None:
+            # If a node is found, retrieve its coordinates and size
+            # Start typing at the center of the element
+
+            x, y = self.resize_coordinates(self.coords1)
+
+            command = "import pyautogui; "
+            command += f"pyautogui.click({x}, {y}); "
+
+            if overwrite:
+                command += (
+                    f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
+                )
+
+            command += f"pyautogui.write({repr(text)}); "
+
+            if enter:
+                command += "pyautogui.press('enter'); "
+        else:
+            # If no element is found, start typing at the current cursor location
+            command = "import pyautogui; "
+
+            if overwrite:
+                command += (
+                    f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
+                )
+
+            command += f"pyautogui.write({repr(text)}); "
+
+            if enter:
+                command += "pyautogui.press('enter'); "
+
+        return command
+
+    @agent_action
+    def save_to_knowledge(self, text: List[str]):
+        """Save facts, elements, texts, etc. to a long-term knowledge bank for reuse during this task. Can be used for copy-pasting text, saving elements, etc.
+        Args:
+            text:List[str] the text to save to the knowledge
+        """
+        self.notes.extend(text)
+        return """WAIT"""
+
+    @agent_action
+    def drag_and_drop(
+        self, starting_description: str, ending_description: str, hold_keys: List = []
+    ):
+        """Drag from the starting description to the ending description
+        Args:
+            starting_description:str, a very detailed description of where to start the drag action. This description should be at least a full sentence.
+            ending_description:str, a very detailed description of where to end the drag action. This description should be at least a full sentence.
+            hold_keys:List list of keys to hold while dragging
+        """
+        x1, y1 = self.resize_coordinates(self.coords1)
+        x2, y2 = self.resize_coordinates(self.coords2)
+
+        command = "import pyautogui; "
+
+        command += f"pyautogui.moveTo({x1}, {y1}); "
+        # TODO: specified duration?
+        for k in hold_keys:
+            command += f"pyautogui.keyDown({repr(k)}); "
+        command += f"pyautogui.dragTo({x2}, {y2}, duration=1.); pyautogui.mouseUp(); "
+        for k in hold_keys:
+            command += f"pyautogui.keyUp({repr(k)}); "
+
+        # Return pyautoguicode to drag and drop the elements
+
+        return command
+
+    @agent_action
+    def highlight_text_span(self, starting_phrase: str, ending_phrase: str):
+        """Highlight a text span between a provided starting phrase and ending phrase. Use this to highlight words, lines, and paragraphs.
+        Args:
+            starting_phrase:str, the phrase that denotes the start of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
+            ending_phrase:str, the phrase that denotes the end of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
+        """
+
+        x1, y1 = self.coords1
+        x2, y2 = self.coords2
+
+        command = "import pyautogui; "
+        command += f"pyautogui.moveTo({x1}, {y1}); "
+        command += f"pyautogui.dragTo({x2}, {y2}, duration=1.); pyautogui.mouseUp(); "
+
+        # Return pyautoguicode to drag and drop the elements
+        return command
+
+    @agent_action
+    def set_cell_values(
+        self, cell_values: Dict[str, Any], app_name: str, sheet_name: str
+    ):
+        """Use this to set individual cell values in a spreadsheet. For example, setting A2 to "hello" would be done by passing {"A2": "hello"} as cell_values. The sheet must be opened before this command can be used.
+        Args:
+            cell_values: Dict[str, Any], A dictionary of cell values to set in the spreadsheet. The keys are the cell coordinates in the format "A1", "B2", etc.
+                Supported value types include: float, int, string, bool, formulas.
+            app_name: str, The name of the spreadsheet application. For example, "Some_sheet.xlsx".
+            sheet_name: str, The name of the sheet in the spreadsheet. For example, "Sheet1".
+        """
+        return SET_CELL_VALUES_CMD.format(
+            cell_values=cell_values, app_name=app_name, sheet_name=sheet_name
+        )
+
+    @agent_action
+    def scroll(self, element_description: str, clicks: int, shift: bool = False):
+        """Scroll the element in the specified direction
+        Args:
+            element_description:str, a very detailed description of which element to enter scroll in. This description should be at least a full sentence.
+            clicks:int, the number of clicks to scroll can be positive (up) or negative (down).
+            shift:bool, whether to use shift+scroll for horizontal scrolling
+        """
+
+        x, y = self.resize_coordinates(self.coords1)
+
+        if shift:
+            return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.hscroll({clicks})"
+        else:
+            return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.vscroll({clicks})"
+
+    @agent_action
+    def hotkey(self, keys: List):
+        """Press a hotkey combination
+        Args:
+            keys:List the keys to press in combination in a list format (e.g. ['ctrl', 'c'])
+        """
+        # add quotes around the keys
+        keys = [f"'{key}'" for key in keys]
+        return f"import pyautogui; pyautogui.hotkey({', '.join(keys)})"
+
+    @agent_action
+    def hold_and_press(self, hold_keys: List, press_keys: List):
+        """Hold a list of keys and press a list of keys
+        Args:
+            hold_keys:List, list of keys to hold
+            press_keys:List, list of keys to press in a sequence
+        """
+
+        press_keys_str = "[" + ", ".join([f"'{key}'" for key in press_keys]) + "]"
+        command = "import pyautogui; "
+        for k in hold_keys:
+            command += f"pyautogui.keyDown({repr(k)}); "
+        command += f"pyautogui.press({press_keys_str}); "
+        for k in hold_keys:
+            command += f"pyautogui.keyUp({repr(k)}); "
+
+        return command
+
+    @agent_action
+    def wait(self, time: float):
+        """Wait for a specified amount of time
+        Args:
+            time:float the amount of time to wait in seconds
+        """
+        return f"""import time; time.sleep({time})"""
+
+    @agent_action
+    def done(
+        self,
+        return_value: Optional[Union[Dict, str, List, Tuple, int, float, bool]] = None,
+    ):
+        """End the current task with a success and the required return value"""
+        self.returned_info = return_value
+        return """DONE"""
+
+    @agent_action
+    def fail(self):
+        """End the current task with a failure, and replan the whole task."""
+        return """FAIL"""

gui_agents/s2/agents/manager.py

@@ -0,0 +1,321 @@
+import logging
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+import platform
+
+from gui_agents.s2.agents.grounding import ACI
+from gui_agents.s2.core.module import BaseModule
+from gui_agents.s2.core.knowledge import KnowledgeBase
+from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
+from gui_agents.s2.core.engine import OpenAIEmbeddingEngine
+from gui_agents.s2.utils.common_utils import (
+    Dag,
+    Node,
+    calculate_tokens,
+    call_llm_safe,
+    parse_dag,
+)
+
+logger = logging.getLogger("desktopenv.agent")
+
+NUM_IMAGE_TOKEN = 1105  # Value set of screen of size 1920x1080 for openai vision
+
+
+class Manager(BaseModule):
+    def __init__(
+        self,
+        engine_params: Dict,
+        grounding_agent: ACI,
+        local_kb_path: str,
+        embedding_engine,
+        search_engine: Optional[str] = None,
+        multi_round: bool = False,
+        platform: str = platform.system().lower(),
+    ):
+        # TODO: move the prompt to Procedural Memory
+        super().__init__(engine_params, platform)
+
+        # Initialize the ACI
+        self.grounding_agent = grounding_agent
+
+        # Initialize the planner
+        sys_prompt = PROCEDURAL_MEMORY.COMBINED_MANAGER_PROMPT
+
+        self.generator_agent = self._create_agent(sys_prompt)
+
+        # Initialize the remaining modules
+        self.dag_translator_agent = self._create_agent(
+            PROCEDURAL_MEMORY.DAG_TRANSLATOR_PROMPT
+        )
+        self.narrative_summarization_agent = self._create_agent(
+            PROCEDURAL_MEMORY.TASK_SUMMARIZATION_PROMPT
+        )
+        self.episode_summarization_agent = self._create_agent(
+            PROCEDURAL_MEMORY.SUBTASK_SUMMARIZATION_PROMPT
+        )
+
+        self.local_kb_path = local_kb_path
+
+        self.embedding_engine = embedding_engine
+        self.knowledge_base = KnowledgeBase(
+            embedding_engine=self.embedding_engine,
+            local_kb_path=self.local_kb_path,
+            platform=platform,
+            engine_params=engine_params,
+        )
+
+        self.planner_history = []
+
+        self.turn_count = 0
+        self.search_engine = search_engine
+        self.multi_round = multi_round
+
+    def summarize_episode(self, trajectory):
+        """Summarize the episode experience for lifelong learning reflection
+        Args:
+            trajectory: str: The episode experience to be summarized
+        """
+
+        # Create Reflection on whole trajectories for next round trial, keep earlier messages as exemplars
+        self.episode_summarization_agent.add_message(trajectory, role="user")
+        subtask_summarization = call_llm_safe(self.episode_summarization_agent)
+        self.episode_summarization_agent.add_message(
+            subtask_summarization, role="assistant"
+        )
+
+        return subtask_summarization
+
+    def summarize_narrative(self, trajectory):
+        """Summarize the narrative experience for lifelong learning reflection
+        Args:
+            trajectory: str: The narrative experience to be summarized
+        """
+        # Create Reflection on whole trajectories for next round trial
+        self.narrative_summarization_agent.add_message(trajectory, role="user")
+        lifelong_learning_reflection = call_llm_safe(self.narrative_summarization_agent)
+
+        return lifelong_learning_reflection
+
+    def _generate_step_by_step_plan(
+        self,
+        observation: Dict,
+        instruction: str,
+        failed_subtask: Optional[Node] = None,
+        completed_subtasks_list: List[Node] = [],
+        remaining_subtasks_list: List[Node] = [],
+    ) -> Tuple[Dict, str]:
+        agent = self.grounding_agent
+
+        # Converts a list of DAG Nodes into a natural langauge list
+        def format_subtask_list(subtasks: List[Node]) -> str:
+            res = ""
+            for idx, node in enumerate(subtasks):
+                res += f"{idx+1}. **{node.name}**:\n"
+                bullets = re.split(r"(?<=[.!?;]) +", node.info)
+                for bullet in bullets:
+                    res += f"   - {bullet}\n"
+                res += "\n"
+            return res
+
+        # Perform Retrieval only at the first planning step
+        if self.turn_count == 0:
+
+            self.search_query = self.knowledge_base.formulate_query(
+                instruction, observation
+            )
+
+            most_similar_task = ""
+            retrieved_experience = ""
+            integrated_knowledge = ""
+            # Retrieve most similar narrative (task) experience
+            most_similar_task, retrieved_experience = (
+                self.knowledge_base.retrieve_narrative_experience(instruction)
+            )
+            logger.info(
+                "SIMILAR TASK EXPERIENCE: %s",
+                most_similar_task + "\n" + retrieved_experience.strip(),
+            )
+
+            # Retrieve knowledge from the web if search_engine is provided
+            if self.search_engine is not None:
+                retrieved_knowledge = self.knowledge_base.retrieve_knowledge(
+                    instruction=instruction,
+                    search_query=self.search_query,
+                    search_engine=self.search_engine,
+                )
+                logger.info("RETRIEVED KNOWLEDGE: %s", retrieved_knowledge)
+
+                if retrieved_knowledge is not None:
+                    # Fuse the retrieved knowledge and experience
+                    integrated_knowledge = self.knowledge_base.knowledge_fusion(
+                        observation=observation,
+                        instruction=instruction,
+                        web_knowledge=retrieved_knowledge,
+                        similar_task=most_similar_task,
+                        experience=retrieved_experience,
+                    )
+                    logger.info("INTEGRATED KNOWLEDGE: %s", integrated_knowledge)
+
+            integrated_knowledge = integrated_knowledge or retrieved_experience
+
+            # Add the integrated knowledge to the task instruction in the system prompt
+            if integrated_knowledge:
+                instruction += f"\nYou may refer to some retrieved knowledge if you think they are useful.{integrated_knowledge}"
+
+            self.generator_agent.add_system_prompt(
+                self.generator_agent.system_prompt.replace(
+                    "TASK_DESCRIPTION", instruction
+                )
+            )
+
+        # Re-plan on failure case
+        if failed_subtask:
+            generator_message = (
+                f"The subtask {failed_subtask} cannot be completed. Please generate a new plan for the remainder of the trajectory.\n\n"
+                f"Successfully Completed Subtasks:\n{format_subtask_list(completed_subtasks_list)}\n"
+            )
+        # Re-plan on subtask completion case
+        elif len(completed_subtasks_list) + len(remaining_subtasks_list) > 0:
+            generator_message = (
+                "The current trajectory and desktop state is provided. Please revise the plan for the following trajectory.\n\n"
+                f"Successfully Completed Subtasks:\n{format_subtask_list(completed_subtasks_list)}\n"
+                f"Future Remaining Subtasks:\n{format_subtask_list(remaining_subtasks_list)}\n"
+            )
+        # Initial plan case
+        else:
+            generator_message = "Please generate the initial plan for the task.\n"
+
+        logger.info("GENERATOR MESSAGE: %s", generator_message)
+
+        self.generator_agent.add_message(
+            generator_message,
+            image_content=observation.get("screenshot", None),
+            role="user",
+        )
+
+        logger.info("GENERATING HIGH LEVEL PLAN")
+
+        plan = call_llm_safe(self.generator_agent)
+        if plan == "":
+            raise Exception("Plan Generation Failed - Fix the Prompt")
+
+        logger.info("HIGH LEVEL STEP BY STEP PLAN: %s", plan)
+
+        self.generator_agent.add_message(plan, role="assistant")
+        self.planner_history.append(plan)
+        self.turn_count += 1
+
+        # Set Cost based on GPT-4o
+        input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
+        cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
+
+        planner_info = {
+            "search_query": self.search_query,
+            "goal_plan": plan,
+            "num_input_tokens_plan": input_tokens,
+            "num_output_tokens_plan": output_tokens,
+            "goal_plan_cost": cost,
+        }
+
+        assert type(plan) == str
+
+        return planner_info, plan
+
+    def _generate_dag(self, instruction: str, plan: str) -> Tuple[Dict, Dag]:
+        # For the re-planning case, remove the prior input since this should only translate the new plan
+        self.dag_translator_agent.reset()
+
+        # Add initial instruction and plan to the agent's message history
+        self.dag_translator_agent.add_message(
+            f"Instruction: {instruction}\nPlan: {plan}", role="user"
+        )
+
+        logger.info("GENERATING DAG")
+
+        # Generate DAG
+        dag_raw = call_llm_safe(self.dag_translator_agent)
+
+        dag = parse_dag(dag_raw)
+
+        logger.info("Generated DAG: %s", dag_raw)
+
+        self.dag_translator_agent.add_message(dag_raw, role="assistant")
+
+        input_tokens, output_tokens = calculate_tokens(
+            self.dag_translator_agent.messages
+        )
+
+        # Set Cost based on GPT-4o
+        cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
+
+        dag_info = {
+            "dag": dag_raw,
+            "num_input_tokens_dag": input_tokens,
+            "num_output_tokens_dag": output_tokens,
+            "dag_cost": cost,
+        }
+
+        assert type(dag) == Dag
+
+        return dag_info, dag
+
+    def _topological_sort(self, dag: Dag) -> List[Node]:
+        """Topological sort of the DAG using DFS
+        dag: Dag: Object representation of the DAG with nodes and edges
+        """
+
+        def dfs(node_name, visited, stack):
+            visited[node_name] = True
+            for neighbor in adj_list[node_name]:
+                if not visited[neighbor]:
+                    dfs(neighbor, visited, stack)
+            stack.append(node_name)
+
+        # Convert edges to adjacency list
+        adj_list = defaultdict(list)
+        for u, v in dag.edges:
+            adj_list[u.name].append(v.name)
+
+        visited = {node.name: False for node in dag.nodes}
+        stack = []
+
+        for node in dag.nodes:
+            if not visited[node.name]:
+                dfs(node.name, visited, stack)
+
+        # Return the nodes in topologically sorted order
+        sorted_nodes = [
+            next(n for n in dag.nodes if n.name == name) for name in stack[::-1]
+        ]
+        return sorted_nodes
+
+    def get_action_queue(
+        self,
+        instruction: str,
+        observation: Dict,
+        failed_subtask: Optional[Node] = None,
+        completed_subtasks_list: List[Node] = [],
+        remaining_subtasks_list: List[Node] = [],
+    ):
+        """Generate the action list based on the instruction
+        instruction:str: Instruction for the task
+        """
+
+        planner_info, plan = self._generate_step_by_step_plan(
+            observation,
+            instruction,
+            failed_subtask,
+            completed_subtasks_list,
+            remaining_subtasks_list,
+        )
+
+        # Generate the DAG
+        dag_info, dag = self._generate_dag(instruction, plan)
+
+        # Topological sort of the DAG
+        action_queue = self._topological_sort(dag)
+
+        planner_info.update(dag_info)
+
+        return planner_info, action_queue

gui_agents/s2/agents/worker.py

@@ -0,0 +1,254 @@
+import logging
+import re
+import textwrap
+from typing import Dict, List, Tuple
+import platform
+
+from gui_agents.s2.agents.grounding import ACI
+from gui_agents.s2.core.module import BaseModule
+from gui_agents.s2.core.knowledge import KnowledgeBase
+from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
+from gui_agents.s2.utils.common_utils import (
+    Node,
+    calculate_tokens,
+    call_llm_safe,
+    parse_single_code_from_string,
+    sanitize_code,
+    extract_first_agent_function,
+)
+
+logger = logging.getLogger("desktopenv.agent")
+
+
+class Worker(BaseModule):
+    def __init__(
+        self,
+        engine_params: Dict,
+        grounding_agent: ACI,
+        local_kb_path: str,
+        embedding_engine,
+        platform: str = platform.system().lower(),
+        enable_reflection: bool = True,
+        use_subtask_experience: bool = True,
+    ):
+        """
+        Worker receives a subtask list and active subtask and generates the next action for the to execute.
+        Args:
+            engine_params: Dict
+                Parameters for the multimodal engine
+            grounding_agent: Agent
+                The grounding agent to use
+            local_kb_path: str
+                Path to knowledge base
+            platform: str
+                OS platform the agent runs on (darwin, linux, windows)
+            enable_reflection: bool
+                Whether to enable reflection
+            use_subtask_experience: bool
+                Whether to use subtask experience
+        """
+        super().__init__(engine_params, platform)
+
+        self.grounding_agent = grounding_agent
+        self.local_kb_path = local_kb_path
+        self.embedding_engine = embedding_engine
+        self.enable_reflection = enable_reflection
+        self.use_subtask_experience = use_subtask_experience
+        self.reset()
+
+    def reset(self):
+        if self.platform != "linux":
+            skipped_actions = ["set_cell_values"]
+        else:
+            skipped_actions = []
+
+        sys_prompt = PROCEDURAL_MEMORY.construct_worker_procedural_memory(
+            type(self.grounding_agent), skipped_actions=skipped_actions
+        ).replace("CURRENT_OS", self.platform)
+
+        self.generator_agent = self._create_agent(sys_prompt)
+        self.reflection_agent = self._create_agent(
+            PROCEDURAL_MEMORY.REFLECTION_ON_TRAJECTORY
+        )
+
+        self.knowledge_base = KnowledgeBase(
+            embedding_engine=self.embedding_engine,
+            local_kb_path=self.local_kb_path,
+            platform=self.platform,
+            engine_params=self.engine_params,
+        )
+
+        self.turn_count = 0
+        self.worker_history = []
+        self.reflections = []
+        self.cost_this_turn = 0
+        self.screenshot_inputs = []
+        self.planner_history = []
+        self.max_trajector_length = 8
+
+    def flush_messages(self):
+        # generator msgs are alternating [user, assistant], so 2 per round
+        if len(self.generator_agent.messages) > 2 * self.max_trajector_length + 1:
+            self.generator_agent.remove_message_at(1)
+            self.generator_agent.remove_message_at(1)
+        # reflector msgs are all [(user text, user image)], so 1 per round
+        if len(self.reflection_agent.messages) > self.max_trajector_length + 1:
+            self.reflection_agent.remove_message_at(1)
+
+    def generate_next_action(
+        self,
+        instruction: str,
+        search_query: str,
+        subtask: str,
+        subtask_info: Dict,
+        future_tasks: List[Node],
+        done_task: List[Node],
+        obs: Dict,
+    ) -> Tuple[Dict, List]:
+        """
+        Predict the next action(s) based on the current observation.
+        """
+        # Provide the top_app to the Grounding Agent to remove all other applications from the tree. At t=0, top_app is None
+        agent = self.grounding_agent
+
+        # Get RAG knowledge, only update system message at t=0
+        if self.turn_count == 0:
+            if self.use_subtask_experience:
+                subtask_query_key = (
+                    "Task:\n"
+                    + search_query
+                    + "\n\nSubtask: "
+                    + subtask
+                    + "\nSubtask Instruction: "
+                    + subtask_info
+                )
+                retrieved_similar_subtask, retrieved_subtask_experience = (
+                    self.knowledge_base.retrieve_episodic_experience(subtask_query_key)
+                )
+
+                # Dirty fix to replace id with element description during subtask retrieval
+                pattern = r"\(\d+"
+                retrieved_subtask_experience = re.sub(
+                    pattern, "(element_description", retrieved_subtask_experience
+                )
+                retrieved_subtask_experience = retrieved_subtask_experience.replace(
+                    "_id", "_description"
+                )
+
+                logger.info(
+                    "SIMILAR SUBTASK EXPERIENCE: %s",
+                    retrieved_similar_subtask
+                    + "\n"
+                    + retrieved_subtask_experience.strip(),
+                )
+                instruction += "\nYou may refer to some similar subtask experience if you think they are useful. {}".format(
+                    retrieved_similar_subtask + "\n" + retrieved_subtask_experience
+                )
+
+            self.generator_agent.add_system_prompt(
+                self.generator_agent.system_prompt.replace(
+                    "SUBTASK_DESCRIPTION", subtask
+                )
+                .replace("TASK_DESCRIPTION", instruction)
+                .replace("FUTURE_TASKS", ", ".join([f.name for f in future_tasks]))
+                .replace("DONE_TASKS", ",".join(d.name for d in done_task))
+            )
+
+        # Reflection generation does not add its own response, it only gets the trajectory
+        reflection = None
+        if self.enable_reflection:
+            # Load the initial subtask info
+            if self.turn_count == 0:
+                text_content = textwrap.dedent(f"""
+                    Subtask Description: {subtask}
+                    Subtask Information: {subtask_info}
+                    Current Trajectory below:
+                    """)
+                updated_sys_prompt = (
+                    self.reflection_agent.system_prompt + "\n" + text_content
+                )
+                self.reflection_agent.add_system_prompt(updated_sys_prompt)
+                self.reflection_agent.add_message(
+                    text_content="The initial screen is provided. No action has been taken yet.",
+                    image_content=obs["screenshot"],
+                    role="user",
+                )
+            # Load the latest action
+            else:
+                text_content = self.clean_worker_generation_for_reflection(
+                    self.planner_history[-1]
+                )
+                self.reflection_agent.add_message(
+                    text_content=text_content,
+                    image_content=obs["screenshot"],
+                    role="user",
+                )
+                reflection = call_llm_safe(self.reflection_agent)
+                self.reflections.append(reflection)
+                logger.info("REFLECTION: %s", reflection)
+
+        generator_message = (
+            f"\nYou may use this reflection on the previous action and overall trajectory: {reflection}\n"
+            if reflection and self.turn_count > 0
+            else ""
+        ) + f"Text Buffer = [{','.join(agent.notes)}]."
+
+        # Only provide subinfo in the very first message to avoid over influence and redundancy
+        if self.turn_count == 0:
+            generator_message += f"Remember only complete the subtask: {subtask}\n"
+            generator_message += f"You can use this extra information for completing the current subtask: {subtask_info}.\n"
+
+        # logger.info("GENERATOR MESSAGE: %s", generator_message)
+
+        self.generator_agent.add_message(
+            generator_message, image_content=obs["screenshot"], role="user"
+        )
+
+        plan = call_llm_safe(self.generator_agent)
+        self.planner_history.append(plan)
+        logger.info("PLAN: %s", plan)
+        self.generator_agent.add_message(plan, role="assistant")
+
+        # Calculate input/output tokens and gpt-4o cost
+        input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
+        cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
+        self.cost_this_turn += cost
+        logger.info("EXECTUOR COST: %s", self.cost_this_turn)
+
+        # Use the DescriptionBasedACI to convert agent_action("desc") into agent_action([x, y])
+        try:
+            agent.assign_coordinates(plan, obs)
+            plan_code = parse_single_code_from_string(plan.split("Grounded Action")[-1])
+            plan_code = sanitize_code(plan_code)
+            plan_code = extract_first_agent_function(plan_code)
+            exec_code = eval(plan_code)
+        except Exception as e:
+            logger.error("Error in parsing plan code: %s", e)
+            plan_code = "agent.wait(1.0)"
+            exec_code = eval(plan_code)
+
+        executor_info = {
+            "current_subtask": subtask,
+            "current_subtask_info": subtask_info,
+            "executor_plan": plan,
+            "plan_code": plan_code,
+            "reflection": reflection,
+            "num_input_tokens_executor": input_tokens,
+            "num_output_tokens_executor": output_tokens,
+        }
+        self.turn_count += 1
+
+        self.screenshot_inputs.append(obs["screenshot"])
+        self.flush_messages()
+
+        return executor_info, [exec_code]
+
+    # Removes the previous action verification, and removes any extraneous grounded actions
+    def clean_worker_generation_for_reflection(self, worker_generation: str) -> str:
+        # Remove the previous action verification
+        res = worker_generation[worker_generation.find("(Screenshot Analysis)") :]
+        action = extract_first_agent_function(worker_generation)
+        # Cut off extra grounded actions
+        res = res[: res.find("(Grounded Action)")]
+        res += f"(Grounded Action)\n```python\n{action}\n```\n"
+        return res

gui_agents/s2/cli_app.py

@@ -0,0 +1,382 @@
+import argparse
+import datetime
+import io
+import logging
+import os
+import platform
+import pyautogui
+import signal
+import sys
+import time
+
+from PIL import Image
+
+from gui_agents.s2.agents.grounding import OSWorldACI
+from gui_agents.s2.agents.agent_s import AgentS2
+
+current_platform = platform.system().lower()
+
+# Global flag to track pause state for debugging
+paused = False
+
+
+def get_char():
+    """Get a single character from stdin without pressing Enter"""
+    try:
+        # Import termios and tty on Unix-like systems
+        if platform.system() in ["Darwin", "Linux"]:
+            import termios
+            import tty
+
+            fd = sys.stdin.fileno()
+            old_settings = termios.tcgetattr(fd)
+            try:
+                tty.setraw(sys.stdin.fileno())
+                ch = sys.stdin.read(1)
+            finally:
+                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+            return ch
+        else:
+            # Windows fallback
+            import msvcrt
+
+            return msvcrt.getch().decode("utf-8", errors="ignore")
+    except:
+        return input()  # Fallback for non-terminal environments
+
+
+def signal_handler(signum, frame):
+    """Handle Ctrl+C signal for debugging during agent execution"""
+    global paused
+
+    if not paused:
+        print("\n\n🔸 Agent-S Workflow Paused 🔸")
+        print("=" * 50)
+        print("Options:")
+        print("  • Press Ctrl+C again to quit")
+        print("  • Press Esc to resume workflow")
+        print("=" * 50)
+
+        paused = True
+
+        while paused:
+            try:
+                print("\n[PAUSED] Waiting for input... ", end="", flush=True)
+                char = get_char()
+
+                if ord(char) == 3:  # Ctrl+C
+                    print("\n\n🛑 Exiting Agent-S...")
+                    sys.exit(0)
+                elif ord(char) == 27:  # Esc
+                    print("\n\n▶️  Resuming Agent-S workflow...")
+                    paused = False
+                    break
+                else:
+                    print(f"\n   Unknown command: '{char}' (ord: {ord(char)})")
+
+            except KeyboardInterrupt:
+                print("\n\n🛑 Exiting Agent-S...")
+                sys.exit(0)
+    else:
+        # Already paused, second Ctrl+C means quit
+        print("\n\n🛑 Exiting Agent-S...")
+        sys.exit(0)
+
+
+# Set up signal handler for Ctrl+C
+signal.signal(signal.SIGINT, signal_handler)
+
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+
+datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
+
+log_dir = "logs"
+os.makedirs(log_dir, exist_ok=True)
+
+file_handler = logging.FileHandler(
+    os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+debug_handler = logging.FileHandler(
+    os.path.join("logs", "debug-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+stdout_handler = logging.StreamHandler(sys.stdout)
+sdebug_handler = logging.FileHandler(
+    os.path.join("logs", "sdebug-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+
+file_handler.setLevel(logging.INFO)
+debug_handler.setLevel(logging.DEBUG)
+stdout_handler.setLevel(logging.INFO)
+sdebug_handler.setLevel(logging.DEBUG)
+
+formatter = logging.Formatter(
+    fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
+)
+file_handler.setFormatter(formatter)
+debug_handler.setFormatter(formatter)
+stdout_handler.setFormatter(formatter)
+sdebug_handler.setFormatter(formatter)
+
+stdout_handler.addFilter(logging.Filter("desktopenv"))
+sdebug_handler.addFilter(logging.Filter("desktopenv"))
+
+logger.addHandler(file_handler)
+logger.addHandler(debug_handler)
+logger.addHandler(stdout_handler)
+logger.addHandler(sdebug_handler)
+
+platform_os = platform.system()
+
+
+def show_permission_dialog(code: str, action_description: str):
+    """Show a platform-specific permission dialog and return True if approved."""
+    if platform.system() == "Darwin":
+        result = os.system(
+            f'osascript -e \'display dialog "Do you want to execute this action?\n\n{code} which will try to {action_description}" with title "Action Permission" buttons {{"Cancel", "OK"}} default button "OK" cancel button "Cancel"\''
+        )
+        return result == 0
+    elif platform.system() == "Linux":
+        result = os.system(
+            f'zenity --question --title="Action Permission" --text="Do you want to execute this action?\n\n{code}" --width=400 --height=200'
+        )
+        return result == 0
+    return False
+
+
+def scale_screen_dimensions(width: int, height: int, max_dim_size: int):
+    scale_factor = min(max_dim_size / width, max_dim_size / height, 1)
+    safe_width = int(width * scale_factor)
+    safe_height = int(height * scale_factor)
+    return safe_width, safe_height
+
+
+def run_agent(agent, instruction: str, scaled_width: int, scaled_height: int):
+    global paused
+    obs = {}
+    traj = "Task:\n" + instruction
+    subtask_traj = ""
+    for step in range(15):
+        # Check if we're in paused state and wait
+        while paused:
+            time.sleep(0.1)
+
+        # Get screen shot using pyautogui
+        screenshot = pyautogui.screenshot()
+        screenshot = screenshot.resize((scaled_width, scaled_height), Image.LANCZOS)
+
+        # Save the screenshot to a BytesIO object
+        buffered = io.BytesIO()
+        screenshot.save(buffered, format="PNG")
+
+        # Get the byte value of the screenshot
+        screenshot_bytes = buffered.getvalue()
+        # Convert to base64 string.
+        obs["screenshot"] = screenshot_bytes
+
+        # Check again for pause state before prediction
+        while paused:
+            time.sleep(0.1)
+
+        print(f"\n🔄 Step {step + 1}/15: Getting next action from agent...")
+
+        # Get next action code from the agent
+        info, code = agent.predict(instruction=instruction, observation=obs)
+
+        if "done" in code[0].lower() or "fail" in code[0].lower():
+            if platform.system() == "Darwin":
+                os.system(
+                    f'osascript -e \'display dialog "Task Completed" with title "OpenACI Agent" buttons "OK" default button "OK"\''
+                )
+            elif platform.system() == "Linux":
+                os.system(
+                    f'zenity --info --title="OpenACI Agent" --text="Task Completed" --width=200 --height=100'
+                )
+
+            agent.update_narrative_memory(traj)
+            break
+
+        if "next" in code[0].lower():
+            continue
+
+        if "wait" in code[0].lower():
+            print("⏳ Agent requested wait...")
+            time.sleep(5)
+            continue
+
+        else:
+            time.sleep(1.0)
+            print("EXECUTING CODE:", code[0])
+
+            # Check for pause state before execution
+            while paused:
+                time.sleep(0.1)
+
+            # Ask for permission before executing
+            exec(code[0])
+            time.sleep(1.0)
+
+            # Update task and subtask trajectories and optionally the episodic memory
+            traj += (
+                "\n\nReflection:\n"
+                + str(info["reflection"])
+                + "\n\n----------------------\n\nPlan:\n"
+                + info["executor_plan"]
+            )
+            subtask_traj = agent.update_episodic_memory(info, subtask_traj)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run AgentS2 with specified model.")
+    parser.add_argument(
+        "--provider",
+        type=str,
+        default="anthropic",
+        help="Specify the provider to use (e.g., openai, anthropic, etc.)",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="claude-3-7-sonnet-20250219",
+        help="Specify the model to use (e.g., gpt-4o)",
+    )
+    parser.add_argument(
+        "--model_url",
+        type=str,
+        default="",
+        help="The URL of the main generation model API.",
+    )
+    parser.add_argument(
+        "--model_api_key",
+        type=str,
+        default="",
+        help="The API key of the main generation model.",
+    )
+
+    # Grounding model config option 1: API based
+    parser.add_argument(
+        "--grounding_model_provider",
+        type=str,
+        default="anthropic",
+        help="Specify the provider to use for the grounding model (e.g., openai, anthropic, etc.)",
+    )
+    parser.add_argument(
+        "--grounding_model",
+        type=str,
+        default="claude-3-7-sonnet-20250219",
+        help="Specify the grounding model to use (e.g., claude-3-5-sonnet-20241022)",
+    )
+    parser.add_argument(
+        "--grounding_model_resize_width",
+        type=int,
+        default=1366,
+        help="Width of screenshot image after processor rescaling",
+    )
+    parser.add_argument(
+        "--grounding_model_resize_height",
+        type=int,
+        default=None,
+        help="Height of screenshot image after processor rescaling",
+    )
+
+    # Grounding model config option 2: Self-hosted endpoint based
+    parser.add_argument(
+        "--endpoint_provider",
+        type=str,
+        default="",
+        help="Specify the endpoint provider for your grounding model, only HuggingFace TGI support for now",
+    )
+    parser.add_argument(
+        "--endpoint_url",
+        type=str,
+        default="",
+        help="Specify the endpoint URL for your grounding model",
+    )
+    parser.add_argument(
+        "--endpoint_api_key",
+        type=str,
+        default="",
+        help="The API key of the grounding model.",
+    )
+
+    parser.add_argument(
+        "--embedding_engine_type",
+        type=str,
+        default="openai",
+        help="Specify the embedding engine type (supports openai, gemini)",
+    )
+
+    args = parser.parse_args()
+    assert (
+        args.grounding_model_provider and args.grounding_model
+    ) or args.endpoint_url, "Error: No grounding model was provided. Either provide an API based model, or a self-hosted HuggingFace endpoint"
+
+    # Re-scales screenshot size to ensure it fits in UI-TARS context limit
+    screen_width, screen_height = pyautogui.size()
+    scaled_width, scaled_height = scale_screen_dimensions(
+        screen_width, screen_height, max_dim_size=2400
+    )
+
+    # Load the general engine params
+    engine_params = {
+        "engine_type": args.provider,
+        "model": args.model,
+        "base_url": args.model_url,
+        "api_key": args.model_api_key,
+    }
+
+    # Load the grounding engine from a HuggingFace TGI endpoint
+    if args.endpoint_url:
+        engine_params_for_grounding = {
+            "engine_type": args.endpoint_provider,
+            "base_url": args.endpoint_url,
+            "api_key": args.endpoint_api_key,
+        }
+    else:
+        grounding_height = args.grounding_model_resize_height
+        # If not provided, use the aspect ratio of the screen to compute the height
+        if grounding_height is None:
+            grounding_height = (
+                screen_height * args.grounding_model_resize_width / screen_width
+            )
+
+        engine_params_for_grounding = {
+            "engine_type": args.grounding_model_provider,
+            "model": args.grounding_model,
+            "grounding_width": args.grounding_model_resize_width,
+            "grounding_height": grounding_height,
+        }
+
+    grounding_agent = OSWorldACI(
+        platform=current_platform,
+        engine_params_for_generation=engine_params,
+        engine_params_for_grounding=engine_params_for_grounding,
+        width=screen_width,
+        height=screen_height,
+    )
+
+    agent = AgentS2(
+        engine_params,
+        grounding_agent,
+        platform=current_platform,
+        action_space="pyautogui",
+        observation_type="mixed",
+        search_engine=None,
+        embedding_engine_type=args.embedding_engine_type,
+    )
+
+    while True:
+        query = input("Query: ")
+
+        agent.reset()
+
+        # Run the agent on your own device
+        run_agent(agent, query, scaled_width, scaled_height)
+
+        response = input("Would you like to provide another query? (y/n): ")
+        if response.lower() != "y":
+            break
+
+
+if __name__ == "__main__":
+    main()

gui_agents/s2/core/engine.py

@@ -0,0 +1,485 @@
+import os
+
+import backoff
+import numpy as np
+from anthropic import Anthropic
+from openai import (
+    AzureOpenAI,
+    APIConnectionError,
+    APIError,
+    AzureOpenAI,
+    OpenAI,
+    RateLimitError,
+)
+from google import genai
+from google.genai import types
+
+
+class LMMEngine:
+    pass
+
+
+class OpenAIEmbeddingEngine(LMMEngine):
+    def __init__(
+        self,
+        embedding_model: str = "text-embedding-3-small",
+        api_key=None,
+    ):
+        """Init an OpenAI Embedding engine
+
+        Args:
+            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
+            api_key (_type_, optional): Auth key from OpenAI. Defaults to None.
+        """
+        self.model = embedding_model
+        self.api_key = api_key
+
+    @backoff.on_exception(
+        backoff.expo,
+        (
+            APIError,
+            RateLimitError,
+            APIConnectionError,
+        ),
+    )
+    def get_embeddings(self, text: str) -> np.ndarray:
+        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
+            )
+        client = OpenAI(api_key=api_key)
+        response = client.embeddings.create(model=self.model, input=text)
+        return np.array([data.embedding for data in response.data])
+
+
+class GeminiEmbeddingEngine(LMMEngine):
+    def __init__(
+        self,
+        embedding_model: str = "text-embedding-004",
+        api_key=None,
+    ):
+        """Init an Gemini Embedding engine
+
+        Args:
+            embedding_model (str, optional): Model name. Defaults to "text-embedding-004".
+            api_key (_type_, optional): Auth key from Gemini. Defaults to None.
+        """
+        self.model = embedding_model
+        self.api_key = api_key
+
+    @backoff.on_exception(
+        backoff.expo,
+        (
+            APIError,
+            RateLimitError,
+            APIConnectionError,
+        ),
+    )
+    def get_embeddings(self, text: str) -> np.ndarray:
+        api_key = self.api_key or os.getenv("GEMINI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
+            )
+        client = genai.Client(api_key=api_key)
+
+        result = client.models.embed_content(
+            model=self.model,
+            contents=text,
+            config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY"),
+        )
+
+        return np.array([i.values for i in result.embeddings])
+
+
+class AzureOpenAIEmbeddingEngine(LMMEngine):
+    def __init__(
+        self,
+        embedding_model: str = "text-embedding-3-small",
+        api_key=None,
+        api_version=None,
+        endpoint_url=None,
+    ):
+        """Init an Azure OpenAI Embedding engine
+
+        Args:
+            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
+            api_key (_type_, optional): Auth key from Azure OpenAI. Defaults to None.
+            api_version (_type_, optional): API version. Defaults to None.
+            endpoint_url (_type_, optional): Endpoint URL. Defaults to None.
+        """
+        self.model = embedding_model
+        self.api_key = api_key
+        self.api_version = api_version
+        self.endpoint_url = endpoint_url
+
+    @backoff.on_exception(
+        backoff.expo,
+        (
+            APIError,
+            RateLimitError,
+            APIConnectionError,
+        ),
+    )
+    def get_embeddings(self, text: str) -> np.ndarray:
+        api_key = self.api_key or os.getenv("AZURE_OPENAI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
+            )
+        api_version = self.api_version or os.getenv("OPENAI_API_VERSION")
+        if api_version is None:
+            raise ValueError(
+                "An API Version needs to be provided in either the api_version parameter or as an environment variable named OPENAI_API_VERSION"
+            )
+        endpoint_url = self.endpoint_url or os.getenv("AZURE_OPENAI_ENDPOINT")
+        if endpoint_url is None:
+            raise ValueError(
+                "An Endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
+            )
+        client = AzureOpenAI(
+            api_key=api_key,
+            api_version=api_version,
+            azure_endpoint=endpoint_url,
+        )
+        response = client.embeddings.create(input=text, model=self.model)
+        return np.array([data.embedding for data in response.data])
+
+
+class LMMEngineOpenAI(LMMEngine):
+    def __init__(
+        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.base_url = base_url
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
+            )
+        if not self.llm_client:
+            if not self.base_url:
+                self.llm_client = OpenAI(api_key=api_key)
+            else:
+                self.llm_client = OpenAI(base_url=self.base_url, api_key=api_key)
+        return (
+            self.llm_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .choices[0]
+            .message.content
+        )
+
+
+class LMMEngineAnthropic(LMMEngine):
+    def __init__(
+        self, base_url=None, api_key=None, model=None, thinking=False, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.thinking = thinking
+        self.api_key = api_key
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("ANTHROPIC_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ANTHROPIC_API_KEY"
+            )
+        if not self.llm_client:
+            self.llm_client = Anthropic(api_key=api_key)
+        if self.thinking:
+            full_response = self.llm_client.messages.create(
+                system=messages[0]["content"][0]["text"],
+                model=self.model,
+                messages=messages[1:],
+                max_tokens=8192,
+                thinking={"type": "enabled", "budget_tokens": 4096},
+                **kwargs,
+            )
+            thoughts = full_response.content[0].thinking
+            print("CLAUDE 3.7 THOUGHTS:", thoughts)
+            return full_response.content[1].text
+        return (
+            self.llm_client.messages.create(
+                system=messages[0]["content"][0]["text"],
+                model=self.model,
+                messages=messages[1:],
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .content[0]
+            .text
+        )
+
+
+class LMMEngineGemini(LMMEngine):
+    def __init__(
+        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.base_url = base_url
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("GEMINI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
+            )
+        base_url = self.base_url or os.getenv("GEMINI_ENDPOINT_URL")
+        if base_url is None:
+            raise ValueError(
+                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named GEMINI_ENDPOINT_URL"
+            )
+        if not self.llm_client:
+            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
+        return (
+            self.llm_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .choices[0]
+            .message.content
+        )
+
+
+class LMMEngineOpenRouter(LMMEngine):
+    def __init__(
+        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.base_url = base_url
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("OPENROUTER_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENROUTER_API_KEY"
+            )
+        base_url = self.base_url or os.getenv("OPEN_ROUTER_ENDPOINT_URL")
+        if base_url is None:
+            raise ValueError(
+                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named OPEN_ROUTER_ENDPOINT_URL"
+            )
+        if not self.llm_client:
+            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
+        return (
+            self.llm_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .choices[0]
+            .message.content
+        )
+
+
+class LMMEngineAzureOpenAI(LMMEngine):
+    def __init__(
+        self,
+        base_url=None,
+        api_key=None,
+        azure_endpoint=None,
+        model=None,
+        api_version=None,
+        rate_limit=-1,
+        **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.api_version = api_version
+        self.api_key = api_key
+        self.azure_endpoint = azure_endpoint
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+        self.cost = 0.0
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("AZURE_OPENAI_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
+            )
+        api_version = self.api_version or os.getenv("OPENAI_API_VERSION")
+        if api_version is None:
+            raise ValueError(
+                "api_version must be provided either as a parameter or as an environment variable named OPENAI_API_VERSION"
+            )
+        azure_endpoint = self.azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
+        if azure_endpoint is None:
+            raise ValueError(
+                "An Azure API endpoint needs to be provided in either the azure_endpoint parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
+            )
+        if not self.llm_client:
+            self.llm_client = AzureOpenAI(
+                azure_endpoint=azure_endpoint,
+                api_key=api_key,
+                api_version=api_version,
+            )
+        completion = self.llm_client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            max_tokens=max_new_tokens if max_new_tokens else 4096,
+            temperature=temperature,
+            **kwargs,
+        )
+        total_tokens = completion.usage.total_tokens
+        self.cost += 0.02 * ((total_tokens + 500) / 1000)
+        return completion.choices[0].message.content
+
+
+class LMMEnginevLLM(LMMEngine):
+    def __init__(
+        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
+    ):
+        assert model is not None, "model must be provided"
+        self.model = model
+        self.api_key = api_key
+        self.base_url = base_url
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(
+        self,
+        messages,
+        temperature=0.0,
+        top_p=0.8,
+        repetition_penalty=1.05,
+        max_new_tokens=512,
+        **kwargs
+    ):
+        api_key = self.api_key or os.getenv("vLLM_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "A vLLM API key needs to be provided in either the api_key parameter or as an environment variable named vLLM_API_KEY"
+            )
+        base_url = self.base_url or os.getenv("vLLM_ENDPOINT_URL")
+        if base_url is None:
+            raise ValueError(
+                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named vLLM_ENDPOINT_URL"
+            )
+        if not self.llm_client:
+            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
+        completion = self.llm_client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            max_tokens=max_new_tokens if max_new_tokens else 4096,
+            temperature=temperature,
+            top_p=top_p,
+            extra_body={"repetition_penalty": repetition_penalty},
+        )
+        return completion.choices[0].message.content
+
+
+class LMMEngineHuggingFace(LMMEngine):
+    def __init__(self, base_url=None, api_key=None, rate_limit=-1, **kwargs):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("HF_TOKEN")
+        if api_key is None:
+            raise ValueError(
+                "A HuggingFace token needs to be provided in either the api_key parameter or as an environment variable named HF_TOKEN"
+            )
+        base_url = self.base_url
+        if base_url is None:
+            raise ValueError(
+                "HuggingFace endpoint must be provided as base_url parameter."
+            )
+        if not self.llm_client:
+            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
+        return (
+            self.llm_client.chat.completions.create(
+                model="tgi",
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .choices[0]
+            .message.content
+        )
+
+
+class LMMEngineParasail(LMMEngine):
+    def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
+        assert model is not None, "Parasail model id must be provided"
+        self.model = model
+        self.api_key = api_key
+        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
+        self.llm_client = None
+
+    @backoff.on_exception(
+        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
+    )
+    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
+        api_key = self.api_key or os.getenv("PARASAIL_API_KEY")
+        if api_key is None:
+            raise ValueError(
+                "A Parasail API key needs to be provided in either the api_key parameter or as an environment variable named PARASAIL_API_KEY"
+            )
+        if not self.llm_client:
+            self.llm_client = OpenAI(
+                base_url="https://api.parasail.io/v1", api_key=api_key
+            )
+        return (
+            self.llm_client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                max_tokens=max_new_tokens if max_new_tokens else 4096,
+                temperature=temperature,
+                **kwargs,
+            )
+            .choices[0]
+            .message.content
+        )

gui_agents/s2/core/knowledge.py

@@ -0,0 +1,420 @@
+import json
+import os
+from typing import Dict, Tuple
+
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+from gui_agents.s2.core.module import BaseModule
+from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
+from gui_agents.s2.utils.common_utils import (
+    call_llm_safe,
+    load_embeddings,
+    load_knowledge_base,
+    save_embeddings,
+)
+from gui_agents.s2.utils.query_perplexica import query_to_perplexica
+
+
+class KnowledgeBase(BaseModule):
+    def __init__(
+        self,
+        embedding_engine,
+        local_kb_path: str,
+        platform: str,
+        engine_params: Dict,
+        save_knowledge: bool = True,
+    ):
+        super().__init__(engine_params, platform)
+
+        self.local_kb_path = local_kb_path
+
+        # initialize embedding engine
+        self.embedding_engine = embedding_engine
+
+        # Initialize paths for different memory types
+        self.episodic_memory_path = os.path.join(
+            self.local_kb_path, self.platform, "episodic_memory.json"
+        )
+        self.narrative_memory_path = os.path.join(
+            self.local_kb_path, self.platform, "narrative_memory.json"
+        )
+        self.embeddings_path = os.path.join(
+            self.local_kb_path, self.platform, "embeddings.pkl"
+        )
+
+        # Initialize trajectory tracking
+        self.task_trajectory = ""
+        self.current_subtask_trajectory = ""
+        self.current_search_query = ""
+
+        self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
+            "CURRENT_OS", self.platform
+        )
+
+        # All three agents share a generic RAG prompt that asks the agent to provide information for UI automation in CURRENT_OS
+        self.query_formulator = self._create_agent(self.rag_module_system_prompt)
+        self.llm_search_agent = self._create_agent(self.rag_module_system_prompt)
+        self.knowledge_fusion_agent = self._create_agent(self.rag_module_system_prompt)
+
+        self.narrative_summarization_agent = self._create_agent(
+            PROCEDURAL_MEMORY.TASK_SUMMARIZATION_PROMPT
+        )
+        self.episode_summarization_agent = self._create_agent(
+            PROCEDURAL_MEMORY.SUBTASK_SUMMARIZATION_PROMPT
+        )
+
+        self.save_knowledge = save_knowledge
+
+    def retrieve_knowledge(
+        self, instruction: str, search_query: str, search_engine: str = "llm"
+    ) -> Tuple[str, str]:
+        """Retrieve knowledge using search engine
+        Args:
+            instruction (str): task instruction
+            observation (Dict): current observation
+            search_engine (str): search engine to use"""
+
+        # Use search engine to retrieve knowledge based on the formulated query
+        search_results = self._search(instruction, search_query, search_engine)
+
+        return search_query, search_results
+
+    def formulate_query(self, instruction: str, observation: Dict) -> str:
+        """Formulate search query based on instruction and current state"""
+        query_path = os.path.join(
+            self.local_kb_path, self.platform, "formulate_query.json"
+        )
+        try:
+            with open(query_path, "r") as f:
+                formulate_query = json.load(f)
+        except:
+            formulate_query = {}
+
+        if instruction in formulate_query:
+            return formulate_query[instruction]
+
+        self.query_formulator.reset()
+
+        self.query_formulator.add_message(
+            f"The task is: {instruction}\n"
+            "To use google search to get some useful information, first carefully analyze "
+            "the screenshot of the current desktop UI state, then given the task "
+            "instruction, formulate a question that can be used to search on the Internet "
+            "for information in helping with the task execution.\n"
+            "The question should not be too general or too specific. Please ONLY provide "
+            "the question.\nQuestion:",
+            image_content=(
+                observation["screenshot"] if "screenshot" in observation else None
+            ),
+            role="user",
+        )
+
+        search_query = self.query_formulator.get_response().strip().replace('"', "")
+        print("search query: ", search_query)
+        formulate_query[instruction] = search_query
+        with open(query_path, "w") as f:
+            json.dump(formulate_query, f, indent=2)
+
+        return search_query
+
+    def _search(self, instruction: str, search_query: str, search_engine: str) -> str:
+        """Execute search using specified engine"""
+
+        # Default to perplexica rag knowledge to see if the query exists
+        file = os.path.join(
+            self.local_kb_path, self.platform, f"{search_engine}_rag_knowledge.json"
+        )
+
+        try:
+            with open(file, "r") as f:
+                exist_search_results = json.load(f)
+        except:
+            exist_search_results = {}
+
+        if instruction in exist_search_results:
+            return exist_search_results[instruction]
+        if search_engine.lower() == "llm":
+            self.llm_search_agent.reset()
+            # Use LLM's internal knowledge like a search engine
+            self.llm_search_agent.add_message(search_query, role="user")
+            search_results = self.llm_search_agent.get_response()
+        elif search_engine.lower() == "perplexica":
+            # Use perplexica to search for the query
+            search_results = query_to_perplexica(search_query)
+        else:
+            raise ValueError(f"Unsupported search engine: {search_engine}")
+
+        exist_search_results[instruction] = search_results.strip()
+        with open(
+            os.path.join(
+                self.local_kb_path,
+                self.platform,
+                f"{search_engine}_rag_knowledge.json",
+            ),
+            "w",
+        ) as f:
+            json.dump(exist_search_results, f, indent=2)
+
+        return search_results
+
+    def retrieve_narrative_experience(self, instruction: str) -> Tuple[str, str]:
+        """Retrieve narrative experience using embeddings"""
+
+        knowledge_base = load_knowledge_base(self.narrative_memory_path)
+        if not knowledge_base:
+            return "None", "None"
+
+        embeddings = load_embeddings(self.embeddings_path)
+
+        # Get or create instruction embedding
+        instruction_embedding = embeddings.get(instruction)
+
+        if instruction_embedding is None:
+            instruction_embedding = self.embedding_engine.get_embeddings(instruction)
+            embeddings[instruction] = instruction_embedding
+
+        # Get or create embeddings for knowledge base entries
+        candidate_embeddings = []
+        for key in knowledge_base:
+            candidate_embedding = embeddings.get(key)
+            if candidate_embedding is None:
+                candidate_embedding = self.embedding_engine.get_embeddings(key)
+                embeddings[key] = candidate_embedding
+
+            candidate_embeddings.append(candidate_embedding)
+
+        save_embeddings(self.embeddings_path, embeddings)
+
+        similarities = cosine_similarity(
+            instruction_embedding, np.vstack(candidate_embeddings)
+        )[0]
+        sorted_indices = np.argsort(similarities)[::-1]
+
+        keys = list(knowledge_base.keys())
+        idx = 1 if keys[sorted_indices[0]] == instruction else 0
+        return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]
+
+    def retrieve_episodic_experience(self, instruction: str) -> Tuple[str, str]:
+        """Retrieve similar task experience using embeddings"""
+
+        knowledge_base = load_knowledge_base(self.episodic_memory_path)
+        if not knowledge_base:
+            return "None", "None"
+
+        embeddings = load_embeddings(self.embeddings_path)
+
+        # Get or create instruction embedding
+        instruction_embedding = embeddings.get(instruction)
+
+        if instruction_embedding is None:
+            instruction_embedding = self.embedding_engine.get_embeddings(instruction)
+            embeddings[instruction] = instruction_embedding
+
+        # Get or create embeddings for knowledge base entries
+        candidate_embeddings = []
+        for key in knowledge_base:
+            candidate_embedding = embeddings.get(key)
+            if candidate_embedding is None:
+                candidate_embedding = self.embedding_engine.get_embeddings(key)
+                embeddings[key] = candidate_embedding
+
+            candidate_embeddings.append(candidate_embedding)
+
+        save_embeddings(self.embeddings_path, embeddings)
+
+        similarities = cosine_similarity(
+            instruction_embedding, np.vstack(candidate_embeddings)
+        )[0]
+        sorted_indices = np.argsort(similarities)[::-1]
+
+        keys = list(knowledge_base.keys())
+        idx = 1 if keys[sorted_indices[0]] == instruction else 0
+        return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]
+
+    def knowledge_fusion(
+        self,
+        observation: Dict,
+        instruction: str,
+        web_knowledge: str,
+        similar_task: str,
+        experience: str,
+    ) -> str:
+        """Combine web knowledge with similar task experience"""
+
+        self.knowledge_fusion_agent.reset()
+
+        self.knowledge_fusion_agent.add_message(
+            f"Task: {instruction}\n"
+            f"**Web search result**:\n{web_knowledge}\n\n"
+            f"**Retrieved similar task experience**:\n"
+            f"Similar task:{similar_task}\n{experience}\n\n"
+            f"Based on the web search result and the retrieved similar task experience, "
+            f"if you think the similar task experience is indeed useful to the main task, "
+            f"integrate it with the web search result. Provide the final knowledge in a numbered list.",
+            image_content=(
+                observation["screenshot"] if "screenshot" in observation else None
+            ),
+            role="user",
+        )
+        return self.knowledge_fusion_agent.get_response()
+
+    def save_episodic_memory(self, subtask_key: str, subtask_traj: str) -> None:
+        """Save episodic memory (subtask level knowledge).
+
+        Args:
+            subtask_key (str): Key identifying the subtask
+            subtask_traj (str): Trajectory/experience of the subtask
+        """
+        if not self.save_knowledge:
+            return
+
+        try:
+            kb = load_knowledge_base(self.episodic_memory_path)
+        except:
+            kb = {}
+
+        if subtask_key not in kb:
+            subtask_summarization = self.summarize_episode(subtask_traj)
+            kb[subtask_key] = subtask_summarization
+
+            os.makedirs(os.path.dirname(self.episodic_memory_path), exist_ok=True)
+            with open(self.episodic_memory_path, "w") as fout:
+                json.dump(kb, fout, indent=2)
+
+        return kb.get(subtask_key)
+
+    def save_narrative_memory(self, task_key: str, task_traj: str) -> None:
+        """Save narrative memory (task level knowledge).
+
+        Args:
+            task_key (str): Key identifying the task
+            task_traj (str): Full trajectory/experience of the task
+        """
+        if not self.save_knowledge:
+            return
+
+        try:
+            kb = load_knowledge_base(self.narrative_memory_path)
+        except:
+            kb = {}
+
+        if task_key not in kb:
+            task_summarization = self.summarize_narrative(task_traj)
+            kb[task_key] = task_summarization
+
+            os.makedirs(os.path.dirname(self.narrative_memory_path), exist_ok=True)
+            with open(self.narrative_memory_path, "w") as fout:
+                json.dump(kb, fout, indent=2)
+
+        return kb.get(task_key)
+
+    def initialize_task_trajectory(self, instruction: str) -> None:
+        """Initialize a new task trajectory.
+
+        Args:
+            instruction (str): The task instruction
+        """
+        self.task_trajectory = f"Task:\n{instruction}"
+        self.current_search_query = ""
+        self.current_subtask_trajectory = ""
+
+    def update_task_trajectory(self, meta_data: Dict) -> None:
+        """Update the task trajectory with new metadata.
+
+        Args:
+            meta_data (Dict): Metadata from the agent's prediction
+        """
+        if not self.current_search_query and "search_query" in meta_data:
+            self.current_search_query = meta_data["search_query"]
+
+        self.task_trajectory += (
+            "\n\nReflection:\n"
+            + str(meta_data["reflection"])
+            + "\n\n----------------------\n\nPlan:\n"
+            + meta_data["executor_plan"]
+        )
+
+    def handle_subtask_trajectory(self, meta_data: Dict) -> None:
+        """Handle subtask trajectory updates based on subtask status.
+
+        Args:
+            meta_data (Dict): Metadata containing subtask information
+
+        Returns:
+            bool: Whether the subtask was completed
+        """
+        subtask_status = meta_data["subtask_status"]
+        subtask = meta_data["subtask"]
+        subtask_info = meta_data["subtask_info"]
+
+        if subtask_status in ["Start", "Done"]:
+            # If there's an existing subtask trajectory, finalize it
+            if self.current_subtask_trajectory:
+                self.current_subtask_trajectory += "\nSubtask Completed.\n"
+                subtask_key = self.current_subtask_trajectory.split(
+                    "\n----------------------\n\nPlan:\n"
+                )[0]
+                self.save_episodic_memory(subtask_key, self.current_subtask_trajectory)
+                self.current_subtask_trajectory = ""
+                return True
+
+            # Start new subtask trajectory
+            self.current_subtask_trajectory = (
+                f"Task:\n{self.current_search_query}\n\n"
+                f"Subtask: {subtask}\n"
+                f"Subtask Instruction: {subtask_info}\n"
+                f"----------------------\n\n"
+                f'Plan:\n{meta_data["executor_plan"]}\n'
+            )
+            return False
+
+        elif subtask_status == "In":
+            # Continue current subtask trajectory
+            self.current_subtask_trajectory += (
+                f'\n----------------------\n\nPlan:\n{meta_data["executor_plan"]}\n'
+            )
+            return False
+
+    def finalize_task(self) -> None:
+        """Finalize the task by saving any remaining trajectories."""
+        # Save any remaining subtask trajectory
+        if self.current_subtask_trajectory:
+            self.current_subtask_trajectory += "\nSubtask Completed.\n"
+            subtask_key = self.current_subtask_trajectory.split(
+                "\n----------------------\n\nPlan:\n"
+            )[0]
+            self.save_episodic_memory(subtask_key, self.current_subtask_trajectory)
+
+        # Save the complete task trajectory
+        if self.task_trajectory and self.current_search_query:
+            self.save_narrative_memory(self.current_search_query, self.task_trajectory)
+
+        # Reset trajectories
+        self.task_trajectory = ""
+        self.current_subtask_trajectory = ""
+        self.current_search_query = ""
+
+    def summarize_episode(self, trajectory):
+        """Summarize the episode experience for lifelong learning reflection
+        Args:
+            trajectory: str: The episode experience to be summarized
+        """
+
+        # Create Reflection on whole trajectories for next round trial, keep earlier messages as exemplars
+        self.episode_summarization_agent.add_message(trajectory)
+        subtask_summarization = call_llm_safe(self.episode_summarization_agent)
+        self.episode_summarization_agent.add_message(subtask_summarization)
+
+        return subtask_summarization
+
+    def summarize_narrative(self, trajectory):
+        """Summarize the narrative experience for lifelong learning reflection
+        Args:
+            trajectory: str: The narrative experience to be summarized
+        """
+        # Create Reflection on whole trajectories for next round trial
+        self.narrative_summarization_agent.add_message(trajectory)
+        task_summarization = call_llm_safe(self.narrative_summarization_agent)
+
+        return task_summarization

gui_agents/s2/core/mllm.py

@@ -0,0 +1,295 @@
+import base64
+
+import numpy as np
+
+from gui_agents.s2.core.engine import (
+    LMMEngineAnthropic,
+    LMMEngineAzureOpenAI,
+    LMMEngineHuggingFace,
+    LMMEngineOpenAI,
+    LMMEngineOpenRouter,
+    LMMEngineParasail,
+    LMMEnginevLLM,
+    LMMEngineGemini,
+)
+
+
+class LMMAgent:
+    def __init__(self, engine_params=None, system_prompt=None, engine=None):
+        if engine is None:
+            if engine_params is not None:
+                engine_type = engine_params.get("engine_type")
+                if engine_type == "openai":
+                    self.engine = LMMEngineOpenAI(**engine_params)
+                elif engine_type == "anthropic":
+                    self.engine = LMMEngineAnthropic(**engine_params)
+                elif engine_type == "azure":
+                    self.engine = LMMEngineAzureOpenAI(**engine_params)
+                elif engine_type == "vllm":
+                    self.engine = LMMEnginevLLM(**engine_params)
+                elif engine_type == "huggingface":
+                    self.engine = LMMEngineHuggingFace(**engine_params)
+                elif engine_type == "gemini":
+                    self.engine = LMMEngineGemini(**engine_params)
+                elif engine_type == "open_router":
+                    self.engine = LMMEngineOpenRouter(**engine_params)
+                elif engine_type == "parasail":
+                    self.engine = LMMEngineParasail(**engine_params)
+                else:
+                    raise ValueError("engine_type is not supported")
+            else:
+                raise ValueError("engine_params must be provided")
+        else:
+            self.engine = engine
+
+        self.messages = []  # Empty messages
+
+        if system_prompt:
+            self.add_system_prompt(system_prompt)
+        else:
+            self.add_system_prompt("You are a helpful assistant.")
+
+    def encode_image(self, image_content):
+        # if image_content is a path to an image file, check type of the image_content to verify
+        if isinstance(image_content, str):
+            with open(image_content, "rb") as image_file:
+                return base64.b64encode(image_file.read()).decode("utf-8")
+        else:
+            return base64.b64encode(image_content).decode("utf-8")
+
+    def reset(
+        self,
+    ):
+
+        self.messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": self.system_prompt}],
+            }
+        ]
+
+    def add_system_prompt(self, system_prompt):
+        self.system_prompt = system_prompt
+        if len(self.messages) > 0:
+            self.messages[0] = {
+                "role": "system",
+                "content": [{"type": "text", "text": self.system_prompt}],
+            }
+        else:
+            self.messages.append(
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": self.system_prompt}],
+                }
+            )
+
+    def remove_message_at(self, index):
+        """Remove a message at a given index"""
+        if index < len(self.messages):
+            self.messages.pop(index)
+
+    def replace_message_at(
+        self, index, text_content, image_content=None, image_detail="high"
+    ):
+        """Replace a message at a given index"""
+        if index < len(self.messages):
+            self.messages[index] = {
+                "role": self.messages[index]["role"],
+                "content": [{"type": "text", "text": text_content}],
+            }
+            if image_content:
+                base64_image = self.encode_image(image_content)
+                self.messages[index]["content"].append(
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{base64_image}",
+                            "detail": image_detail,
+                        },
+                    }
+                )
+
+    def add_message(
+        self,
+        text_content,
+        image_content=None,
+        role=None,
+        image_detail="high",
+        put_text_last=False,
+    ):
+        """Add a new message to the list of messages"""
+
+        # API-style inference from OpenAI and AzureOpenAI
+        if isinstance(
+            self.engine,
+            (
+                LMMEngineOpenAI,
+                LMMEngineAzureOpenAI,
+                LMMEngineHuggingFace,
+                LMMEngineGemini,
+                LMMEngineOpenRouter,
+                LMMEngineParasail,
+            ),
+        ):
+            # infer role from previous message
+            if role != "user":
+                if self.messages[-1]["role"] == "system":
+                    role = "user"
+                elif self.messages[-1]["role"] == "user":
+                    role = "assistant"
+                elif self.messages[-1]["role"] == "assistant":
+                    role = "user"
+
+            message = {
+                "role": role,
+                "content": [{"type": "text", "text": text_content}],
+            }
+
+            if isinstance(image_content, np.ndarray) or image_content:
+                # Check if image_content is a list or a single image
+                if isinstance(image_content, list):
+                    # If image_content is a list of images, loop through each image
+                    for image in image_content:
+                        base64_image = self.encode_image(image)
+                        message["content"].append(
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/png;base64,{base64_image}",
+                                    "detail": image_detail,
+                                },
+                            }
+                        )
+                else:
+                    # If image_content is a single image, handle it directly
+                    base64_image = self.encode_image(image_content)
+                    message["content"].append(
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{base64_image}",
+                                "detail": image_detail,
+                            },
+                        }
+                    )
+
+            # Rotate text to be the last message if desired
+            if put_text_last:
+                text_content = message["content"].pop(0)
+                message["content"].append(text_content)
+
+            self.messages.append(message)
+
+        # For API-style inference from Anthropic
+        elif isinstance(self.engine, LMMEngineAnthropic):
+            # infer role from previous message
+            if role != "user":
+                if self.messages[-1]["role"] == "system":
+                    role = "user"
+                elif self.messages[-1]["role"] == "user":
+                    role = "assistant"
+                elif self.messages[-1]["role"] == "assistant":
+                    role = "user"
+
+            message = {
+                "role": role,
+                "content": [{"type": "text", "text": text_content}],
+            }
+
+            if image_content:
+                # Check if image_content is a list or a single image
+                if isinstance(image_content, list):
+                    # If image_content is a list of images, loop through each image
+                    for image in image_content:
+                        base64_image = self.encode_image(image)
+                        message["content"].append(
+                            {
+                                "type": "image",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": "image/png",
+                                    "data": base64_image,
+                                },
+                            }
+                        )
+                else:
+                    # If image_content is a single image, handle it directly
+                    base64_image = self.encode_image(image_content)
+                    message["content"].append(
+                        {
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": "image/png",
+                                "data": base64_image,
+                            },
+                        }
+                    )
+            self.messages.append(message)
+
+        # Locally hosted vLLM model inference
+        elif isinstance(self.engine, LMMEnginevLLM):
+            # infer role from previous message
+            if role != "user":
+                if self.messages[-1]["role"] == "system":
+                    role = "user"
+                elif self.messages[-1]["role"] == "user":
+                    role = "assistant"
+                elif self.messages[-1]["role"] == "assistant":
+                    role = "user"
+
+            message = {
+                "role": role,
+                "content": [{"type": "text", "text": text_content}],
+            }
+
+            if image_content:
+                # Check if image_content is a list or a single image
+                if isinstance(image_content, list):
+                    # If image_content is a list of images, loop through each image
+                    for image in image_content:
+                        base64_image = self.encode_image(image)
+                        message["content"].append(
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image;base64,{base64_image}"
+                                },
+                            }
+                        )
+                else:
+                    # If image_content is a single image, handle it directly
+                    base64_image = self.encode_image(image_content)
+                    message["content"].append(
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image;base64,{base64_image}"},
+                        }
+                    )
+
+            self.messages.append(message)
+        else:
+            raise ValueError("engine_type is not supported")
+
+    def get_response(
+        self,
+        user_message=None,
+        messages=None,
+        temperature=0.0,
+        max_new_tokens=None,
+        **kwargs,
+    ):
+        """Generate the next response based on previous messages"""
+        if messages is None:
+            messages = self.messages
+        if user_message:
+            messages.append(
+                {"role": "user", "content": [{"type": "text", "text": user_message}]}
+            )
+
+        return self.engine.generate(
+            messages,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            **kwargs,
+        )

gui_agents/s2/core/module.py

@@ -0,0 +1,17 @@
+from typing import Dict, Optional
+from gui_agents.s2.core.mllm import LMMAgent
+
+
+class BaseModule:
+    def __init__(self, engine_params: Dict, platform: str):
+        self.engine_params = engine_params
+        self.platform = platform
+
+    def _create_agent(
+        self, system_prompt: str = None, engine_params: Optional[Dict] = None
+    ) -> LMMAgent:
+        """Create a new LMMAgent instance"""
+        agent = LMMAgent(engine_params or self.engine_params)
+        if system_prompt:
+            agent.add_system_prompt(system_prompt)
+        return agent

gui_agents/s2/memory/procedural_memory.py

@@ -0,0 +1,246 @@
+import inspect
+import textwrap
+
+
+class PROCEDURAL_MEMORY:
+
+    @staticmethod
+    def construct_worker_procedural_memory(agent_class, skipped_actions):
+        procedural_memory = textwrap.dedent(f"""\
+        You are an expert in graphical user interfaces and Python code. You are responsible for executing the current subtask: `SUBTASK_DESCRIPTION` of the larger goal: `TASK_DESCRIPTION`.
+        IMPORTANT: ** The subtasks: ['DONE_TASKS'] have already been done. The future subtasks ['FUTURE_TASKS'] will be done in the future by me. You must only perform the current subtask: `SUBTASK_DESCRIPTION`. Do not try to do future subtasks. **
+        You are working in CURRENT_OS. You must only complete the subtask provided and not the larger goal.
+        You are provided with:
+        1. A screenshot of the current time step.
+        2. The history of your previous interactions with the UI.
+        3. Access to the following class and methods to interact with the UI:
+        class Agent:
+        """)
+
+        for attr_name in dir(agent_class):
+            if attr_name in skipped_actions:
+                continue
+
+            attr = getattr(agent_class, attr_name)
+            if callable(attr) and hasattr(attr, "is_agent_action"):
+                # Use inspect to get the full function signature
+                signature = inspect.signature(attr)
+                procedural_memory += f"""
+    def {attr_name}{signature}:
+    '''{attr.__doc__}'''
+        """
+
+        procedural_memory += textwrap.dedent("""
+        Your response should be formatted like this:
+        (Previous action verification)
+        Carefully analyze based on the screenshot if the previous action was successful. If the previous action was not successful, provide a reason for the failure.
+
+        (Screenshot Analysis)
+        Closely examine and describe the current state of the desktop along with the currently open applications.
+
+        (Next Action)
+        Based on the current screenshot and the history of your previous interaction with the UI, decide on the next action in natural language to accomplish the given task.
+
+        (Grounded Action)
+        Translate the next action into code using the provided API methods. Format the code like this:
+        ```python
+        agent.click("The menu button at the top right of the window", 1, "left")
+        ```
+        Note for the code:
+        1. Only perform one action at a time.
+        2. Do not put anything other than python code in the block. You can only use one function call at a time. Do not put more than one function call in the block.
+        3. You must use only the available methods provided above to interact with the UI, do not invent new methods.
+        4. Only return one code block every time. There must be a single line of code in the code block.
+        5. If you think the task is already completed, return `agent.done()` in the code block.
+        6. If you think the task cannot be completed, return `agent.fail()` in the code block.
+        7. Do not do anything other than the exact specified task. Return with `agent.done()` immediately after the task is completed or `agent.fail()` if it cannot be completed.
+        8. Whenever possible, your grounded action should use hot-keys with the agent.hotkey() action instead of clicking or dragging.
+        9. My computer's password is 'password', feel free to use it when you need sudo rights.
+        10. Do not use the "command" + "tab" hotkey on MacOS.
+        """)
+
+        return procedural_memory.strip()
+
+    # Manager prompt that generalizes to initial planning, re-planning after subtask completion, and re-planning after failure
+    COMBINED_MANAGER_PROMPT = textwrap.dedent("""
+    You are an expert planning agent for solving GUI navigation tasks. You need to generate a plan for solving the following task: TASK_DESCRIPTION.
+
+    You are provided with:
+    1. The state of the computer screen through a desktop screenshot and other related information
+    2. (If available) A list of successfully completed subtasks
+    3. (If available) A list of future remaining subtasks
+
+    Your responsibilities:
+    1. Generate a new plan or revise the pre-existing plan to complete the task
+    2. Ensure the plan is concise and contains only necessary steps
+    3. Carefully observe and understand the current state of the computer before generating your plan
+    4. Avoid including steps in your plan that the task does not ask for
+
+    Below are important considerations when generating your plan:
+    1. Provide the plan in a step-by-step format with detailed descriptions for each subtask.
+    2. Do not repeat subtasks that have already been successfully completed. Only plan for the remainder of the main task.
+    3. Do not include verification steps in your planning. Steps that confirm or validate other subtasks should not be included.
+    4. Do not include optional steps in your planning. Your plan must be as concise as possible.
+    5. Do not include unnecessary steps in your planning. If you are unsure if a step is necessary, do not include it in your plan.
+    6. When revising an existing plan:
+      - If you feel the trajectory and future subtasks seem correct based on the current state of the desktop, you may re-use future subtasks.
+      - If you feel some future subtasks are not detailed enough, use your observations from the desktop screenshot to update these subtasks to be more detailed.
+      - If you feel some future subtasks are incorrect or unnecessary, feel free to modify or even remove them.
+    """)
+
+    # USED IN OSWORLD EXPERIMENTS
+    RAG_AGENT_OSWORLD = """
+    Given a desktop computer task instruction, you are an agent which should provide useful information as requested, to help another agent follow the instruction and perform the task.
+    The domain of the desktop computer task is from [CURRENT_OS, VLC, LibreOffice, Chrome, Thunderbird, VS Code, GIMP].
+    The task is: TASK_DESCRIPTION
+    The simplified accessibility tree of the current computer UI is: ACCESSIBLITY_TREE
+    """
+
+    RAG_AGENT = """
+    Given a desktop computer task instruction, you are an agent which should provide useful information as requested, to help another agent follow the instruction and perform the task in CURRENT_OS.
+    """
+
+    # For reflection agent, post-action verification mainly for cycle detection
+    REFLECTION_ON_TRAJECTORY = textwrap.dedent("""
+    You are a reflection agent designed to assist in subtask execution by reflecting on the trajectory of a subtask and providing feedback for what the next step should be.
+    You have access to the Subtask Description and the Current Trajectory of another computer agent. The Current Trajectory is a sequence of a desktop image, chain-of-thought reasoning, and a desktop action for each time step. The last image is the screen's display after the last action.
+    Your task is to generate a reflection. Your generated reflection must fall under one of the two cases listed below:
+
+    Case 1. The trajectory is not going according to plan. This is often due to the latest action not being executed correctly, or a cycle of actions being continually repeated with no progress being made. In this case, explicitly highlight why the current trajectory is incorrect, and encourage the computer agent to try a new action. However, DO NOT encourage a specific action in particular.
+    Case 2. The trajectory is going according to plan. In this case, simply tell the agent to continue proceeding as planned. DO NOT encourage a specific action in particular.
+    
+    To be successful, you must follow the rules below:
+    - DO NOT suggest any specific future plans or actions. Your only goal is to provide a reflection, not an actual plan or action.
+    - Any response that falls under Case 1 should explain why the trajectory is not going according to plan. You should especially lookout for cycles of actions that are continually repeated with no progress.
+    - Any response that falls under Case 2 should be concise, since you just need to affirm the agent to continue with the current trajectory.
+    """)
+
+    TASK_SUMMARIZATION_PROMPT = """
+    You are a summarization agent designed to analyze a trajectory of desktop task execution.
+    You have access to the Task Description and Whole Trajectory including plan, verification and reflection at each step.
+    Your summarized information will be referred to by another agent when performing the tasks.
+    You should follow the below instructions:
+    1. If the task is successfully executed, you should summarize the successful plan based on the whole trajectory to finish the task.
+    2. Otherwise, provide the reasons why the task is failed and potential suggestions that may avoid this failure.
+
+    **ATTENTION**
+    1. Only extract the correct plan and do not provide redundant steps.
+    2. Do not contain grounded actions in the plan.
+    3. If there are the successfully used hot-keys, make sure to include them in the plan.
+    4. The suggestions are for another agent not human, so they must be doable through the agent's action.
+    5. Don't generate high-level suggestions (e.g., Implement Error Handling).
+    """
+
+    DAG_TRANSLATOR_PROMPT = """You are a plan to Dependency Graph conversion agent. Your task is to analyze a given plan and generate a structured JSON output representing the plan and its corresponding directed acyclic graph (DAG).
+
+The output should be a valid JSON object wrapped in <json></json> tags, with the following structure:
+
+<json>
+{
+  "dag": {
+    "nodes": [
+      {
+        "name": "Short name or brief description of the step",
+        "info": "Detailed information about executing this step"
+      }
+    ],
+    "edges": [
+      [
+        {"name": "Name of the source node", "info": "Info of the source node"},
+        {"name": "Name of the target node", "info": "Info of the target node"}
+      ]
+    ]
+  }
+}
+</json>
+
+Important guidelines you must follow:
+1. The "plan" field should contain the entire original plan as a string.
+2. In the "dag" object:
+   a. Each node in the "nodes" array should contain 'name' and 'info' fields.
+   b. 'name' should be a concise, one-line description of the subtask.
+   c. 'info' should contain all available information about executing that subtask from the original plan. Do not remove or edit any information from the 'info' field.
+3. The "edges" array should represent the connections between nodes, showing the order and dependencies of the steps.
+4. If the plan only has one subtask, you MUST construct a graph with a SINGLE node. The "nodes" array should have that single subtask as a node, and the "edges" array should be empty.
+5. The graph must be a directed acyclic graph (DAG) and must be connected.
+6. Do not include completed subtasks in the graph. A completed subtask must not be included in a node or an edge.
+7. Do not include repeated or optional steps in the graph. Any extra information should be incorporated into the 'info' field of the relevant node.
+8. It is okay for the graph to have a single node and no edges, if the provided plan only has one subtask.
+
+Analyze the given plan and provide the output in this JSON format within the <json></json> tags. Ensure the JSON is valid and properly escaped.
+"""
+
+    SUBTASK_SUMMARIZATION_PROMPT = textwrap.dedent("""
+    You are a summarization agent designed to analyze a trajectory of desktop task execution.
+    You will summarize the correct plan and grounded actions based on the whole trajectory of a subtask, ensuring the summarized plan contains only correct and necessary steps.
+
+    **ATTENTION**
+	  1.	Summarize the correct plan and its corresponding grounded actions. Carefully filter out any repeated or incorrect steps based on the verification output in the trajectory. Only include the necessary steps for successfully completing the subtask.
+    2.	Description Replacement in Grounded Actions:
+        When summarizing grounded actions, the agent.click() and agent.drag_and_drop() grounded actions take a description string as an argument.
+        Replace these description strings with placeholders like \"element1_description\", \"element2_description\", etc., while maintaining the total number of parameters.
+        For example, agent.click(\"The menu button in the top row\", 1) should be converted into agent.click(\"element1_description\", 1)
+        Ensure the placeholders (\"element1_description\", \"element2_description\", ...) follow the order of appearance in the grounded actions.
+	  3.	Only generate grounded actions that are explicitly present in the trajectory. Do not introduce any grounded actions that do not exist in the trajectory.
+	  4.	For each step in the plan, provide a corresponding grounded action. Use the exact format:
+    	  Action: [Description of the correct action]
+    	  Grounded Action: [Grounded actions with the \"element1_description\" replacement when needed]
+	  5.	Exclude any other details that are not necessary for completing the task.
+    """)
+
+    STATE_EVALUATOR_SYSTEM_PROMPT = """
+    You are an impartial evaluator to evaluate the completeness of the given desktop computer task, you are also an expert of accessibility tree, os environment and python programming.
+    The task is: TASK_DESCRIPTION, it is executed by a digital agent who can perform the task without knowing whether the task requirements are met.
+    As an evaluator, your task is to judge whether the task is finished and meets the task requirement.
+    You have access to the:
+    1. Task instruction.
+    2. The whole actions performed by the digital agent.
+    3. The accessibility tree at the first step and the last step.
+    4. The screenshot at the first step and the last step.
+
+    You are able to proceed your judgment process in the following ways based on the task instruction:
+    1. By comparing the difference in the accessibility trees of the UI, you should judge whether the task is complete given the task instruction.
+    2. If you cannot judge based on the observations, you can evalaute it by writing and running a python script to do a further examination. For example, you can use the 'subprocess' module to run the external command in a terminal to check whether an application has been installed.
+    You can also call the file system API to do the file check, etc. You can also try to interactive with the environment via other methods or interface you are familiared with.
+
+    **IMPORTANT**
+    1. If no python script is needed, you should provide your analysis and put the judgment at the end of the response in this format: Judgment: Yes/No
+    2. Otherwise, you should format your response into two parts as shown below:
+        ```python
+        # your code script here
+        ```
+
+    **ATTENTION**
+    1. You should only use scripts when you have to.
+    2. When you generate code script, only return one code block every time, the code block should contain the whole script you want to run. You must guarantee that the script is comprehensive and executable, make sure to print out the scripts' results for subsequent judgement.
+    Additionally, the comment of the code is **PROHIBITED**
+    3. You should strictly follow the response format mentioned above.
+
+    **SUBSEQUENCE**
+    If you have generated the python script, I will execute it and return the corresponding result to you (Started with "The output after executing the script is:..."). Then you should judge whether the task has been completed or not comprehensively based on the script and its result,
+    the task information, and the comparison of accessibility trees and screenshots. Provide your analysis and put the judgment at the end of the response in this format: Judgment: Yes/No
+    """
+
+    OBS_EVALUATOR_SYSTEM_PROMPT = """
+    You are an impartial evaluator to evaluate the completeness of the given desktop computer task.
+    The task is: TASK_DESCRIPTION, it is executed by a digital agent who can perform the task without knowing whether the task requirements are met.
+    As an evaluator, your task is to judge whether the task is finished and meets the task requirement.
+    You have access to the task instruction, the whole actions performed by the digital agent, the accessibility tree of the UI and screenshot at the first time step and the last time step.
+    By comparing the difference in the accessibility trees of the UI, you should judge whether the task is complete given the task instruction.
+    Provide your analysis and put the judgment at the end of the response in this format:
+    Judgment: Yes/No
+    Only say Yes or No in the Judgment section. Do not provide any other information in the Judgment section.
+    """
+
+    PHRASE_TO_WORD_COORDS_PROMPT = textwrap.dedent("""
+    You are an expert in graphical user interfaces. Your task is to process a phrase of text, and identify the most relevant word on the computer screen.
+    You are provided with a phrase, a table with all the text on the screen, and a screenshot of the computer screen. You will identify the single word id that is best associated with the provided phrase.
+    This single word must be displayed on the computer screenshot, and its location on the screen should align with the provided phrase.
+    Each row in the text table provides 2 pieces of data in the following order. 1st is the unique word id. 2nd is the corresponding word.
+
+    To be successful, it is very important to follow all these rules:
+    1. First, think step by step and generate your reasoning about which word id to click on.
+    2. Then, output the unique word id. Remember, the word id is the 1st number in each row of the text table.
+    3. If there are multiple occurrences of the same word, use the surrounding context in the phrase to choose the correct one. Pay very close attention to punctuation and capitalization.
+
+    """)

gui_agents/s2/utils/common_utils.py

@@ -0,0 +1,223 @@
+import json
+import re
+from typing import List
+import time
+import tiktoken
+
+from typing import Tuple, List, Union, Dict
+
+from pydantic import BaseModel, ValidationError
+
+import pickle
+
+
+class Node(BaseModel):
+    name: str
+    info: str
+
+
+class Dag(BaseModel):
+    nodes: List[Node]
+    edges: List[List[Node]]
+
+
+NUM_IMAGE_TOKEN = 1105  # Value set of screen of size 1920x1080 for openai vision
+
+
+def call_llm_safe(agent) -> Union[str, Dag]:
+    # Retry if fails
+    max_retries = 3  # Set the maximum number of retries
+    attempt = 0
+    response = ""
+    while attempt < max_retries:
+        try:
+            response = agent.get_response()
+            break  # If successful, break out of the loop
+        except Exception as e:
+            attempt += 1
+            print(f"Attempt {attempt} failed: {e}")
+            if attempt == max_retries:
+                print("Max retries reached. Handling failure.")
+        time.sleep(1.0)
+    return response
+
+
+def calculate_tokens(messages, num_image_token=NUM_IMAGE_TOKEN) -> Tuple[int, int]:
+
+    num_input_images = 0
+    output_message = messages[-1]
+
+    input_message = messages[:-1]
+
+    input_string = """"""
+    for message in input_message:
+        input_string += message["content"][0]["text"] + "\n"
+        if len(message["content"]) > 1:
+            num_input_images += 1
+
+    input_text_tokens = get_input_token_length(input_string)
+
+    input_image_tokens = num_image_token * num_input_images
+
+    output_tokens = get_input_token_length(output_message["content"][0]["text"])
+
+    return (input_text_tokens + input_image_tokens), output_tokens
+
+
+# Code based on https://github.com/xlang-ai/OSWorld/blob/main/mm_agents/agent.py
+
+
+def parse_dag(text):
+    pattern = r"<json>(.*?)</json>"
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        json_str = match.group(1)
+        try:
+            json_data = json.loads(json_str)
+            return Dag(**json_data["dag"])
+        except json.JSONDecodeError:
+            print("Error: Invalid JSON")
+            return None
+        except KeyError:
+            print("Error: 'dag' key not found in JSON")
+            return None
+        except ValidationError as e:
+            print(f"Error: Invalid data structure - {e}")
+            return None
+    else:
+        print("Error: JSON not found")
+        return None
+
+
+def parse_dag(text):
+    """
+    Try extracting JSON from <json>…</json> tags first;
+    if not found, try ```json … ``` Markdown fences.
+    """
+
+    def _extract(pattern):
+        m = re.search(pattern, text, re.DOTALL)
+        return m.group(1).strip() if m else None
+
+    # 1) look for <json>…</json>
+    json_str = _extract(r"<json>(.*?)</json>")
+    # 2) fallback to ```json … ```
+    if json_str is None:
+        json_str = _extract(r"```json\s*(.*?)\s*```")
+
+    if json_str is None:
+        print("Error: JSON not found in either <json> tags or ```json``` fence")
+        return None
+
+    try:
+        payload = json.loads(json_str)
+    except json.JSONDecodeError as e:
+        print(f"Error: Invalid JSON ({e})")
+        return None
+
+    if "dag" not in payload:
+        print("Error: 'dag' key not found in JSON")
+        return None
+
+    try:
+        return Dag(**payload["dag"])
+    except ValidationError as e:
+        print(f"Error: Invalid data structure - {e}")
+        return None
+
+
+def parse_single_code_from_string(input_string):
+    input_string = input_string.strip()
+    if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
+        return input_string.strip()
+
+    # This regular expression will match both ```code``` and ```python code```
+    # and capture the `code` part. It uses a non-greedy match for the content inside.
+    pattern = r"```(?:\w+\s+)?(.*?)```"
+    # Find all non-overlapping matches in the string
+    matches = re.findall(pattern, input_string, re.DOTALL)
+
+    # The regex above captures the content inside the triple backticks.
+    # The `re.DOTALL` flag allows the dot `.` to match newline characters as well,
+    # so the code inside backticks can span multiple lines.
+
+    # matches now contains all the captured code snippets
+
+    codes = []
+
+    for match in matches:
+        match = match.strip()
+        commands = [
+            "WAIT",
+            "DONE",
+            "FAIL",
+        ]  # fixme: updates this part when we have more commands
+
+        if match in commands:
+            codes.append(match.strip())
+        elif match.split("\n")[-1] in commands:
+            if len(match.split("\n")) > 1:
+                codes.append("\n".join(match.split("\n")[:-1]))
+            codes.append(match.split("\n")[-1])
+        else:
+            codes.append(match)
+
+    if len(codes) <= 0:
+        return "fail"
+    return codes[0]
+
+
+def get_input_token_length(input_string):
+    enc = tiktoken.encoding_for_model("gpt-4")
+    tokens = enc.encode(input_string)
+    return len(tokens)
+
+
+def sanitize_code(code):
+    # This pattern captures the outermost double-quoted text
+    if "\n" in code:
+        pattern = r'(".*?")'
+        # Find all matches in the text
+        matches = re.findall(pattern, code, flags=re.DOTALL)
+        if matches:
+            # Replace the first occurrence only
+            first_match = matches[0]
+            code = code.replace(first_match, f'"""{first_match[1:-1]}"""', 1)
+    return code
+
+
+def extract_first_agent_function(code_string):
+    # Regular expression pattern to match 'agent' functions with any arguments, including nested parentheses
+    pattern = r'agent\.[a-zA-Z_]+\((?:[^()\'"]|\'[^\']*\'|"[^"]*")*\)'
+
+    # Find all matches in the string
+    matches = re.findall(pattern, code_string)
+
+    # Return the first match if found, otherwise return None
+    return matches[0] if matches else None
+
+
+def load_knowledge_base(kb_path: str) -> Dict:
+    try:
+        with open(kb_path, "r") as f:
+            return json.load(f)
+    except Exception as e:
+        print(f"Error loading knowledge base: {e}")
+        return {}
+
+
+def load_embeddings(embeddings_path: str) -> Dict:
+    try:
+        with open(embeddings_path, "rb") as f:
+            return pickle.load(f)
+    except Exception as e:
+        print(f"Error loading embeddings: {e}")
+        return {}
+
+
+def save_embeddings(embeddings_path: str, embeddings: Dict):
+    try:
+        with open(embeddings_path, "wb") as f:
+            pickle.dump(embeddings, f)
+    except Exception as e:
+        print(f"Error saving embeddings: {e}")

gui_agents/s2/utils/query_perplexica.py

@@ -0,0 +1,32 @@
+import requests
+import os
+
+
+def query_to_perplexica(query):
+    # Retrieve the URL from an environment variable
+    url = os.getenv("PERPLEXICA_URL")
+    if not url:
+        raise ValueError(
+            "PERPLEXICA_URL environment variable not set. It may take the form: 'http://localhost:{port}/api/search'. The port number is set in the config.toml in the Perplexica directory."
+        )
+
+    # Request Message
+    message = {"focusMode": "webSearch", "query": query, "history": [["human", query]]}
+
+    response = requests.post(url, json=message)
+
+    if response.status_code == 200:
+        return response.json()["message"]
+    elif response.status_code == 400:
+        raise ValueError(
+            "The request is malformed or missing required fields, such as FocusModel or query"
+        )
+    else:
+        raise ValueError("Internal Server Error")
+
+
+# Test Code
+if __name__ == "__main__":
+    query = "What is Agent S?"
+    response = query_to_perplexica(query)
+    print(response)