wehub-resource-sync
205 lines
7.8 KiB
Python
205 lines
7.8 KiB
Python
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]
|