chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,353 @@
|
||||
from functools import partial
|
||||
import logging
|
||||
import textwrap
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from gui_agents.s3.agents.grounding import ACI
|
||||
from gui_agents.s3.core.module import BaseModule
|
||||
from gui_agents.s3.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s3.utils.common_utils import (
|
||||
call_llm_safe,
|
||||
call_llm_formatted,
|
||||
parse_code_from_string,
|
||||
split_thinking_response,
|
||||
create_pyautogui_code,
|
||||
)
|
||||
from gui_agents.s3.utils.formatters import (
|
||||
SINGLE_ACTION_FORMATTER,
|
||||
CODE_VALID_FORMATTER,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class Worker(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
worker_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:
|
||||
worker_engine_params: Dict
|
||||
Parameters for the worker agent
|
||||
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__(worker_engine_params, platform)
|
||||
|
||||
self.temperature = worker_engine_params.get("temperature", 0.0)
|
||||
self.use_thinking = worker_engine_params.get("model", "") in [
|
||||
"claude-opus-4-20250514",
|
||||
"claude-sonnet-4-20250514",
|
||||
"claude-3-7-sonnet-20250219",
|
||||
"claude-sonnet-4-5-20250929",
|
||||
"claude-opus-4-5-20251101",
|
||||
]
|
||||
self.grounding_agent = grounding_agent
|
||||
self.max_trajectory_length = max_trajectory_length
|
||||
self.enable_reflection = enable_reflection
|
||||
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
if self.platform != "linux":
|
||||
skipped_actions = ["set_cell_values"]
|
||||
else:
|
||||
skipped_actions = []
|
||||
|
||||
# Hide code agent action entirely if no env/controller is available
|
||||
if not getattr(self.grounding_agent, "env", None) or not getattr(
|
||||
getattr(self.grounding_agent, "env", None), "controller", None
|
||||
):
|
||||
skipped_actions.append("call_code_agent")
|
||||
|
||||
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 = []
|
||||
|
||||
def flush_messages(self):
|
||||
"""Flush messages based on the model's context limits.
|
||||
|
||||
This method ensures that the agent's message history does not exceed the maximum trajectory length.
|
||||
|
||||
Side Effects:
|
||||
- Modifies the messages of generator, reflection, and bon_judge agents to fit within the context limits.
|
||||
"""
|
||||
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]:
|
||||
if agent is None:
|
||||
continue
|
||||
# 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_reflection(self, instruction: str, obs: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
Generate a reflection based on the current observation and instruction.
|
||||
|
||||
Args:
|
||||
instruction (str): The task instruction.
|
||||
obs (Dict): The current observation containing the screenshot.
|
||||
|
||||
Returns:
|
||||
Optional[str, str]: The generated reflection text and thoughts, if any (turn_count > 0).
|
||||
|
||||
Side Effects:
|
||||
- Updates reflection agent's history
|
||||
- Generates reflection response with API call
|
||||
"""
|
||||
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)
|
||||
logger.info("REFLECTION THOUGHTS: %s", reflection_thoughts)
|
||||
logger.info("REFLECTION: %s", reflection)
|
||||
return reflection, reflection_thoughts
|
||||
|
||||
def generate_next_action(self, instruction: str, obs: Dict) -> Tuple[Dict, List]:
|
||||
"""
|
||||
Predict the next action(s) based on the current observation.
|
||||
"""
|
||||
|
||||
self.grounding_agent.assign_screenshot(obs)
|
||||
self.grounding_agent.set_task_instruction(instruction)
|
||||
|
||||
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:
|
||||
prompt_with_instructions = self.generator_agent.system_prompt.replace(
|
||||
"TASK_DESCRIPTION", instruction
|
||||
)
|
||||
self.generator_agent.add_system_prompt(prompt_with_instructions)
|
||||
|
||||
# Get the per-step reflection
|
||||
reflection, reflection_thoughts = self._generate_reflection(instruction, obs)
|
||||
if reflection:
|
||||
generator_message += f"REFLECTION: You may use this reflection on the previous action and overall trajectory:\n{reflection}\n"
|
||||
|
||||
# Get the grounding agent's knowledge base buffer
|
||||
generator_message += (
|
||||
f"\nCurrent Text Buffer = [{','.join(self.grounding_agent.notes)}]\n"
|
||||
)
|
||||
|
||||
# Add code agent result from previous step if available (from full task or subtask execution)
|
||||
if (
|
||||
hasattr(self.grounding_agent, "last_code_agent_result")
|
||||
and self.grounding_agent.last_code_agent_result is not None
|
||||
):
|
||||
code_result = self.grounding_agent.last_code_agent_result
|
||||
generator_message += f"\nCODE AGENT RESULT:\n"
|
||||
generator_message += (
|
||||
f"Task/Subtask Instruction: {code_result['task_instruction']}\n"
|
||||
)
|
||||
generator_message += f"Steps Completed: {code_result['steps_executed']}\n"
|
||||
generator_message += f"Max Steps: {code_result['budget']}\n"
|
||||
generator_message += (
|
||||
f"Completion Reason: {code_result['completion_reason']}\n"
|
||||
)
|
||||
generator_message += f"Summary: {code_result['summary']}\n"
|
||||
if code_result["execution_history"]:
|
||||
generator_message += f"Execution History:\n"
|
||||
for i, step in enumerate(code_result["execution_history"]):
|
||||
action = step["action"]
|
||||
# Format code snippets with proper backticks
|
||||
if "```python" in action:
|
||||
# Extract Python code and format it
|
||||
code_start = action.find("```python") + 9
|
||||
code_end = action.find("```", code_start)
|
||||
if code_end != -1:
|
||||
python_code = action[code_start:code_end].strip()
|
||||
generator_message += (
|
||||
f"Step {i+1}: \n```python\n{python_code}\n```\n"
|
||||
)
|
||||
else:
|
||||
generator_message += f"Step {i+1}: \n{action}\n"
|
||||
elif "```bash" in action:
|
||||
# Extract Bash code and format it
|
||||
code_start = action.find("```bash") + 7
|
||||
code_end = action.find("```", code_start)
|
||||
if code_end != -1:
|
||||
bash_code = action[code_start:code_end].strip()
|
||||
generator_message += (
|
||||
f"Step {i+1}: \n```bash\n{bash_code}\n```\n"
|
||||
)
|
||||
else:
|
||||
generator_message += f"Step {i+1}: \n{action}\n"
|
||||
else:
|
||||
generator_message += f"Step {i+1}: \n{action}\n"
|
||||
generator_message += "\n"
|
||||
|
||||
# Log the code agent result section for debugging (truncated execution history)
|
||||
log_message = f"\nCODE AGENT RESULT:\n"
|
||||
log_message += (
|
||||
f"Task/Subtask Instruction: {code_result['task_instruction']}\n"
|
||||
)
|
||||
log_message += f"Steps Completed: {code_result['steps_executed']}\n"
|
||||
log_message += f"Max Steps: {code_result['budget']}\n"
|
||||
log_message += f"Completion Reason: {code_result['completion_reason']}\n"
|
||||
log_message += f"Summary: {code_result['summary']}\n"
|
||||
if code_result["execution_history"]:
|
||||
log_message += f"Execution History (truncated):\n"
|
||||
# Only log first 3 steps and last 2 steps to keep logs manageable
|
||||
total_steps = len(code_result["execution_history"])
|
||||
for i, step in enumerate(code_result["execution_history"]):
|
||||
if i < 3 or i >= total_steps - 2: # First 3 and last 2 steps
|
||||
action = step["action"]
|
||||
if "```python" in action:
|
||||
code_start = action.find("```python") + 9
|
||||
code_end = action.find("```", code_start)
|
||||
if code_end != -1:
|
||||
python_code = action[code_start:code_end].strip()
|
||||
log_message += (
|
||||
f"Step {i+1}: ```python\n{python_code}\n```\n"
|
||||
)
|
||||
else:
|
||||
log_message += f"Step {i+1}: {action}\n"
|
||||
elif "```bash" in action:
|
||||
code_start = action.find("```bash") + 7
|
||||
code_end = action.find("```", code_start)
|
||||
if code_end != -1:
|
||||
bash_code = action[code_start:code_end].strip()
|
||||
log_message += (
|
||||
f"Step {i+1}: ```bash\n{bash_code}\n```\n"
|
||||
)
|
||||
else:
|
||||
log_message += f"Step {i+1}: {action}\n"
|
||||
else:
|
||||
log_message += f"Step {i+1}: {action}\n"
|
||||
elif i == 3 and total_steps > 5:
|
||||
log_message += f"... (truncated {total_steps - 5} steps) ...\n"
|
||||
|
||||
logger.info(
|
||||
f"WORKER_CODE_AGENT_RESULT_SECTION - Step {self.turn_count + 1}: Code agent result added to generator message:\n{log_message}"
|
||||
)
|
||||
|
||||
# Reset the code agent result after adding it to context
|
||||
self.grounding_agent.last_code_agent_result = None
|
||||
|
||||
# Finalize the generator message
|
||||
self.generator_agent.add_message(
|
||||
generator_message, image_content=obs["screenshot"], role="user"
|
||||
)
|
||||
|
||||
# Generate the plan and next action
|
||||
format_checkers = [
|
||||
SINGLE_ACTION_FORMATTER,
|
||||
partial(CODE_VALID_FORMATTER, self.grounding_agent, obs),
|
||||
]
|
||||
plan = call_llm_formatted(
|
||||
self.generator_agent,
|
||||
format_checkers,
|
||||
temperature=self.temperature,
|
||||
use_thinking=self.use_thinking,
|
||||
)
|
||||
self.worker_history.append(plan)
|
||||
self.generator_agent.add_message(plan, role="assistant")
|
||||
logger.info("PLAN:\n %s", plan)
|
||||
|
||||
# Extract the next action from the plan
|
||||
plan_code = parse_code_from_string(plan)
|
||||
try:
|
||||
assert plan_code, "Plan code should not be empty"
|
||||
exec_code = create_pyautogui_code(self.grounding_agent, plan_code, obs)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Could not evaluate the following plan code:\n{plan_code}\nError: {e}"
|
||||
)
|
||||
exec_code = self.grounding_agent.wait(
|
||||
1.333
|
||||
) # Skip a turn if the code cannot be evaluated
|
||||
|
||||
executor_info = {
|
||||
"plan": plan,
|
||||
"plan_code": plan_code,
|
||||
"exec_code": exec_code,
|
||||
"reflection": reflection,
|
||||
"reflection_thoughts": reflection_thoughts,
|
||||
"code_agent_output": (
|
||||
self.grounding_agent.last_code_agent_result
|
||||
if hasattr(self.grounding_agent, "last_code_agent_result")
|
||||
and self.grounding_agent.last_code_agent_result is not None
|
||||
else None
|
||||
),
|
||||
}
|
||||
self.turn_count += 1
|
||||
self.screenshot_inputs.append(obs["screenshot"])
|
||||
self.flush_messages()
|
||||
return executor_info, [exec_code]
|
||||
Reference in New Issue
Block a user