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