chore: import upstream snapshot with attribution

gui_agents/s1/core/AgentS.py

@@ -0,0 +1,382 @@
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Tuple
+import platform
+
+from gui_agents.s1.aci.ACI import ACI
+from gui_agents.s1.core.Manager import Manager
+from gui_agents.s1.core.Worker import Worker
+from gui_agents.s1.utils.common_utils import Node
+from gui_agents.utils import download_kb_data
+
+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 GraphSearchAgent(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 = "pyatuogui",
+        observation_type: str = "mixed",
+        search_engine: Optional[str] = None,
+        memory_root_path: str = os.getcwd(),
+        memory_folder_name: str = "kb_s1",
+        kb_release_tag: str = "v0.2.2",
+    ):
+        """Initialize GraphSearchAgent
+
+        Args:
+            engine_params: Configuration parameters for the LLM engine
+            grounding_agent: Instance of ACI class for UI interaction
+            platform: Operating system platform (macos, ubuntu)
+            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)
+            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".
+        """
+        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.
+        print("Downloading knowledge base initial Agent-S knowledge...")
+        self.local_kb_path = os.path.join(
+            self.memory_root_path, self.memory_folder_name
+        )
+
+        if not os.path.exists(self.local_kb_path):
+            download_kb_data(
+                version="s1",
+                release_tag=kb_release_tag,
+                download_dir=self.local_kb_path,
+                platform=self.platform,
+            )
+            print(
+                f"Successfully completed download of knowledge base for version s1, 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."
+            )
+
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset agent state and initialize components"""
+        # Initialize core components
+        self.planner = Manager(
+            self.engine_params,
+            self.grounding_agent,
+            platform=self.platform,
+            search_engine=self.engine,
+            local_kb_path=self.local_kb_path,
+        )
+        self.executor = Worker(
+            self.engine_params,
+            self.grounding_agent,
+            platform=self.platform,
+            local_kb_path=self.local_kb_path,
+        )
+
+        # 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_feedback: str = ""
+        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]]:
+        """Predict next UI action sequence
+
+        Args:
+            instruction: Natural language instruction
+            observation: Current UI state observation Dictionary {"accessibility_tree": str, "screenshot": bytes}
+            info: Dictionary containing additional information.
+
+        Returns:
+            Tuple of (agent info dict, list of actions)
+        """
+        # 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, then true again after a failed plan
+            if self.requires_replan:
+                logger.info("(RE)PLANNING...")
+                # failure feedback is the reason for the failure of the previous plan
+                planner_info, self.subtasks = self.planner.get_action_queue(
+                    instruction=instruction,
+                    observation=observation,
+                    failure_feedback=self.failure_feedback,
+                )
+
+                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...")
+                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
+            if "FAIL" in actions:
+                self.requires_replan = True
+                # set the failure feedback to the evaluator feedback
+                self.failure_feedback = f"Completed subtasks: {self.completed_tasks}. The subtask {self.current_subtask} cannot be completed. Please try another approach. {executor_info['plan_code']}. Please replan."
+                self.needs_next_subtask = True
+
+                # 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
+
+            elif "DONE" in actions:
+                self.requires_replan = False
+                self.completed_tasks.append(self.current_subtask)
+                self.needs_next_subtask = True
+                if self.subtasks:
+                    self.should_send_action = False
+                self.subtask_status = "Done"
+
+                self.reset_executor_state()
+
+            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/s1/core/BaseModule.py

@@ -0,0 +1,18 @@
+from typing import Dict, Optional
+
+from gui_agents.s1.mllm.MultimodalAgent 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/s1/core/Knowledge.py

@@ -0,0 +1,250 @@
+import json
+import os
+from typing import Dict, Tuple
+
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+from gui_agents.s1.core.BaseModule import BaseModule
+from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
+from gui_agents.s1.mllm.MultimodalEngine import OpenAIEmbeddingEngine
+from gui_agents.s1.utils.common_utils import (
+    load_embeddings,
+    load_knowledge_base,
+    save_embeddings,
+)
+from gui_agents.s1.utils.query_perplexica import query_to_perplexica
+
+
+class KnowledgeBase(BaseModule):
+    def __init__(
+        self,
+        local_kb_path: str,
+        platform: str,
+        engine_params: Dict,
+        use_image_for_search: bool = False,
+    ):
+        super().__init__(engine_params, platform)
+
+        self.local_kb_path = local_kb_path
+
+        # initialize embedding engine
+        # TODO: Support other embedding engines
+        self.embedding_engine = OpenAIEmbeddingEngine(
+            api_key=(
+                engine_params["api_key"]
+                if "api_key" in engine_params
+                else os.getenv("OPENAI_API_KEY")
+            )
+        )
+
+        # 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"
+        )
+
+        self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
+            "CURRENT_OS", self.platform
+        )
+
+        # All three agent share a generic RAG prompt that ask 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.use_image_for_search = use_image_for_search
+
+    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.add_message(
+            f"The task is: {instruction}\n"
+            f"Accessibility tree of the current desktop UI state: {observation['linearized_accessibility_tree']}\n"
+            "To use google search to get some useful information, first carefully analyze "
+            "the accessibility tree 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 self.use_image_for_search and "screenshot" in observation
+                else None
+            ),
+        )
+
+        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":
+            # Use LLM's internal knowledge like a search engine
+            self.llm_search_agent.add_message(search_query)
+            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.add_message(
+            f"Task: {instruction}\n"
+            f"Accessibility tree of the current desktop UI state: {observation['linearized_accessibility_tree']}\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 self.use_image_for_search and "screenshot" in observation
+                else None
+            ),
+        )
+        return self.knowledge_fusion_agent.get_response()

gui_agents/s1/core/Manager.py

@@ -0,0 +1,280 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+import platform
+
+from gui_agents.s1.aci.ACI import ACI
+from gui_agents.s1.core.BaseModule import BaseModule
+from gui_agents.s1.core.Knowledge import KnowledgeBase
+from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
+from gui_agents.s1.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,
+        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 submodules of the Manager
+        self.generator_agent = self._create_agent(PROCEDURAL_MEMORY.MANAGER_PROMPT)
+        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.knowledge_base = KnowledgeBase(self.local_kb_path, platform, engine_params)
+
+        self.planner_history = []
+
+        self.turn_count = 0
+        self.search_engine = search_engine
+        self.multi_round = multi_round
+        self.platform = platform
+
+    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)
+        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, failure_feedback: str = ""
+    ) -> Tuple[Dict, str]:
+        agent = self.grounding_agent
+
+        self.active_apps = agent.get_active_apps(observation)
+
+        tree_input = agent.linearize_and_annotate_tree(observation)
+        observation["linearized_accessibility_tree"] = tree_input
+
+        # Perform Retrieval only at the first planning step
+        if self.turn_count == 0:
+
+            self.search_query = self.knowledge_base.formulate_query(
+                instruction, observation
+            )
+
+            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
+                )
+            )
+
+        generator_message = (
+            f"Accessibility Tree: {tree_input}\n"
+            f"The clipboard contains: {agent.clipboard}."
+            f"The current open applications are {agent.get_active_apps(observation)}"
+            + (
+                f" Previous plan failed at step: {failure_feedback}"
+                if failure_feedback
+                else ""
+            )
+        )
+
+        self.generator_agent.add_message(
+            generator_message, image_content=observation.get("screenshot", None)
+        )
+
+        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)
+
+        self.planner_history.append(plan)
+
+        self.turn_count += 1
+
+        input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
+
+        # Set Cost based on GPT-4o
+        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]:
+        # Add initial instruction and plan to the agent's message history
+        self.dag_translator_agent.add_message(
+            f"Instruction: {instruction}\nPlan: {plan}"
+        )
+
+        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)
+
+        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,
+        failure_feedback: str = None,
+    ):
+        """Generate the action list based on the instruction
+        instruction:str: Instruction for the task
+        """
+        # Generate the high level plan
+        planner_info, plan = self._generate_step_by_step_plan(
+            observation, instruction, failure_feedback
+        )
+
+        # 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/s1/core/ProceduralMemory.py

@@ -0,0 +1,203 @@
+import inspect
+import textwrap
+
+
+class PROCEDURAL_MEMORY:
+    @staticmethod
+    def construct_worker_procedural_memory(agent_class):
+        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 simplified accessibility tree of the UI at the current time step.
+        2. A screenshot of the current time step.
+        3. The history of your previous interactions with the UI.
+        4. Access to the following class and methods to interact with the UI:
+        class Agent:
+        """)
+
+        for attr_name in dir(agent_class):
+            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 and the accessibility tree 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, the accessibility tree 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(123, 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.
+        3. Only return one code block every time. There must be a single line of code in the code block.
+        4. Please only use the available methods provided above to interact with the UI.
+        5. If you think the task is already completed, you can return `agent.done()` in the code block.
+        6. If you think the task cannot be completed, you can 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 use hot-keys or typing rather than mouse clicks.
+        9. My computer's password is 'password', feel free to use it when you need sudo rights
+        """)
+        return procedural_memory.strip()
+
+    # MANAGER_PROMPT = """You are a planning agent for solving GUI navigation tasks. You will be provided the initial configuration of a system including accessibility, screenshot and other information. You need to solve the following task: TASK_DESCRIPTION. You will describe in as much detail as possible the steps required to complete the task by a GUI agent. Please do not include any verification steps in your plan that is not your responsibility. IMPORTANT: Your plan should be as concize as possible and should not include any unnecessary steps. Do not fine-tune, or embellish anything or cause any side effects. Generate the plan that can be accomplished in the shortest time. Please take the current state into account when generating the plan. Please provide the plan in a step-by-step format and make sure you do not include anything that's already done in the GUI in your plan."""
+
+    # TODO: exploring this prompt
+    MANAGER_PROMPT = """You are a planning agent for solving GUI navigation tasks. You will be provided the initial configuration of a system including accessibility, screenshot and other information. You need to solve the following task: TASK_DESCRIPTION. You will describe in as much detail as possible the steps required to complete the task by a GUI agent. Please do not include any verification steps in your plan that is not your responsibility. IMPORTANT: Your plan should be as concize as possible and should not include any unnecessary steps. Do not fine-tune, or embellish anything or cause any side effects. Generate the plan that can be accomplished in the shortest time. Please take the current state into account when generating the plan. Please provide the plan in a step-by-step format and make sure you do not include anything that's already done in the GUI in your plan. You don't need to arrange the steps in order just list out everything that needs to be done. You may follow a dependency structure. Note that the execution agent that will complete your plan can't actually see everything thats visible to you."""
+
+    # NOTE: below prompt results in suboptimal initial plans
+    # MANAGER_PROMPT = """You are an expert planning agent for GUI tasks. You will be provided with an initial state of the system including accessibility, screenshot and other information and the final state represented by the task: TASK_DESCRIPTION. Tell me everything that needs to be done in order to reach the goal state. You don't need to arrange the steps in order just list out everything that needs to be done. You may follow a dependency structure."""
+
+    # 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.
+    """
+
+    # TODO: confirm this prompt
+    REFLECTION_ON_TRAJECTORY = """
+    You are a reflection agent designed to assist in task execution by analyzing a trajectory of task execution until this time step and providing feedback for the next step prediction.
+    You have access to the Task Description and Current Trajectory, and the image for each step. The most recent image is what happened after the latest action in the trajectory.
+    You should ONLY provide informative reflection feedback (potential mitigation alternatives) based on your expertise for the planning agent when you observe the abnormal trajectory (e.g., contain consecutive failures).
+    Otherwise, let the agent continue to proceed as planned.
+    Make sure to avoid providing any information about specific planning or actions and avoid generating repeated reflection feedbacks.
+    Assume the grounded action is correct, do not judge about it.
+    """
+
+    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. You will be provided a plan and you will generate a directed acyclic graph in the specified format for the plan. Each node in your graph should contain two fields name and subinfo. name is a one line description of each subtask. subinfo is all available information about executing that subtask available in the step by step plan. Please do not remove or edit any information out of the subinfo. The graph must be a directed acyclic graph. The graph must be connected. Do not include any repeated or optional steps in the graph, any extra info must go in the subinfo.
+    # """
+
+    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>
+
+Guidelines:
+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. The graph must be a directed acyclic graph (DAG) and must be connected.
+5. Do not include repeated or optional steps in the graph. Any extra information should be incorporated into the 'info' field of the relevant node.
+
+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 = """
+    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.	ID Replacement in Grounded Actions:
+    When summarizing grounded actions, replace all actual IDs with placeholders element1_id, element2_id, etc., while maintaining the total number of parameters.
+    Ensure the placeholders (element1_id, element2_id, …) 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 element_id replacement]
+	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.
+    """

gui_agents/s1/core/Worker.py

@@ -0,0 +1,256 @@
+import logging
+import os
+import re
+from typing import Dict, List, Tuple
+import platform
+
+from gui_agents.s1.aci.ACI import ACI
+from gui_agents.s1.core.BaseModule import BaseModule
+from gui_agents.s1.core.Knowledge import KnowledgeBase
+from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
+from gui_agents.s1.utils import common_utils
+from gui_agents.s1.utils.common_utils import Node, calculate_tokens, call_llm_safe
+
+logger = logging.getLogger("desktopenv.agent")
+
+
+class Worker(BaseModule):
+    def __init__(
+        self,
+        engine_params: Dict,
+        grounding_agent: ACI,
+        local_kb_path: str,
+        platform: str = platform.system().lower(),
+        search_engine: str = "perplexica",
+        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
+            search_engine: str
+                The search engine to use
+            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.enable_reflection = enable_reflection
+        self.search_engine = search_engine
+        self.use_subtask_experience = use_subtask_experience
+        self.reset()
+
+    def flush_messages(self, n):
+        # After every max_trajectory_length trajectories, remove messages from the start except the system prompt
+        for agent in [self.generator_agent]:
+            if len(agent.messages) > 2 * n + 1:
+                # Remove the user message and assistant message, both are 1 because the elements will move back after 1 pop
+                agent.remove_message_at(1)
+                agent.remove_message_at(1)
+
+    def reset(self):
+        self.generator_agent = self._create_agent(
+            PROCEDURAL_MEMORY.construct_worker_procedural_memory(
+                type(self.grounding_agent)
+            ).replace("CURRENT_OS", self.platform)
+        )
+        self.reflection_agent = self._create_agent(
+            PROCEDURAL_MEMORY.REFLECTION_ON_TRAJECTORY
+        )
+
+        self.knowledge_base = KnowledgeBase(
+            local_kb_path=self.local_kb_path,
+            platform=self.platform,
+            engine_params=self.engine_params,
+        )
+
+        self.turn_count = 0
+        self.planner_history = []
+        self.reflections = []
+        self.cost_this_turn = 0
+        self.tree_inputs = []
+        self.screenshot_inputs = []
+
+    # TODO: Experimental
+    def remove_ids_from_history(self):
+        for message in self.generator_agent.messages:
+            if message["role"] == "user":
+                for content in message["content"]:
+                    if content["type"] == "text":
+                        # Regex pattern to match lines that start with a number followed by spaces and remove the number
+                        pattern = r"^\d+\s+"
+
+                        # Apply the regex substitution on each line
+                        processed_lines = [
+                            re.sub(pattern, "", line)
+                            for line in content["text"].splitlines()
+                        ]
+
+                        # Join the processed lines back into a single string
+                        result = "\n".join(processed_lines)
+
+                        result = result.replace("id\t", "")
+
+                        # replace message content
+                        content["text"] = result
+
+    def generate_next_action(
+        self,
+        instruction: str,
+        search_query: str,
+        subtask: str,
+        subtask_info: str,
+        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
+
+        self.active_apps = agent.get_active_apps(obs)
+
+        # Get RAG knowledge, only update system message at t=0
+        if self.turn_count == 0:
+            # TODO: uncomment and fix for subtask level RAG
+            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)
+                )
+                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))
+            )
+
+        # Clear older messages - we keep full context. if you want to keep only the last n messages, you can use the flush_messages function
+        # self.flush_messages(3) # flushes generator messages
+
+        # Reflection generation
+        reflection = None
+        if self.enable_reflection and self.turn_count > 0:
+            # TODO: reuse planner history
+            self.reflection_agent.add_message(
+                "Task Description: "
+                + subtask
+                + " Instruction: "
+                + subtask_info
+                + "\n"
+                + "Current Trajectory: "
+                + "\n\n".join(self.planner_history)
+                + "\n"
+            )
+            reflection = call_llm_safe(self.reflection_agent)
+            self.reflections.append(reflection)
+            self.reflection_agent.add_message(reflection)
+
+            logger.info("REFLECTION: %s", reflection)
+
+        # Plan Generation
+        tree_input = agent.linearize_and_annotate_tree(obs)
+
+        self.remove_ids_from_history()
+
+        # Bash terminal message.
+        generator_message = (
+            (
+                f"\nYou may use the reflection on the previous trajectory: {reflection}\n"
+                if reflection
+                else ""
+            )
+            + f"Accessibility Tree: {tree_input}\n"
+            f"Text Buffer = [{','.join(agent.notes)}]. "
+            f"The current open applications are {agent.get_active_apps(obs)} and the active app is {agent.get_top_app(obs)}.\n"
+        )
+
+        print("ACTIVE APP IS: ", agent.get_top_app(obs))
+        # Only provide subinfo in the very first message to avoid over influence and redundancy
+        if self.turn_count == 0:
+            generator_message += f"Remeber 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"]
+        )
+
+        plan = call_llm_safe(self.generator_agent)
+        self.planner_history.append(plan)
+        logger.info("PLAN: %s", plan)
+
+        self.generator_agent.add_message(plan)
+
+        # Calculate input and output tokens
+        input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
+
+        # Set Cost based on GPT-4o
+        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)
+
+        # Extract code block from the plan
+        plan_code = common_utils.parse_single_code_from_string(
+            plan.split("Grounded Action")[-1]
+        )
+        plan_code = common_utils.sanitize_code(plan_code)
+        plan_code = common_utils.extract_first_agent_function(plan_code)
+        exec_code = eval(plan_code)
+
+        # If agent selects an element that was out of range, it should not be executed just send a WAIT command.
+        # TODO: should provide this as code feedback to the agent?
+        if agent.index_out_of_range_flag:
+            plan_code = "agent.wait(1.0)"
+            exec_code = eval(plan_code)
+            agent.index_out_of_range_flag = False
+
+        executor_info = {
+            "current_subtask": subtask,
+            "current_subtask_info": subtask_info,
+            "executor_plan": plan,
+            "linearized_accessibility_tree": tree_input,
+            "plan_code": plan_code,
+            "reflection": reflection,
+            "num_input_tokens_executor": input_tokens,
+            "num_output_tokens_executor": output_tokens,
+            "executor_cost": cost,
+        }
+        self.turn_count += 1
+
+        self.tree_inputs.append(tree_input)
+        self.screenshot_inputs.append(obs["screenshot"])
+
+        return executor_info, [exec_code]