chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,94 @@
|
||||
import logging
|
||||
import platform
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from gui_agents.s3.agents.grounding import ACI
|
||||
from gui_agents.s3.agents.worker import Worker
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class UIAgent:
|
||||
"""Base class for UI automation agents"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
worker_engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
platform: str = platform.system().lower(),
|
||||
):
|
||||
"""Initialize UIAgent
|
||||
|
||||
Args:
|
||||
worker_engine_params: Configuration parameters for the worker LLM agent
|
||||
grounding_agent: Instance of ACI class for UI interaction
|
||||
platform: Operating system platform (macos, linux, windows)
|
||||
"""
|
||||
self.worker_engine_params = worker_engine_params
|
||||
self.grounding_agent = grounding_agent
|
||||
self.platform = platform
|
||||
|
||||
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
|
||||
|
||||
|
||||
class AgentS3(UIAgent):
|
||||
"""Agent that uses no hierarchy for less inference time"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
worker_engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
platform: str = platform.system().lower(),
|
||||
max_trajectory_length: int = 8,
|
||||
enable_reflection: bool = True,
|
||||
):
|
||||
"""Initialize a minimalist AgentS2 without hierarchy
|
||||
|
||||
Args:
|
||||
worker_engine_params: Configuration parameters for the worker agent.
|
||||
grounding_agent: Instance of ACI class for UI interaction
|
||||
platform: Operating system platform (darwin, linux, windows)
|
||||
max_trajectory_length: Maximum number of image turns to keep
|
||||
enable_reflection: Creates a reflection agent to assist the worker agent
|
||||
"""
|
||||
|
||||
super().__init__(worker_engine_params, grounding_agent, platform)
|
||||
self.max_trajectory_length = max_trajectory_length
|
||||
self.enable_reflection = enable_reflection
|
||||
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset agent state and initialize components"""
|
||||
self.executor = Worker(
|
||||
worker_engine_params=self.worker_engine_params,
|
||||
grounding_agent=self.grounding_agent,
|
||||
platform=self.platform,
|
||||
max_trajectory_length=self.max_trajectory_length,
|
||||
enable_reflection=self.enable_reflection,
|
||||
)
|
||||
|
||||
def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
|
||||
# Initialize the three info dictionaries
|
||||
executor_info, actions = self.executor.generate_next_action(
|
||||
instruction=instruction, obs=observation
|
||||
)
|
||||
|
||||
# concatenate the three info dictionaries
|
||||
info = {**{k: v for d in [executor_info or {}] for k, v in d.items()}}
|
||||
|
||||
return info, actions
|
||||
@@ -0,0 +1,333 @@
|
||||
import logging
|
||||
from typing import Dict, List, Tuple, Optional
|
||||
|
||||
from gui_agents.s3.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s3.utils.common_utils import call_llm_safe, split_thinking_response
|
||||
from gui_agents.s3.core.mllm import LMMAgent
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
def extract_code_block(action: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""Extract code and determine type from action string."""
|
||||
if "```python" in action:
|
||||
code_type = "python"
|
||||
code = action.split("```python")[1].split("```")[0].strip()
|
||||
elif "```bash" in action:
|
||||
code_type = "bash"
|
||||
code = action.split("```bash")[1].split("```")[0].strip()
|
||||
elif "```" in action:
|
||||
code_type = None
|
||||
code = action.split("```")[1].split("```")[0].strip()
|
||||
else:
|
||||
code_type = None
|
||||
code = None
|
||||
|
||||
logger.debug(
|
||||
f"Extracted code block: type={code_type}, length={len(code) if code else 0}"
|
||||
)
|
||||
return code_type, code
|
||||
|
||||
|
||||
def execute_code(code_type: str, code: str, env_controller) -> Dict:
|
||||
"""Execute code based on its type."""
|
||||
# Log the full code being executed (untruncated)
|
||||
logger.info(f"CODING_AGENT_CODE_EXECUTION - Type: {code_type}\nCode:\n{code}")
|
||||
|
||||
try:
|
||||
if code_type == "bash":
|
||||
result = env_controller.run_bash_script(code, timeout=30)
|
||||
elif code_type == "python":
|
||||
result = env_controller.run_python_script(code)
|
||||
else:
|
||||
result = {"status": "error", "error": f"Unknown code type: {code_type}"}
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing {code_type} code: {e}")
|
||||
return {"status": "error", "error": str(e)}
|
||||
|
||||
|
||||
def format_result(result: Dict, step_count: int) -> str:
|
||||
"""Format execution result into context string."""
|
||||
if not result:
|
||||
logger.warning(f"Step {step_count + 1}: No result returned from execution")
|
||||
return f"""
|
||||
Step {step_count + 1} Error:
|
||||
Error: No result returned from execution
|
||||
"""
|
||||
|
||||
status = result.get("status", "unknown")
|
||||
return_code = result.get("returncode", result.get("return_code", -1))
|
||||
|
||||
# Handle different response structures for bash vs python
|
||||
if "returncode" in result:
|
||||
# Bash script response
|
||||
output = result.get("output", "") # Contains both stdout and stderr merged
|
||||
error = result.get("error", "") # Always empty for bash
|
||||
else:
|
||||
# Python script response
|
||||
output = result.get("output", "") # stdout only
|
||||
error = result.get("error", "") # stderr only
|
||||
|
||||
logger.debug(f"Step {step_count + 1}: Status={status}, Return Code={return_code}")
|
||||
|
||||
# Format with better structure for multi-line outputs
|
||||
result_text = f"Step {step_count + 1} Result:\n"
|
||||
result_text += f"Status: {status}\n"
|
||||
result_text += f"Return Code: {return_code}\n"
|
||||
|
||||
if output:
|
||||
result_text += f"Output:\n{output}\n"
|
||||
|
||||
if error:
|
||||
result_text += f"Error:\n{error}\n"
|
||||
|
||||
return result_text
|
||||
|
||||
|
||||
class CodeAgent:
|
||||
"""A dedicated agent for executing code with a budget of steps."""
|
||||
|
||||
def __init__(self, engine_params: Dict, budget: int = 20):
|
||||
"""Initialize the CodeAgent."""
|
||||
if not engine_params:
|
||||
raise ValueError("engine_params cannot be None or empty")
|
||||
|
||||
self.engine_params = engine_params
|
||||
self.budget = budget
|
||||
self.agent = None
|
||||
|
||||
logger.info(f"CodeAgent initialized with budget={budget}")
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
"""Reset the code agent state."""
|
||||
logger.debug("Resetting CodeAgent state")
|
||||
self.agent = LMMAgent(
|
||||
engine_params=self.engine_params,
|
||||
system_prompt=PROCEDURAL_MEMORY.CODE_AGENT_PROMPT,
|
||||
)
|
||||
|
||||
def execute(self, task_instruction: str, screenshot: str, env_controller) -> Dict:
|
||||
"""Execute code for the given task with a budget of steps."""
|
||||
if env_controller is None:
|
||||
raise ValueError("env_controller is required for code execution")
|
||||
|
||||
print(f"\n🚀 STARTING CODE EXECUTION")
|
||||
print("=" * 60)
|
||||
print(f"Task: {task_instruction}")
|
||||
print(f"Budget: {self.budget} steps")
|
||||
print("=" * 60)
|
||||
|
||||
logger.info(f"Starting code execution for task: {task_instruction}")
|
||||
logger.info(f"Budget: {self.budget} steps")
|
||||
|
||||
self.reset()
|
||||
|
||||
# Add initial task instruction and screenshot context as user message
|
||||
context = (
|
||||
f"Task: {task_instruction}\n\nCurrent screenshot is provided for context."
|
||||
)
|
||||
self.agent.add_message(context, image_content=screenshot, role="user")
|
||||
|
||||
step_count = 0
|
||||
execution_history = []
|
||||
|
||||
while step_count < self.budget:
|
||||
logger.info(f"Step {step_count + 1}/{self.budget}")
|
||||
|
||||
# Get assistant response (thoughts and code)
|
||||
response = call_llm_safe(self.agent, temperature=1)
|
||||
|
||||
# Print to terminal for immediate visibility
|
||||
print(f"\n🤖 CODING AGENT RESPONSE - Step {step_count + 1}/{self.budget}")
|
||||
print("=" * 60)
|
||||
print(response)
|
||||
print("=" * 60)
|
||||
|
||||
# Log the latest message from the coding agent (untruncated)
|
||||
logger.info(
|
||||
f"CODING_AGENT_LATEST_MESSAGE - Step {step_count + 1}:\n{response}"
|
||||
)
|
||||
|
||||
# Check if response is None or empty
|
||||
if not response or response.strip() == "":
|
||||
error_msg = f"Step {step_count + 1}: LLM returned empty response"
|
||||
logger.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
# Parse the response to extract action
|
||||
action, thoughts = split_thinking_response(response)
|
||||
|
||||
execution_history.append(
|
||||
{"step": step_count + 1, "action": action, "thoughts": thoughts}
|
||||
)
|
||||
|
||||
# Check for completion signals
|
||||
action_upper = action.upper().strip()
|
||||
if action_upper == "DONE":
|
||||
print(f"\n✅ TASK COMPLETED - Step {step_count + 1}")
|
||||
print("=" * 60)
|
||||
print("Agent signaled task completion")
|
||||
print("=" * 60)
|
||||
logger.info(f"Step {step_count + 1}: Task completed successfully")
|
||||
completion_reason = "DONE"
|
||||
break
|
||||
elif action_upper == "FAIL":
|
||||
print(f"\n❌ TASK FAILED - Step {step_count + 1}")
|
||||
print("=" * 60)
|
||||
print("Agent signaled task failure")
|
||||
print("=" * 60)
|
||||
logger.info(f"Step {step_count + 1}: Task failed by agent request")
|
||||
completion_reason = "FAIL"
|
||||
break
|
||||
|
||||
# Extract and execute code
|
||||
code_type, code = extract_code_block(action)
|
||||
|
||||
if code:
|
||||
result = execute_code(code_type, code, env_controller)
|
||||
# Prepare formatted output and error for logging
|
||||
output = result.get("output", "")
|
||||
error = result.get("error", "")
|
||||
message = result.get("message", "")
|
||||
status = result.get("status", "")
|
||||
|
||||
# Print execution result to terminal for immediate visibility
|
||||
print(f"\n⚡ CODE EXECUTION RESULT - Step {step_count + 1}")
|
||||
print("-" * 50)
|
||||
print(f"Status: {status}")
|
||||
if output:
|
||||
print(f"Output:\n{output}")
|
||||
if error:
|
||||
print(f"Error:\n{error}")
|
||||
if message and not output and not error:
|
||||
print(f"Message:\n{message}")
|
||||
print("-" * 50)
|
||||
|
||||
log_lines = [
|
||||
f"CODING_AGENT_EXECUTION_RESULT - Step {step_count + 1}:",
|
||||
f"Status: {status}" if status else None,
|
||||
]
|
||||
|
||||
if output:
|
||||
log_lines.append(
|
||||
"Output:\n" + ("-" * 40) + f"\n{output}\n" + ("-" * 40)
|
||||
)
|
||||
if error:
|
||||
log_lines.append(
|
||||
"Error:\n" + ("!" * 40) + f"\n{error}\n" + ("!" * 40)
|
||||
)
|
||||
if message and not output and not error:
|
||||
log_lines.append(
|
||||
"Message:\n" + ("-" * 40) + f"\n{message}\n" + ("-" * 40)
|
||||
)
|
||||
|
||||
# Remove None entries and join
|
||||
formatted_log = "\n".join([line for line in log_lines if line])
|
||||
logger.info(formatted_log)
|
||||
else:
|
||||
print(f"\n⚠️ NO CODE BLOCK FOUND - Step {step_count + 1}")
|
||||
print("-" * 50)
|
||||
print("Action did not contain executable code")
|
||||
print("-" * 50)
|
||||
|
||||
logger.warning(f"Step {step_count + 1}: No code block found in action")
|
||||
result = {"status": "skipped", "message": "No code block found"}
|
||||
logger.info(
|
||||
f"CODING_AGENT_EXECUTION_RESULT - Step {step_count + 1}:\n"
|
||||
f"Status: skipped\n"
|
||||
f"Message:\n{'-' * 40}\n{result['message']}\n{'-' * 40}"
|
||||
)
|
||||
# Add assistant's thoughts and code to message history
|
||||
self.agent.add_message(response, role="assistant")
|
||||
|
||||
# Process result and add formatted environment results as user message
|
||||
result_context = format_result(result, step_count)
|
||||
self.agent.add_message(result_context, role="user")
|
||||
|
||||
step_count += 1
|
||||
|
||||
# Handle budget exhaustion
|
||||
if "completion_reason" not in locals():
|
||||
print(f"\n⏰ BUDGET EXHAUSTED - {step_count} steps completed")
|
||||
print("=" * 60)
|
||||
print(f"Maximum budget of {self.budget} steps reached")
|
||||
print("=" * 60)
|
||||
logger.info(f"Budget exhausted after {step_count} steps")
|
||||
completion_reason = f"BUDGET_EXHAUSTED_AFTER_{step_count}_STEPS"
|
||||
|
||||
# Generate final summary
|
||||
logger.info("Generating execution summary")
|
||||
summary = self._generate_summary(execution_history, task_instruction)
|
||||
|
||||
result = {
|
||||
"task_instruction": task_instruction,
|
||||
"completion_reason": completion_reason,
|
||||
"summary": summary,
|
||||
"execution_history": execution_history,
|
||||
"steps_executed": step_count,
|
||||
"budget": self.budget,
|
||||
}
|
||||
|
||||
logger.info(f"Code execution completed: steps={step_count}")
|
||||
return result
|
||||
|
||||
def _generate_summary(
|
||||
self, execution_history: List[Dict], task_instruction: str
|
||||
) -> str:
|
||||
"""Generate summary of code execution session."""
|
||||
if not execution_history:
|
||||
logger.info("No execution history to summarize")
|
||||
return "No actions were executed."
|
||||
|
||||
logger.info(f"Generated summary for {len(execution_history)} steps")
|
||||
|
||||
# Build detailed execution context for summary agent
|
||||
execution_context = f"Task: {task_instruction}\n\nExecution Steps:\n"
|
||||
|
||||
for step in execution_history:
|
||||
step_num = step["step"]
|
||||
thoughts = step.get("thoughts", "")
|
||||
action = step.get("action", "")
|
||||
|
||||
execution_context += f"\nStep {step_num}:\n"
|
||||
if thoughts:
|
||||
execution_context += f"Thoughts: {thoughts}\n"
|
||||
execution_context += f"Code: {action}\n"
|
||||
|
||||
# Create summary prompt with same context as coding agent
|
||||
summary_prompt = f"""
|
||||
{execution_context}
|
||||
|
||||
Please provide a concise summary of the code execution session. Focus on:
|
||||
|
||||
1. The code logic implemented at each step
|
||||
2. The outputs and results produced by each code execution
|
||||
3. The progression of the solution approach
|
||||
|
||||
Do not make judgments about success or failure. Simply describe what was attempted and what resulted.
|
||||
|
||||
Keep the summary under 150 words and use clear, factual language.
|
||||
"""
|
||||
|
||||
# Generate summary using LLM with dedicated summary system prompt
|
||||
try:
|
||||
summary_agent = LMMAgent(
|
||||
engine_params=self.engine_params,
|
||||
system_prompt=PROCEDURAL_MEMORY.CODE_SUMMARY_AGENT_PROMPT,
|
||||
)
|
||||
summary_agent.add_message(summary_prompt, role="user")
|
||||
summary = call_llm_safe(summary_agent, temperature=1)
|
||||
|
||||
if not summary or summary.strip() == "":
|
||||
summary = "Summary generation failed - no response from LLM"
|
||||
logger.warning("Summary generation failed - empty response from LLM")
|
||||
|
||||
except Exception as e:
|
||||
summary = f"Summary generation failed: {str(e)}"
|
||||
logger.error(f"Error generating summary: {e}")
|
||||
|
||||
return summary
|
||||
@@ -0,0 +1,667 @@
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
from pytesseract import Output
|
||||
|
||||
from gui_agents.s3.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s3.core.mllm import LMMAgent
|
||||
from gui_agents.s3.utils.common_utils import call_llm_safe
|
||||
from gui_agents.s3.agents.code_agent import CodeAgent
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class ACI:
|
||||
def __init__(self):
|
||||
self.notes: List[str] = []
|
||||
|
||||
|
||||
# Agent action decorator
|
||||
def agent_action(func):
|
||||
func.is_agent_action = True
|
||||
return func
|
||||
|
||||
|
||||
UBUNTU_APP_SETUP = f"""import subprocess;
|
||||
import difflib;
|
||||
import pyautogui;
|
||||
import time;
|
||||
pyautogui.press('escape');
|
||||
time.sleep(0.5);
|
||||
output = subprocess.check_output(['wmctrl', '-lx']);
|
||||
output = output.decode('utf-8').splitlines();
|
||||
window_titles = [line.split(None, 4)[2] for line in output];
|
||||
closest_matches = difflib.get_close_matches('APP_NAME', window_titles, n=1, cutoff=0.1);
|
||||
if closest_matches:
|
||||
closest_match = closest_matches[0];
|
||||
for line in output:
|
||||
if closest_match in line:
|
||||
window_id = line.split()[0]
|
||||
break;
|
||||
subprocess.run(['wmctrl', '-ia', window_id])
|
||||
subprocess.run(['wmctrl', '-ir', window_id, '-b', 'add,maximized_vert,maximized_horz'])
|
||||
"""
|
||||
|
||||
|
||||
SET_CELL_VALUES_CMD = """import uno
|
||||
import subprocess
|
||||
import unicodedata, json
|
||||
|
||||
def identify_document_type(component):
|
||||
if component.supportsService("com.sun.star.sheet.SpreadsheetDocument"):
|
||||
return "Calc"
|
||||
|
||||
if component.supportsService("com.sun.star.text.TextDocument"):
|
||||
return "Writer"
|
||||
|
||||
if component.supportsService("com.sun.star.sheet.PresentationDocument"):
|
||||
return "Impress"
|
||||
|
||||
return None
|
||||
|
||||
def _norm_name(s: str | None) -> str | None:
|
||||
if s is None:
|
||||
return None
|
||||
if "\\\\u" in s or "\\\\U" in s or "\\\\x" in s:
|
||||
try:
|
||||
# json.loads handles all the escape forms safely
|
||||
s = json.loads(f"{{s}}")
|
||||
except Exception:
|
||||
# fallback: best-effort
|
||||
try:
|
||||
s = s.encode("utf-8").decode("unicode_escape")
|
||||
except Exception:
|
||||
pass
|
||||
# Normalize (NFC works well across platforms)
|
||||
return unicodedata.normalize("NFC", s)
|
||||
|
||||
def cell_ref_to_indices(cell_ref):
|
||||
column_letters = ''.join(filter(str.isalpha, cell_ref))
|
||||
row_number = ''.join(filter(str.isdigit, cell_ref))
|
||||
|
||||
col = sum((ord(char.upper()) - ord('A') + 1) * (26**idx) for idx, char in enumerate(reversed(column_letters))) - 1
|
||||
row = int(row_number) - 1
|
||||
return col, row
|
||||
|
||||
def set_cell_values(new_cell_values: dict[str, str], app_name: str = "Untitled 1", sheet_name: str = "Sheet1"):
|
||||
app_name = _norm_name(app_name)
|
||||
sheet_name = _norm_name(sheet_name)
|
||||
|
||||
new_cell_values_idx = {{}}
|
||||
for k, v in new_cell_values.items():
|
||||
try:
|
||||
col, row = cell_ref_to_indices(k)
|
||||
except:
|
||||
col = row = None
|
||||
|
||||
if col is not None and row is not None:
|
||||
new_cell_values_idx[(col, row)] = v
|
||||
|
||||
# Clean up previous TCP connections.
|
||||
subprocess.run(
|
||||
'echo \"osworld-public-evaluation\" | sudo -S ss --kill --tcp state TIME-WAIT sport = :2002',
|
||||
shell=True,
|
||||
check=True,
|
||||
text=True,
|
||||
capture_output=True
|
||||
)
|
||||
|
||||
# Dynamically allow soffice to listen on port 2002.
|
||||
subprocess.run(
|
||||
[
|
||||
"soffice",
|
||||
"--accept=socket,host=localhost,port=2002;urp;StarOffice.Service"
|
||||
]
|
||||
)
|
||||
|
||||
local_context = uno.getComponentContext()
|
||||
resolver = local_context.ServiceManager.createInstanceWithContext(
|
||||
"com.sun.star.bridge.UnoUrlResolver", local_context
|
||||
)
|
||||
context = resolver.resolve(
|
||||
f"uno:socket,host=localhost,port=2002;urp;StarOffice.ComponentContext"
|
||||
)
|
||||
desktop = context.ServiceManager.createInstanceWithContext(
|
||||
"com.sun.star.frame.Desktop", context
|
||||
)
|
||||
|
||||
# Collect all LibreOffice-related opened windows.
|
||||
documents = []
|
||||
for i, component in enumerate(desktop.Components):
|
||||
title = component.Title
|
||||
doc_type = identify_document_type(component)
|
||||
documents.append((i, component, title, doc_type))
|
||||
|
||||
# Find the LibreOffice Calc app and the sheet of interest.
|
||||
spreadsheet = [doc for doc in documents if doc[3] == "Calc"]
|
||||
selected_spreadsheet = [doc for doc in spreadsheet if doc[2] == app_name]
|
||||
if spreadsheet:
|
||||
try:
|
||||
if selected_spreadsheet:
|
||||
spreadsheet = selected_spreadsheet[0][1]
|
||||
else:
|
||||
spreadsheet = spreadsheet[0][1]
|
||||
|
||||
sheet = spreadsheet.Sheets.getByName(sheet_name)
|
||||
except:
|
||||
raise ValueError(f"Could not find sheet {{sheet_name}} in {{app_name}}.")
|
||||
|
||||
for (col, row), value in new_cell_values_idx.items():
|
||||
cell = sheet.getCellByPosition(col, row)
|
||||
|
||||
# Set the cell value.
|
||||
if isinstance(value, (int, float)):
|
||||
cell.Value = value
|
||||
elif isinstance(value, str):
|
||||
if value.startswith("="):
|
||||
cell.Formula = value
|
||||
else:
|
||||
cell.String = value
|
||||
elif isinstance(value, bool):
|
||||
cell.Value = 1 if value else 0
|
||||
elif value is None:
|
||||
cell.clearContents(0)
|
||||
else:
|
||||
raise ValueError(f"Unsupported cell value type: {{type(value)}}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Could not find LibreOffice Calc app corresponding to {{app_name}}.")
|
||||
|
||||
set_cell_values(new_cell_values={cell_values}, app_name="{app_name}", sheet_name="{sheet_name}")
|
||||
"""
|
||||
|
||||
|
||||
# ACI primitives are parameterized by description, and coordinate generation uses a pretrained grounding model
|
||||
class OSWorldACI(ACI):
|
||||
def __init__(
|
||||
self,
|
||||
env,
|
||||
platform: str,
|
||||
engine_params_for_generation: Dict,
|
||||
engine_params_for_grounding: Dict,
|
||||
width: int = 1920,
|
||||
height: int = 1080,
|
||||
code_agent_budget: int = 20,
|
||||
code_agent_engine_params: Dict = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.env = env
|
||||
self.platform = (
|
||||
platform # Dictates how the switch_applications agent action works.
|
||||
)
|
||||
|
||||
# Configure scaling
|
||||
self.width = width
|
||||
self.height = height
|
||||
|
||||
# Maintain state for save_to_knowledge
|
||||
self.notes = []
|
||||
|
||||
# Screenshot used during ACI execution
|
||||
self.obs = None
|
||||
|
||||
# Configure the visual grounding model responsible for coordinate generation
|
||||
self.grounding_model = LMMAgent(engine_params_for_grounding)
|
||||
self.engine_params_for_grounding = engine_params_for_grounding
|
||||
|
||||
# Configure text grounding agent
|
||||
self.text_span_agent = LMMAgent(
|
||||
engine_params=engine_params_for_generation,
|
||||
system_prompt=PROCEDURAL_MEMORY.PHRASE_TO_WORD_COORDS_PROMPT,
|
||||
)
|
||||
|
||||
# Configure code agent
|
||||
code_agent_engine_params = (
|
||||
code_agent_engine_params or engine_params_for_generation
|
||||
)
|
||||
self.code_agent = CodeAgent(code_agent_engine_params, code_agent_budget)
|
||||
|
||||
# Store task instruction for code agent
|
||||
self.current_task_instruction = None
|
||||
self.last_code_agent_result = None
|
||||
|
||||
# Given the state and worker's referring expression, use the grounding model to generate (x,y)
|
||||
def generate_coords(self, ref_expr: str, obs: Dict) -> List[int]:
|
||||
|
||||
# Reset the grounding model state
|
||||
self.grounding_model.reset()
|
||||
|
||||
# Configure the context, UI-TARS demo does not use system prompt
|
||||
prompt = f"Query:{ref_expr}\nOutput only the coordinate of one point in your response.\n"
|
||||
self.grounding_model.add_message(
|
||||
text_content=prompt, image_content=obs["screenshot"], put_text_last=True
|
||||
)
|
||||
|
||||
# Generate and parse coordinates
|
||||
response = call_llm_safe(self.grounding_model)
|
||||
print("RAW GROUNDING MODEL RESPONSE:", response)
|
||||
numericals = re.findall(r"\d+", response)
|
||||
assert len(numericals) >= 2
|
||||
return [int(numericals[0]), int(numericals[1])]
|
||||
|
||||
# Calls pytesseract to generate word level bounding boxes for text grounding
|
||||
def get_ocr_elements(self, b64_image_data: str) -> Tuple[str, List]:
|
||||
image = Image.open(BytesIO(b64_image_data))
|
||||
image_data = pytesseract.image_to_data(image, output_type=Output.DICT)
|
||||
|
||||
# Clean text by removing leading and trailing spaces and non-alphabetical characters, but keeping punctuation
|
||||
for i, word in enumerate(image_data["text"]):
|
||||
image_data["text"][i] = re.sub(
|
||||
r"^[^a-zA-Z\s.,!?;:\-\+]+|[^a-zA-Z\s.,!?;:\-\+]+$", "", word
|
||||
)
|
||||
|
||||
ocr_elements = []
|
||||
ocr_table = "Text Table:\nWord id\tText\n"
|
||||
# Obtain the <id, text, group number, word number> for each valid element
|
||||
grouping_map = defaultdict(list)
|
||||
ocr_id = 0
|
||||
for i in range(len(image_data["text"])):
|
||||
block_num = image_data["block_num"][i]
|
||||
if image_data["text"][i]:
|
||||
grouping_map[block_num].append(image_data["text"][i])
|
||||
ocr_table += f"{ocr_id}\t{image_data['text'][i]}\n"
|
||||
ocr_elements.append(
|
||||
{
|
||||
"id": ocr_id,
|
||||
"text": image_data["text"][i],
|
||||
"group_num": block_num,
|
||||
"word_num": len(grouping_map[block_num]),
|
||||
"left": image_data["left"][i],
|
||||
"top": image_data["top"][i],
|
||||
"width": image_data["width"][i],
|
||||
"height": image_data["height"][i],
|
||||
}
|
||||
)
|
||||
ocr_id += 1
|
||||
|
||||
return ocr_table, ocr_elements
|
||||
|
||||
# Given the state and worker's text phrase, generate the coords of the first/last word in the phrase
|
||||
def generate_text_coords(
|
||||
self, phrase: str, obs: Dict, alignment: str = ""
|
||||
) -> List[int]:
|
||||
|
||||
ocr_table, ocr_elements = self.get_ocr_elements(obs["screenshot"])
|
||||
|
||||
alignment_prompt = ""
|
||||
if alignment == "start":
|
||||
alignment_prompt = "**Important**: Output the word id of the FIRST word in the provided phrase.\n"
|
||||
elif alignment == "end":
|
||||
alignment_prompt = "**Important**: Output the word id of the LAST word in the provided phrase.\n"
|
||||
|
||||
# Load LLM prompt
|
||||
self.text_span_agent.reset()
|
||||
self.text_span_agent.add_message(
|
||||
alignment_prompt + "Phrase: " + phrase + "\n" + ocr_table, role="user"
|
||||
)
|
||||
self.text_span_agent.add_message(
|
||||
"Screenshot:\n", image_content=obs["screenshot"], role="user"
|
||||
)
|
||||
|
||||
# Obtain the target element
|
||||
response = call_llm_safe(self.text_span_agent)
|
||||
print("TEXT SPAN AGENT RESPONSE:", response)
|
||||
numericals = re.findall(r"\d+", response)
|
||||
if len(numericals) > 0:
|
||||
text_id = int(numericals[-1])
|
||||
else:
|
||||
text_id = 0
|
||||
elem = ocr_elements[text_id]
|
||||
|
||||
# Compute the element coordinates
|
||||
if alignment == "start":
|
||||
coords = [elem["left"], elem["top"] + (elem["height"] // 2)]
|
||||
elif alignment == "end":
|
||||
coords = [elem["left"] + elem["width"], elem["top"] + (elem["height"] // 2)]
|
||||
else:
|
||||
coords = [
|
||||
elem["left"] + (elem["width"] // 2),
|
||||
elem["top"] + (elem["height"] // 2),
|
||||
]
|
||||
return coords
|
||||
|
||||
def assign_screenshot(self, obs: Dict):
|
||||
self.obs = obs
|
||||
|
||||
def set_task_instruction(self, task_instruction: str):
|
||||
"""Set the current task instruction for the code agent."""
|
||||
self.current_task_instruction = task_instruction
|
||||
|
||||
# Resize from grounding model dim into OSWorld dim (1920 * 1080)
|
||||
def resize_coordinates(self, coordinates: List[int]) -> List[int]:
|
||||
grounding_width = self.engine_params_for_grounding["grounding_width"]
|
||||
grounding_height = self.engine_params_for_grounding["grounding_height"]
|
||||
|
||||
return [
|
||||
round(coordinates[0] * self.width / grounding_width),
|
||||
round(coordinates[1] * self.height / grounding_height),
|
||||
]
|
||||
|
||||
@agent_action
|
||||
def click(
|
||||
self,
|
||||
element_description: str,
|
||||
num_clicks: int = 1,
|
||||
button_type: str = "left",
|
||||
hold_keys: List = [],
|
||||
):
|
||||
"""Click on the element
|
||||
Args:
|
||||
element_description:str, a detailed descriptions of which element to click on. This description should be at least a full sentence.
|
||||
num_clicks:int, number of times to click the element
|
||||
button_type:str, which mouse button to press can be "left", "middle", or "right"
|
||||
hold_keys:List, list of keys to hold while clicking
|
||||
"""
|
||||
coords1 = self.generate_coords(element_description, self.obs)
|
||||
x, y = self.resize_coordinates(coords1)
|
||||
command = "import pyautogui; "
|
||||
|
||||
# TODO: specified duration?
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyDown({repr(k)}); "
|
||||
command += f"""import pyautogui; pyautogui.click({x}, {y}, clicks={num_clicks}, button={repr(button_type)}); """
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyUp({repr(k)}); "
|
||||
# Return pyautoguicode to click on the element
|
||||
return command
|
||||
|
||||
@agent_action
|
||||
def switch_applications(self, app_code):
|
||||
"""Switch to a different application that is already open
|
||||
Args:
|
||||
app_code:str the code name of the application to switch to from the provided list of open applications
|
||||
"""
|
||||
if self.platform == "darwin":
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
|
||||
elif self.platform == "linux":
|
||||
return UBUNTU_APP_SETUP.replace("APP_NAME", app_code)
|
||||
elif self.platform == "windows":
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('win', 'd', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
), f"Unsupported platform: {self.platform}. Supported platforms are: darwin, linux, windows."
|
||||
|
||||
@agent_action
|
||||
def open(self, app_or_filename: str):
|
||||
"""Open any application or file with name app_or_filename. Use this action to open applications or files on the desktop, do not open manually.
|
||||
Args:
|
||||
app_or_filename:str, the name of the application or filename to open
|
||||
"""
|
||||
if self.platform == "linux":
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('win'); time.sleep(0.5); pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); pyautogui.hotkey('enter'); time.sleep(0.5)"
|
||||
elif self.platform == "darwin":
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_or_filename)}); pyautogui.press('enter'); time.sleep(1.0)"
|
||||
elif self.platform == "windows":
|
||||
return (
|
||||
"import pyautogui; import time; "
|
||||
"pyautogui.hotkey('win'); time.sleep(0.5); "
|
||||
f"pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); "
|
||||
"pyautogui.press('enter'); time.sleep(0.5)"
|
||||
)
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
), f"Unsupported platform: {self.platform}. Supported platforms are: darwin, linux, windows."
|
||||
|
||||
@agent_action
|
||||
def type(
|
||||
self,
|
||||
element_description: Optional[str] = None,
|
||||
text: str = "",
|
||||
overwrite: bool = False,
|
||||
enter: bool = False,
|
||||
):
|
||||
"""Type text/unicode into a specific element
|
||||
Args:
|
||||
element_description:str, a detailed description of which element to enter text in. This description should be at least a full sentence.
|
||||
text:str, the text to type
|
||||
overwrite:bool, Assign it to True if the text should overwrite the existing text, otherwise assign it to False. Using this argument clears all text in an element.
|
||||
enter:bool, Assign it to True if the enter key should be pressed after typing the text, otherwise assign it to False.
|
||||
"""
|
||||
command = "import pyautogui; "
|
||||
command += (
|
||||
"\ntry:\n"
|
||||
" import pyperclip\n"
|
||||
"except ImportError:\n"
|
||||
" import subprocess\n"
|
||||
" subprocess.run('echo \"osworld-public-evaluation\" | sudo -S apt-get install -y xclip xsel', shell=True, check=True)\n"
|
||||
" subprocess.check_call([subprocess.sys.executable, '-m', 'pip', 'install', 'pyperclip'])\n"
|
||||
" import pyperclip\n\n"
|
||||
)
|
||||
|
||||
if element_description is not None:
|
||||
coords1 = self.generate_coords(element_description, self.obs)
|
||||
x, y = self.resize_coordinates(coords1)
|
||||
command += f"pyautogui.click({x}, {y}); "
|
||||
|
||||
if overwrite:
|
||||
command += (
|
||||
f"pyautogui.hotkey({repr('command' if self.platform == 'darwin' else 'ctrl')}, 'a'); "
|
||||
"pyautogui.press('backspace'); "
|
||||
)
|
||||
|
||||
# Check if text contains Unicode characters that pyautogui.write() can't handle
|
||||
has_unicode = any(ord(char) > 127 for char in text)
|
||||
|
||||
if has_unicode:
|
||||
# Use clipboard method for Unicode characters
|
||||
command += f"pyperclip.copy({repr(text)}); "
|
||||
command += f"pyautogui.hotkey({repr('command' if self.platform == 'darwin' else 'ctrl')}, 'v'); "
|
||||
else:
|
||||
# Use regular pyautogui.write() for ASCII text
|
||||
command += f"pyautogui.write({repr(text)}); "
|
||||
|
||||
if enter:
|
||||
command += "pyautogui.press('enter'); "
|
||||
return command
|
||||
|
||||
@agent_action
|
||||
def save_to_knowledge(self, text: List[str]):
|
||||
"""Save facts, elements, texts, etc. to a long-term knowledge bank for reuse during this task. Can be used for copy-pasting text, saving elements, etc.
|
||||
Args:
|
||||
text:List[str] the text to save to the knowledge
|
||||
"""
|
||||
self.notes.extend(text)
|
||||
return """WAIT"""
|
||||
|
||||
@agent_action
|
||||
def drag_and_drop(
|
||||
self, starting_description: str, ending_description: str, hold_keys: List = []
|
||||
):
|
||||
"""Drag from the starting description to the ending description
|
||||
Args:
|
||||
starting_description:str, a very detailed description of where to start the drag action. This description should be at least a full sentence.
|
||||
ending_description:str, a very detailed description of where to end the drag action. This description should be at least a full sentence.
|
||||
hold_keys:List list of keys to hold while dragging
|
||||
"""
|
||||
coords1 = self.generate_coords(starting_description, self.obs)
|
||||
coords2 = self.generate_coords(ending_description, self.obs)
|
||||
x1, y1 = self.resize_coordinates(coords1)
|
||||
x2, y2 = self.resize_coordinates(coords2)
|
||||
|
||||
command = "import pyautogui; "
|
||||
|
||||
command += f"pyautogui.moveTo({x1}, {y1}); "
|
||||
# TODO: specified duration?
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyDown({repr(k)}); "
|
||||
command += f"pyautogui.dragTo({x2}, {y2}, duration=1., button='left'); pyautogui.mouseUp(); "
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyUp({repr(k)}); "
|
||||
|
||||
# Return pyautoguicode to drag and drop the elements
|
||||
|
||||
return command
|
||||
|
||||
@agent_action
|
||||
def highlight_text_span(
|
||||
self, starting_phrase: str, ending_phrase: str, button: str = "left"
|
||||
):
|
||||
"""Highlight a text span between a provided starting phrase and ending phrase. Use this to highlight words, lines, and paragraphs.
|
||||
Args:
|
||||
starting_phrase:str, the phrase that denotes the start of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
|
||||
ending_phrase:str, the phrase that denotes the end of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
|
||||
button:str, the button to use to highlight the text span. Defaults to "left". Can be "left", "right", or "middle".
|
||||
"""
|
||||
coords1 = self.generate_text_coords(
|
||||
starting_phrase, self.obs, alignment="start"
|
||||
)
|
||||
coords2 = self.generate_text_coords(ending_phrase, self.obs, alignment="end")
|
||||
x1, y1 = coords1
|
||||
x2, y2 = coords2
|
||||
|
||||
command = "import pyautogui; "
|
||||
command += f"pyautogui.moveTo({x1}, {y1}); "
|
||||
command += f"pyautogui.dragTo({x2}, {y2}, duration=1., button='{button}'); pyautogui.mouseUp(); "
|
||||
|
||||
# Return pyautoguicode to drag and drop the elements
|
||||
return command
|
||||
|
||||
@agent_action
|
||||
def set_cell_values(
|
||||
self, cell_values: Dict[str, Any], app_name: str, sheet_name: str
|
||||
):
|
||||
"""Use this to set individual cell values in a spreadsheet. For example, setting A2 to "hello" would be done by passing {"A2": "hello"} as cell_values. The sheet must be opened before this command can be used.
|
||||
Args:
|
||||
cell_values: Dict[str, Any], A dictionary of cell values to set in the spreadsheet. The keys are the cell coordinates in the format "A1", "B2", etc.
|
||||
Supported value types include: float, int, string, bool, formulas.
|
||||
app_name: str, The name of the spreadsheet application. For example, "Some_sheet.xlsx".
|
||||
sheet_name: str, The name of the sheet in the spreadsheet. For example, "Sheet1".
|
||||
"""
|
||||
return SET_CELL_VALUES_CMD.format(
|
||||
cell_values=cell_values, app_name=app_name, sheet_name=sheet_name
|
||||
)
|
||||
|
||||
@agent_action
|
||||
def call_code_agent(self, task: str = None):
|
||||
"""Call the code agent to execute code for tasks or subtasks that can be completed solely with coding.
|
||||
|
||||
Args:
|
||||
task: str, the task or subtask to execute. If None, uses the current full task instruction.
|
||||
|
||||
**🚨 CRITICAL GUIDELINES:**
|
||||
- **ONLY pass a task parameter for SPECIFIC subtasks** (e.g., "Calculate sum of column B", "Filter data by date")
|
||||
- **NEVER pass a task parameter for full tasks** - let it default to the original task instruction
|
||||
- **NEVER rephrase or modify the original task** - this prevents hallucination corruption
|
||||
- **If unsure, omit the task parameter entirely** to use the original task instruction
|
||||
|
||||
Use this for tasks that can be fully accomplished through code execution, particularly for:
|
||||
- Spreadsheet applications (LibreOffice Calc, Excel): data processing, filtering, sorting, calculations, formulas, data analysis
|
||||
- Document editors (LibreOffice Writer, Word): text processing, content editing, formatting, document manipulation
|
||||
- Code editors (VS Code, text editors): code editing, file processing, text manipulation, configuration
|
||||
- Data analysis tools: statistical analysis, data transformation, reporting
|
||||
- File management: bulk operations, file processing, content extraction
|
||||
- System utilities: configuration, setup, automation
|
||||
"""
|
||||
logger.info("=" * 50)
|
||||
logger.info("GROUNDING AGENT: Calling Code Agent")
|
||||
logger.info("=" * 50)
|
||||
|
||||
# **CRITICAL**: Only use provided task for specific subtasks, otherwise use original task instruction
|
||||
if task is not None:
|
||||
# This is a subtask - use the provided task
|
||||
task_to_execute = task
|
||||
logger.info(f"Executing SUBTASK: {task_to_execute}")
|
||||
else:
|
||||
# This is a full task - use the original task instruction to prevent hallucination
|
||||
task_to_execute = self.current_task_instruction
|
||||
logger.info(f"Executing FULL TASK: {task_to_execute}")
|
||||
|
||||
if task_to_execute:
|
||||
print("obs keys: ", self.obs.keys())
|
||||
screenshot = self.obs.get("screenshot", "") if self.obs else ""
|
||||
logger.info(f"Screenshot available: {'Yes' if screenshot else 'No'}")
|
||||
|
||||
logger.info("Executing code agent...")
|
||||
result = self.code_agent.execute(
|
||||
task_to_execute, screenshot, self.env.controller
|
||||
)
|
||||
|
||||
# Store the result for the worker to access
|
||||
self.last_code_agent_result = result
|
||||
|
||||
logger.info("Code agent execution completed")
|
||||
logger.info(f"Result - Completion reason: {result['completion_reason']}")
|
||||
logger.info(f"Steps executed: {result['steps_executed']}")
|
||||
logger.info(f"Summary: {result['summary']}")
|
||||
|
||||
logger.info("=" * 50)
|
||||
logger.info("GROUNDING AGENT: Code Agent Call Finished")
|
||||
logger.info("=" * 50)
|
||||
|
||||
# Return code to be executed in the environment
|
||||
return "import time; time.sleep(2.222)"
|
||||
else:
|
||||
logger.warning("No task instruction available for code agent call")
|
||||
return "import time; time.sleep(1.111)"
|
||||
|
||||
@agent_action
|
||||
def scroll(self, element_description: str, clicks: int, shift: bool = False):
|
||||
"""Scroll the element in the specified direction
|
||||
Args:
|
||||
element_description:str, a very detailed description of which element to enter scroll in. This description should be at least a full sentence.
|
||||
clicks:int, the number of clicks to scroll can be positive (up) or negative (down).
|
||||
shift:bool, whether to use shift+scroll for horizontal scrolling
|
||||
"""
|
||||
coords1 = self.generate_coords(element_description, self.obs)
|
||||
x, y = self.resize_coordinates(coords1)
|
||||
|
||||
if shift:
|
||||
return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.hscroll({clicks})"
|
||||
else:
|
||||
return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.vscroll({clicks})"
|
||||
|
||||
@agent_action
|
||||
def hotkey(self, keys: List):
|
||||
"""Press a hotkey combination
|
||||
Args:
|
||||
keys:List the keys to press in combination in a list format (e.g. ['ctrl', 'c'])
|
||||
"""
|
||||
# add quotes around the keys
|
||||
keys = [f"'{key}'" for key in keys]
|
||||
return f"import pyautogui; pyautogui.hotkey({', '.join(keys)})"
|
||||
|
||||
@agent_action
|
||||
def hold_and_press(self, hold_keys: List, press_keys: List):
|
||||
"""Hold a list of keys and press a list of keys
|
||||
Args:
|
||||
hold_keys:List, list of keys to hold
|
||||
press_keys:List, list of keys to press in a sequence
|
||||
"""
|
||||
|
||||
press_keys_str = "[" + ", ".join([f"'{key}'" for key in press_keys]) + "]"
|
||||
command = "import pyautogui; "
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyDown({repr(k)}); "
|
||||
command += f"pyautogui.press({press_keys_str}); "
|
||||
for k in hold_keys:
|
||||
command += f"pyautogui.keyUp({repr(k)}); "
|
||||
|
||||
return command
|
||||
|
||||
@agent_action
|
||||
def wait(self, time: float):
|
||||
"""Wait for a specified amount of time
|
||||
Args:
|
||||
time:float the amount of time to wait in seconds
|
||||
"""
|
||||
return f"""import time; time.sleep({time})"""
|
||||
|
||||
@agent_action
|
||||
def done(
|
||||
self,
|
||||
):
|
||||
"""End the current task with a success. Use this when you believe the entire task has been fully completed."""
|
||||
return """DONE"""
|
||||
|
||||
@agent_action
|
||||
def fail(self):
|
||||
"""End the current task with a failure. Use this when you believe the entire task is impossible to complete."""
|
||||
return """FAIL"""
|
||||
@@ -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