import logging import textwrap from typing import Dict, List, Tuple from gui_agents.s2_5.agents.grounding import ACI from gui_agents.s2_5.core.module import BaseModule from gui_agents.s2_5.memory.procedural_memory import PROCEDURAL_MEMORY from gui_agents.s2_5.utils.common_utils import ( call_llm_safe, extract_first_agent_function, parse_single_code_from_string, sanitize_code, split_thinking_response, ) logger = logging.getLogger("desktopenv.agent") class Worker(BaseModule): def __init__( self, engine_params: Dict, grounding_agent: ACI, platform: str = "ubuntu", max_trajectory_length: int = 8, enable_reflection: bool = True, ): """ Worker receives the main task and generates actions, without the need of hierarchical planning Args: engine_params: Dict Parameters for the multimodal engine grounding_agent: Agent The grounding agent to use platform: str OS platform the agent runs on (darwin, linux, windows) max_trajectory_length: int The amount of images turns to keep enable_reflection: bool Whether to enable reflection """ super().__init__(engine_params, platform) self.grounding_agent = grounding_agent self.max_trajectory_length = max_trajectory_length self.enable_reflection = enable_reflection self.temperature = engine_params.get("temperature", 0.0) self.use_thinking = engine_params.get("model", "") in [ "claude-3-7-sonnet-20250219" ] self.reset() def reset(self): if self.platform != "linux": skipped_actions = ["set_cell_values"] else: skipped_actions = [] sys_prompt = PROCEDURAL_MEMORY.construct_simple_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.turn_count = 0 self.worker_history = [] self.reflections = [] self.cost_this_turn = 0 self.screenshot_inputs = [] # Flushing strategy dependant on model context limits def flush_messages(self): engine_type = self.engine_params.get("engine_type", "") # Flush strategy for long-context models: keep all text, only keep latest images if engine_type in ["anthropic", "openai", "gemini"]: max_images = self.max_trajectory_length for agent in [self.generator_agent, self.reflection_agent]: # keep latest k images img_count = 0 for i in range(len(agent.messages) - 1, -1, -1): for j in range(len(agent.messages[i]["content"])): if "image" in agent.messages[i]["content"][j].get("type", ""): img_count += 1 if img_count > max_images: del agent.messages[i]["content"][j] # Flush strategy for non-long-context models: drop full turns else: # generator msgs are alternating [user, assistant], so 2 per round if len(self.generator_agent.messages) > 2 * self.max_trajectory_length + 1: self.generator_agent.messages.pop(1) self.generator_agent.messages.pop(1) # reflector msgs are all [(user text, user image)], so 1 per round if len(self.reflection_agent.messages) > self.max_trajectory_length + 1: self.reflection_agent.messages.pop(1) def generate_next_action( self, instruction: str, obs: Dict, ) -> Tuple[Dict, List]: """ Predict the next action(s) based on the current observation. """ agent = self.grounding_agent generator_message = ( "" if self.turn_count > 0 else "The initial screen is provided. No action has been taken yet." ) # Load the task into the system prompt if self.turn_count == 0: self.generator_agent.add_system_prompt( self.generator_agent.system_prompt.replace( "TASK_DESCRIPTION", instruction ) ) # Get the per-step reflection reflection = None reflection_thoughts = None if self.enable_reflection: # Load the initial message if self.turn_count == 0: text_content = textwrap.dedent(f""" Task Description: {instruction} 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: self.reflection_agent.add_message( text_content=self.worker_history[-1], image_content=obs["screenshot"], role="user", ) full_reflection = call_llm_safe( self.reflection_agent, temperature=self.temperature, use_thinking=self.use_thinking, ) reflection, reflection_thoughts = split_thinking_response( full_reflection ) self.reflections.append(reflection) generator_message += f"REFLECTION: You may use this reflection on the previous action and overall trajectory:\n{reflection}\n" logger.info("REFLECTION: %s", reflection) # Add finalized message to conversation generator_message += f"\nCurrent Text Buffer = [{','.join(agent.notes)}]\n" self.generator_agent.add_message( generator_message, image_content=obs["screenshot"], role="user" ) full_plan = call_llm_safe( self.generator_agent, temperature=self.temperature, use_thinking=self.use_thinking, ) plan, plan_thoughts = split_thinking_response(full_plan) # NOTE: currently dropping thinking tokens from context self.worker_history.append(plan) logger.info("FULL PLAN:\n %s", full_plan) self.generator_agent.add_message(plan, role="assistant") # Use the grounding agent 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 = { "full_plan": full_plan, "executor_plan": plan, "plan_thoughts": plan_thoughts, "plan_code": plan_code, "reflection": reflection, "reflection_thoughts": reflection_thoughts, } self.turn_count += 1 self.screenshot_inputs.append(obs["screenshot"]) self.flush_messages() return executor_info, [exec_code]