chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,425 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from gui_agents.s2.agents.grounding import ACI
|
||||
from gui_agents.s2.agents.worker import Worker
|
||||
from gui_agents.s2.agents.manager import Manager
|
||||
from gui_agents.s2.utils.common_utils import Node
|
||||
from gui_agents.utils import download_kb_data
|
||||
from gui_agents.s2.core.engine import (
|
||||
OpenAIEmbeddingEngine,
|
||||
GeminiEmbeddingEngine,
|
||||
AzureOpenAIEmbeddingEngine,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class UIAgent:
|
||||
"""Base class for UI automation agents"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
platform: str = platform.system().lower(),
|
||||
action_space: str = "pyautogui",
|
||||
observation_type: str = "a11y_tree",
|
||||
search_engine: str = "perplexica",
|
||||
):
|
||||
"""Initialize UIAgent
|
||||
|
||||
Args:
|
||||
engine_params: Configuration parameters for the LLM engine
|
||||
grounding_agent: Instance of ACI class for UI interaction
|
||||
platform: Operating system platform (macos, linux, windows)
|
||||
action_space: Type of action space to use (pyautogui, aci)
|
||||
observation_type: Type of observations to use (a11y_tree, mixed)
|
||||
engine: Search engine to use (perplexica, LLM)
|
||||
"""
|
||||
self.engine_params = engine_params
|
||||
self.grounding_agent = grounding_agent
|
||||
self.platform = platform
|
||||
self.action_space = action_space
|
||||
self.observation_type = observation_type
|
||||
self.engine = search_engine
|
||||
|
||||
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
|
||||
|
||||
def update_narrative_memory(self, trajectory: str) -> None:
|
||||
"""Update narrative memory with task trajectory
|
||||
|
||||
Args:
|
||||
trajectory: String containing task execution trajectory
|
||||
"""
|
||||
pass
|
||||
|
||||
def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str:
|
||||
"""Update episodic memory with subtask trajectory
|
||||
|
||||
Args:
|
||||
meta_data: Metadata about current subtask execution
|
||||
subtask_trajectory: String containing subtask execution trajectory
|
||||
|
||||
Returns:
|
||||
Updated subtask trajectory
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class AgentS2(UIAgent):
|
||||
"""Agent that uses hierarchical planning and directed acyclic graph modeling for UI automation"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
platform: str = platform.system().lower(),
|
||||
action_space: str = "pyautogui",
|
||||
observation_type: str = "mixed",
|
||||
search_engine: Optional[str] = None,
|
||||
memory_root_path: str = os.getcwd(),
|
||||
use_default_kb: bool = False,
|
||||
memory_folder_name: str = "kb_s2",
|
||||
kb_release_tag: str = "v0.2.2",
|
||||
embedding_engine_type: str = "openai",
|
||||
embedding_engine_params: Dict = {},
|
||||
):
|
||||
"""Initialize AgentS2
|
||||
|
||||
Args:
|
||||
engine_params: Configuration parameters for the LLM engine
|
||||
grounding_agent: Instance of ACI class for UI interaction
|
||||
platform: Operating system platform (darwin, linux, windows)
|
||||
action_space: Type of action space to use (pyautogui, other)
|
||||
observation_type: Type of observations to use (a11y_tree, screenshot, mixed)
|
||||
search_engine: Search engine to use (LLM, perplexica)
|
||||
use_default_kb: True to use the default OpenAI kb.
|
||||
memory_root_path: Path to memory directory. Defaults to current working directory.
|
||||
memory_folder_name: Name of memory folder. Defaults to "kb_s2".
|
||||
kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
|
||||
embedding_engine_type: Embedding engine to use for knowledge base. Defaults to "openai". Supports "openai" and "gemini".
|
||||
embedding_engine_params: Parameters for embedding engine. Defaults to {}.
|
||||
"""
|
||||
super().__init__(
|
||||
engine_params,
|
||||
grounding_agent,
|
||||
platform,
|
||||
action_space,
|
||||
observation_type,
|
||||
search_engine,
|
||||
)
|
||||
|
||||
self.memory_root_path = memory_root_path
|
||||
self.memory_folder_name = memory_folder_name
|
||||
self.kb_release_tag = kb_release_tag
|
||||
|
||||
# Initialize agent's knowledge base on user's current working directory.
|
||||
self.local_kb_path = os.path.join(
|
||||
self.memory_root_path, self.memory_folder_name
|
||||
)
|
||||
|
||||
if use_default_kb:
|
||||
if not os.path.exists(os.path.join(self.local_kb_path, self.platform)):
|
||||
print("Downloading Agent S2's default knowledge base...")
|
||||
download_kb_data(
|
||||
version="s2",
|
||||
release_tag=kb_release_tag,
|
||||
download_dir=self.local_kb_path,
|
||||
platform=self.platform,
|
||||
)
|
||||
print(
|
||||
f"Successfully completed download of knowledge base for version s2, tag {self.kb_release_tag}, platform {self.platform}."
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"Path local_kb_path {self.local_kb_path} already exists. Skipping download."
|
||||
)
|
||||
print(
|
||||
f"If you'd like to re-download the initial knowledge base, please delete the existing knowledge base at {self.local_kb_path}."
|
||||
)
|
||||
print(
|
||||
"Note, the knowledge is continually updated during inference. Deleting the knowledge base will wipe out all experience gained since the last knowledge base download."
|
||||
)
|
||||
|
||||
if embedding_engine_type == "openai":
|
||||
self.embedding_engine = OpenAIEmbeddingEngine(**embedding_engine_params)
|
||||
elif embedding_engine_type == "gemini":
|
||||
self.embedding_engine = GeminiEmbeddingEngine(**embedding_engine_params)
|
||||
elif embedding_engine_type == "azure":
|
||||
self.embedding_engine = AzureOpenAIEmbeddingEngine(
|
||||
**embedding_engine_params
|
||||
)
|
||||
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset agent state and initialize components"""
|
||||
# Initialize core components
|
||||
self.planner = Manager(
|
||||
engine_params=self.engine_params,
|
||||
grounding_agent=self.grounding_agent,
|
||||
local_kb_path=self.local_kb_path,
|
||||
embedding_engine=self.embedding_engine,
|
||||
search_engine=self.engine,
|
||||
platform=self.platform,
|
||||
)
|
||||
self.executor = Worker(
|
||||
engine_params=self.engine_params,
|
||||
grounding_agent=self.grounding_agent,
|
||||
local_kb_path=self.local_kb_path,
|
||||
embedding_engine=self.embedding_engine,
|
||||
platform=self.platform,
|
||||
)
|
||||
|
||||
# Reset state variables
|
||||
self.requires_replan: bool = True
|
||||
self.needs_next_subtask: bool = True
|
||||
self.step_count: int = 0
|
||||
self.turn_count: int = 0
|
||||
self.failure_subtask: Optional[Node] = None
|
||||
self.should_send_action: bool = False
|
||||
self.completed_tasks: List[Node] = []
|
||||
self.current_subtask: Optional[Node] = None
|
||||
self.subtasks: List[Node] = []
|
||||
self.search_query: str = ""
|
||||
self.subtask_status: str = "Start"
|
||||
|
||||
def reset_executor_state(self) -> None:
|
||||
"""Reset executor and step counter"""
|
||||
self.executor.reset()
|
||||
self.step_count = 0
|
||||
|
||||
def predict(self, instruction: str, observation: Dict) -> Tuple[Dict, List[str]]:
|
||||
# Initialize the three info dictionaries
|
||||
planner_info = {}
|
||||
executor_info = {}
|
||||
evaluator_info = {
|
||||
"obs_evaluator_response": "",
|
||||
"num_input_tokens_evaluator": 0,
|
||||
"num_output_tokens_evaluator": 0,
|
||||
"evaluator_cost": 0.0,
|
||||
}
|
||||
actions = []
|
||||
|
||||
# If the DONE response by the executor is for a subtask, then the agent should continue with the next subtask without sending the action to the environment
|
||||
while not self.should_send_action:
|
||||
self.subtask_status = "In"
|
||||
# If replan is true, generate a new plan. True at start, after a failed plan, or after subtask completion
|
||||
if self.requires_replan:
|
||||
logger.info("(RE)PLANNING...")
|
||||
planner_info, self.subtasks = self.planner.get_action_queue(
|
||||
instruction=instruction,
|
||||
observation=observation,
|
||||
failed_subtask=self.failure_subtask,
|
||||
completed_subtasks_list=self.completed_tasks,
|
||||
remaining_subtasks_list=self.subtasks,
|
||||
)
|
||||
|
||||
self.requires_replan = False
|
||||
if "search_query" in planner_info:
|
||||
self.search_query = planner_info["search_query"]
|
||||
else:
|
||||
self.search_query = ""
|
||||
|
||||
# use the exectuor to complete the topmost subtask
|
||||
if self.needs_next_subtask:
|
||||
logger.info("GETTING NEXT SUBTASK...")
|
||||
|
||||
# this can be empty if the DAG planner deems that all subtasks are completed
|
||||
if len(self.subtasks) <= 0:
|
||||
self.requires_replan = True
|
||||
self.needs_next_subtask = True
|
||||
self.failure_subtask = None
|
||||
self.completed_tasks.append(self.current_subtask)
|
||||
|
||||
# reset executor state
|
||||
self.reset_executor_state()
|
||||
self.should_send_action = True
|
||||
self.subtask_status = "Done"
|
||||
executor_info = {
|
||||
"executor_plan": "agent.done()",
|
||||
"plan_code": "agent.done()",
|
||||
"reflection": "agent.done()",
|
||||
}
|
||||
actions = ["DONE"]
|
||||
break
|
||||
|
||||
self.current_subtask = self.subtasks.pop(0)
|
||||
logger.info(f"NEXT SUBTASK: {self.current_subtask}")
|
||||
self.needs_next_subtask = False
|
||||
self.subtask_status = "Start"
|
||||
|
||||
# get the next action from the executor
|
||||
executor_info, actions = self.executor.generate_next_action(
|
||||
instruction=instruction,
|
||||
search_query=self.search_query,
|
||||
subtask=self.current_subtask.name,
|
||||
subtask_info=self.current_subtask.info,
|
||||
future_tasks=self.subtasks,
|
||||
done_task=self.completed_tasks,
|
||||
obs=observation,
|
||||
)
|
||||
|
||||
self.step_count += 1
|
||||
|
||||
# set the should_send_action flag to True if the executor returns an action
|
||||
self.should_send_action = True
|
||||
|
||||
# replan on failure
|
||||
if "FAIL" in actions:
|
||||
self.requires_replan = True
|
||||
self.needs_next_subtask = True
|
||||
|
||||
# assign the failed subtask
|
||||
self.failure_subtask = self.current_subtask
|
||||
|
||||
# reset the step count, executor, and evaluator
|
||||
self.reset_executor_state()
|
||||
|
||||
# if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
|
||||
if self.subtasks:
|
||||
self.should_send_action = False
|
||||
|
||||
# replan on subtask completion
|
||||
elif "DONE" in actions:
|
||||
self.requires_replan = True
|
||||
self.needs_next_subtask = True
|
||||
self.failure_subtask = None
|
||||
self.completed_tasks.append(self.current_subtask)
|
||||
|
||||
# reset the step count, executor, and evaluator
|
||||
self.reset_executor_state()
|
||||
|
||||
# if more subtasks are remaining, we don't want to send DONE to the environment but move on to the next subtask
|
||||
if self.subtasks:
|
||||
self.should_send_action = False
|
||||
self.subtask_status = "Done"
|
||||
|
||||
self.turn_count += 1
|
||||
|
||||
# reset the should_send_action flag for next iteration
|
||||
self.should_send_action = False
|
||||
|
||||
# concatenate the three info dictionaries
|
||||
info = {
|
||||
**{
|
||||
k: v
|
||||
for d in [planner_info or {}, executor_info or {}, evaluator_info or {}]
|
||||
for k, v in d.items()
|
||||
}
|
||||
}
|
||||
info.update(
|
||||
{
|
||||
"subtask": self.current_subtask.name,
|
||||
"subtask_info": self.current_subtask.info,
|
||||
"subtask_status": self.subtask_status,
|
||||
}
|
||||
)
|
||||
|
||||
return info, actions
|
||||
|
||||
def update_narrative_memory(self, trajectory: str) -> None:
|
||||
"""Update narrative memory from task trajectory
|
||||
|
||||
Args:
|
||||
trajectory: String containing task execution trajectory
|
||||
"""
|
||||
try:
|
||||
reflection_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "narrative_memory.json"
|
||||
)
|
||||
try:
|
||||
reflections = json.load(open(reflection_path))
|
||||
except:
|
||||
reflections = {}
|
||||
|
||||
if self.search_query not in reflections:
|
||||
reflection = self.planner.summarize_narrative(trajectory)
|
||||
reflections[self.search_query] = reflection
|
||||
|
||||
with open(reflection_path, "w") as f:
|
||||
json.dump(reflections, f, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update narrative memory: {e}")
|
||||
|
||||
def update_episodic_memory(self, meta_data: Dict, subtask_trajectory: str) -> str:
|
||||
"""Update episodic memory from subtask trajectory
|
||||
|
||||
Args:
|
||||
meta_data: Metadata about current subtask execution
|
||||
subtask_trajectory: String containing subtask execution trajectory
|
||||
|
||||
Returns:
|
||||
Updated subtask trajectory
|
||||
"""
|
||||
subtask = meta_data["subtask"]
|
||||
subtask_info = meta_data["subtask_info"]
|
||||
subtask_status = meta_data["subtask_status"]
|
||||
# Handle subtask trajectory
|
||||
if subtask_status == "Start" or subtask_status == "Done":
|
||||
# If it's a new subtask start, finalize the previous subtask trajectory if it exists
|
||||
if subtask_trajectory:
|
||||
subtask_trajectory += "\nSubtask Completed.\n"
|
||||
subtask_key = subtask_trajectory.split(
|
||||
"\n----------------------\n\nPlan:\n"
|
||||
)[0]
|
||||
try:
|
||||
subtask_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "episodic_memory.json"
|
||||
)
|
||||
kb = json.load(open(subtask_path))
|
||||
except:
|
||||
kb = {}
|
||||
if subtask_key not in kb.keys():
|
||||
subtask_summarization = self.planner.summarize_episode(
|
||||
subtask_trajectory
|
||||
)
|
||||
kb[subtask_key] = subtask_summarization
|
||||
else:
|
||||
subtask_summarization = kb[subtask_key]
|
||||
logger.info("subtask_key: %s", subtask_key)
|
||||
logger.info("subtask_summarization: %s", subtask_summarization)
|
||||
with open(subtask_path, "w") as fout:
|
||||
json.dump(kb, fout, indent=2)
|
||||
# Reset for the next subtask
|
||||
subtask_trajectory = ""
|
||||
# Start a new subtask trajectory
|
||||
subtask_trajectory = (
|
||||
"Task:\n"
|
||||
+ self.search_query
|
||||
+ "\n\nSubtask: "
|
||||
+ subtask
|
||||
+ "\nSubtask Instruction: "
|
||||
+ subtask_info
|
||||
+ "\n----------------------\n\nPlan:\n"
|
||||
+ meta_data["executor_plan"]
|
||||
+ "\n"
|
||||
)
|
||||
elif subtask_status == "In":
|
||||
# Continue appending to the current subtask trajectory if it's still ongoing
|
||||
subtask_trajectory += (
|
||||
"\n----------------------\n\nPlan:\n"
|
||||
+ meta_data["executor_plan"]
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
return subtask_trajectory
|
||||
@@ -0,0 +1,600 @@
|
||||
import ast
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
from pytesseract import Output
|
||||
|
||||
from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s2.core.mllm import LMMAgent
|
||||
from gui_agents.s2.utils.common_utils import (
|
||||
call_llm_safe,
|
||||
parse_single_code_from_string,
|
||||
)
|
||||
|
||||
|
||||
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;
|
||||
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
|
||||
|
||||
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 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"):
|
||||
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 \"password\" | 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,
|
||||
platform: str,
|
||||
engine_params_for_generation: Dict,
|
||||
engine_params_for_grounding: Dict,
|
||||
width: int = 1920,
|
||||
height: int = 1080,
|
||||
):
|
||||
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 = []
|
||||
|
||||
# Coordinates used during ACI execution
|
||||
self.coords1 = None
|
||||
self.coords2 = 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,
|
||||
)
|
||||
|
||||
# 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
|
||||
|
||||
# Takes a description based action and assigns the coordinates for any coordinate based action
|
||||
# Raises an error if function can't be parsed
|
||||
def assign_coordinates(self, plan: str, obs: Dict):
|
||||
|
||||
# Reset coords from previous action generation
|
||||
self.coords1, self.coords2 = None, None
|
||||
|
||||
try:
|
||||
# Extract the function name and args
|
||||
action = parse_single_code_from_string(plan.split("Grounded Action")[-1])
|
||||
function_name = re.match(r"(\w+\.\w+)\(", action).group(1)
|
||||
args = self.parse_function_args(action)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error in parsing grounded action: {e}") from e
|
||||
|
||||
# arg0 is a description
|
||||
if (
|
||||
function_name in ["agent.click", "agent.type", "agent.scroll"]
|
||||
and len(args) >= 1
|
||||
and args[0] != None
|
||||
):
|
||||
self.coords1 = self.generate_coords(args[0], obs)
|
||||
# arg0 and arg1 are descriptions
|
||||
elif function_name == "agent.drag_and_drop" and len(args) >= 2:
|
||||
self.coords1 = self.generate_coords(args[0], obs)
|
||||
self.coords2 = self.generate_coords(args[1], obs)
|
||||
# arg0 and arg1 are text phrases
|
||||
elif function_name == "agent.highlight_text_span" and len(args) >= 2:
|
||||
self.coords1 = self.generate_text_coords(args[0], obs, alignment="start")
|
||||
self.coords2 = self.generate_text_coords(args[1], obs, alignment="end")
|
||||
|
||||
# Resize from grounding model dim into OSWorld dim (1920 * 1080)
|
||||
def resize_coordinates(self, coordinates: List[int]) -> List[int]:
|
||||
# User explicitly passes the grounding model dimensions
|
||||
if {"grounding_width", "grounding_height"}.issubset(
|
||||
self.engine_params_for_grounding
|
||||
):
|
||||
grounding_width = self.engine_params_for_grounding["grounding_width"]
|
||||
grounding_height = self.engine_params_for_grounding["grounding_height"]
|
||||
# Default to (1000, 1000), which is UI-TARS resizing
|
||||
else:
|
||||
grounding_width = 1000
|
||||
grounding_height = 1000
|
||||
|
||||
return [
|
||||
round(coordinates[0] * self.width / grounding_width),
|
||||
round(coordinates[1] * self.height / grounding_height),
|
||||
]
|
||||
|
||||
# Given a generated ACI function, returns a list of argument values, where descriptions are at the front of the list
|
||||
def parse_function_args(self, function: str) -> List[str]:
|
||||
tree = ast.parse(function)
|
||||
call_node = tree.body[0].value
|
||||
|
||||
def safe_eval(node):
|
||||
if isinstance(
|
||||
node, ast.Constant
|
||||
): # Handles literals like numbers, strings, etc.
|
||||
return node.value
|
||||
else:
|
||||
return ast.unparse(node) # Return as a string if not a literal
|
||||
|
||||
positional_args = [safe_eval(arg) for arg in call_node.args]
|
||||
keyword_args = {kw.arg: safe_eval(kw.value) for kw in call_node.keywords}
|
||||
|
||||
res = []
|
||||
|
||||
for key, val in keyword_args.items():
|
||||
if "description" in key:
|
||||
res.append(val)
|
||||
|
||||
for arg in positional_args:
|
||||
res.append(arg)
|
||||
|
||||
return res
|
||||
|
||||
@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
|
||||
"""
|
||||
x, y = self.resize_coordinates(self.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)"
|
||||
|
||||
@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
|
||||
"""
|
||||
return f"import pyautogui; pyautogui.hotkey('win'); time.sleep(0.5); pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); pyautogui.hotkey('enter'); time.sleep(0.5)"
|
||||
|
||||
@agent_action
|
||||
def type(
|
||||
self,
|
||||
element_description: Optional[str] = None,
|
||||
text: str = "",
|
||||
overwrite: bool = False,
|
||||
enter: bool = False,
|
||||
):
|
||||
"""Type text 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.
|
||||
"""
|
||||
|
||||
if self.coords1 is not None:
|
||||
# If a node is found, retrieve its coordinates and size
|
||||
# Start typing at the center of the element
|
||||
|
||||
x, y = self.resize_coordinates(self.coords1)
|
||||
|
||||
command = "import pyautogui; "
|
||||
command += f"pyautogui.click({x}, {y}); "
|
||||
|
||||
if overwrite:
|
||||
command += (
|
||||
f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
|
||||
)
|
||||
|
||||
command += f"pyautogui.write({repr(text)}); "
|
||||
|
||||
if enter:
|
||||
command += "pyautogui.press('enter'); "
|
||||
else:
|
||||
# If no element is found, start typing at the current cursor location
|
||||
command = "import pyautogui; "
|
||||
|
||||
if overwrite:
|
||||
command += (
|
||||
f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
|
||||
)
|
||||
|
||||
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
|
||||
"""
|
||||
x1, y1 = self.resize_coordinates(self.coords1)
|
||||
x2, y2 = self.resize_coordinates(self.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.); 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):
|
||||
"""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.
|
||||
"""
|
||||
|
||||
x1, y1 = self.coords1
|
||||
x2, y2 = self.coords2
|
||||
|
||||
command = "import pyautogui; "
|
||||
command += f"pyautogui.moveTo({x1}, {y1}); "
|
||||
command += f"pyautogui.dragTo({x2}, {y2}, duration=1.); 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 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
|
||||
"""
|
||||
|
||||
x, y = self.resize_coordinates(self.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,
|
||||
return_value: Optional[Union[Dict, str, List, Tuple, int, float, bool]] = None,
|
||||
):
|
||||
"""End the current task with a success and the required return value"""
|
||||
self.returned_info = return_value
|
||||
return """DONE"""
|
||||
|
||||
@agent_action
|
||||
def fail(self):
|
||||
"""End the current task with a failure, and replan the whole task."""
|
||||
return """FAIL"""
|
||||
@@ -0,0 +1,321 @@
|
||||
import logging
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s2.agents.grounding import ACI
|
||||
from gui_agents.s2.core.module import BaseModule
|
||||
from gui_agents.s2.core.knowledge import KnowledgeBase
|
||||
from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s2.core.engine import OpenAIEmbeddingEngine
|
||||
from gui_agents.s2.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,
|
||||
embedding_engine,
|
||||
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 planner
|
||||
sys_prompt = PROCEDURAL_MEMORY.COMBINED_MANAGER_PROMPT
|
||||
|
||||
self.generator_agent = self._create_agent(sys_prompt)
|
||||
|
||||
# Initialize the remaining modules
|
||||
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.embedding_engine = embedding_engine
|
||||
self.knowledge_base = KnowledgeBase(
|
||||
embedding_engine=self.embedding_engine,
|
||||
local_kb_path=self.local_kb_path,
|
||||
platform=platform,
|
||||
engine_params=engine_params,
|
||||
)
|
||||
|
||||
self.planner_history = []
|
||||
|
||||
self.turn_count = 0
|
||||
self.search_engine = search_engine
|
||||
self.multi_round = multi_round
|
||||
|
||||
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, role="user")
|
||||
subtask_summarization = call_llm_safe(self.episode_summarization_agent)
|
||||
self.episode_summarization_agent.add_message(
|
||||
subtask_summarization, role="assistant"
|
||||
)
|
||||
|
||||
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, role="user")
|
||||
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,
|
||||
failed_subtask: Optional[Node] = None,
|
||||
completed_subtasks_list: List[Node] = [],
|
||||
remaining_subtasks_list: List[Node] = [],
|
||||
) -> Tuple[Dict, str]:
|
||||
agent = self.grounding_agent
|
||||
|
||||
# Converts a list of DAG Nodes into a natural langauge list
|
||||
def format_subtask_list(subtasks: List[Node]) -> str:
|
||||
res = ""
|
||||
for idx, node in enumerate(subtasks):
|
||||
res += f"{idx+1}. **{node.name}**:\n"
|
||||
bullets = re.split(r"(?<=[.!?;]) +", node.info)
|
||||
for bullet in bullets:
|
||||
res += f" - {bullet}\n"
|
||||
res += "\n"
|
||||
return res
|
||||
|
||||
# Perform Retrieval only at the first planning step
|
||||
if self.turn_count == 0:
|
||||
|
||||
self.search_query = self.knowledge_base.formulate_query(
|
||||
instruction, observation
|
||||
)
|
||||
|
||||
most_similar_task = ""
|
||||
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
|
||||
)
|
||||
)
|
||||
|
||||
# Re-plan on failure case
|
||||
if failed_subtask:
|
||||
generator_message = (
|
||||
f"The subtask {failed_subtask} cannot be completed. Please generate a new plan for the remainder of the trajectory.\n\n"
|
||||
f"Successfully Completed Subtasks:\n{format_subtask_list(completed_subtasks_list)}\n"
|
||||
)
|
||||
# Re-plan on subtask completion case
|
||||
elif len(completed_subtasks_list) + len(remaining_subtasks_list) > 0:
|
||||
generator_message = (
|
||||
"The current trajectory and desktop state is provided. Please revise the plan for the following trajectory.\n\n"
|
||||
f"Successfully Completed Subtasks:\n{format_subtask_list(completed_subtasks_list)}\n"
|
||||
f"Future Remaining Subtasks:\n{format_subtask_list(remaining_subtasks_list)}\n"
|
||||
)
|
||||
# Initial plan case
|
||||
else:
|
||||
generator_message = "Please generate the initial plan for the task.\n"
|
||||
|
||||
logger.info("GENERATOR MESSAGE: %s", generator_message)
|
||||
|
||||
self.generator_agent.add_message(
|
||||
generator_message,
|
||||
image_content=observation.get("screenshot", None),
|
||||
role="user",
|
||||
)
|
||||
|
||||
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, role="assistant")
|
||||
self.planner_history.append(plan)
|
||||
self.turn_count += 1
|
||||
|
||||
# Set Cost based on GPT-4o
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
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]:
|
||||
# For the re-planning case, remove the prior input since this should only translate the new plan
|
||||
self.dag_translator_agent.reset()
|
||||
|
||||
# Add initial instruction and plan to the agent's message history
|
||||
self.dag_translator_agent.add_message(
|
||||
f"Instruction: {instruction}\nPlan: {plan}", role="user"
|
||||
)
|
||||
|
||||
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, role="assistant")
|
||||
|
||||
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,
|
||||
failed_subtask: Optional[Node] = None,
|
||||
completed_subtasks_list: List[Node] = [],
|
||||
remaining_subtasks_list: List[Node] = [],
|
||||
):
|
||||
"""Generate the action list based on the instruction
|
||||
instruction:str: Instruction for the task
|
||||
"""
|
||||
|
||||
planner_info, plan = self._generate_step_by_step_plan(
|
||||
observation,
|
||||
instruction,
|
||||
failed_subtask,
|
||||
completed_subtasks_list,
|
||||
remaining_subtasks_list,
|
||||
)
|
||||
|
||||
# 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
|
||||
@@ -0,0 +1,254 @@
|
||||
import logging
|
||||
import re
|
||||
import textwrap
|
||||
from typing import Dict, List, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s2.agents.grounding import ACI
|
||||
from gui_agents.s2.core.module import BaseModule
|
||||
from gui_agents.s2.core.knowledge import KnowledgeBase
|
||||
from gui_agents.s2.memory.procedural_memory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s2.utils.common_utils import (
|
||||
Node,
|
||||
calculate_tokens,
|
||||
call_llm_safe,
|
||||
parse_single_code_from_string,
|
||||
sanitize_code,
|
||||
extract_first_agent_function,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class Worker(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
local_kb_path: str,
|
||||
embedding_engine,
|
||||
platform: str = platform.system().lower(),
|
||||
enable_reflection: bool = True,
|
||||
use_subtask_experience: bool = True,
|
||||
):
|
||||
"""
|
||||
Worker receives a subtask list and active subtask and generates the next action for the to execute.
|
||||
Args:
|
||||
engine_params: Dict
|
||||
Parameters for the multimodal engine
|
||||
grounding_agent: Agent
|
||||
The grounding agent to use
|
||||
local_kb_path: str
|
||||
Path to knowledge base
|
||||
platform: str
|
||||
OS platform the agent runs on (darwin, linux, windows)
|
||||
enable_reflection: bool
|
||||
Whether to enable reflection
|
||||
use_subtask_experience: bool
|
||||
Whether to use subtask experience
|
||||
"""
|
||||
super().__init__(engine_params, platform)
|
||||
|
||||
self.grounding_agent = grounding_agent
|
||||
self.local_kb_path = local_kb_path
|
||||
self.embedding_engine = embedding_engine
|
||||
self.enable_reflection = enable_reflection
|
||||
self.use_subtask_experience = use_subtask_experience
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
if self.platform != "linux":
|
||||
skipped_actions = ["set_cell_values"]
|
||||
else:
|
||||
skipped_actions = []
|
||||
|
||||
sys_prompt = PROCEDURAL_MEMORY.construct_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.knowledge_base = KnowledgeBase(
|
||||
embedding_engine=self.embedding_engine,
|
||||
local_kb_path=self.local_kb_path,
|
||||
platform=self.platform,
|
||||
engine_params=self.engine_params,
|
||||
)
|
||||
|
||||
self.turn_count = 0
|
||||
self.worker_history = []
|
||||
self.reflections = []
|
||||
self.cost_this_turn = 0
|
||||
self.screenshot_inputs = []
|
||||
self.planner_history = []
|
||||
self.max_trajector_length = 8
|
||||
|
||||
def flush_messages(self):
|
||||
# generator msgs are alternating [user, assistant], so 2 per round
|
||||
if len(self.generator_agent.messages) > 2 * self.max_trajector_length + 1:
|
||||
self.generator_agent.remove_message_at(1)
|
||||
self.generator_agent.remove_message_at(1)
|
||||
# reflector msgs are all [(user text, user image)], so 1 per round
|
||||
if len(self.reflection_agent.messages) > self.max_trajector_length + 1:
|
||||
self.reflection_agent.remove_message_at(1)
|
||||
|
||||
def generate_next_action(
|
||||
self,
|
||||
instruction: str,
|
||||
search_query: str,
|
||||
subtask: str,
|
||||
subtask_info: Dict,
|
||||
future_tasks: List[Node],
|
||||
done_task: List[Node],
|
||||
obs: Dict,
|
||||
) -> Tuple[Dict, List]:
|
||||
"""
|
||||
Predict the next action(s) based on the current observation.
|
||||
"""
|
||||
# Provide the top_app to the Grounding Agent to remove all other applications from the tree. At t=0, top_app is None
|
||||
agent = self.grounding_agent
|
||||
|
||||
# Get RAG knowledge, only update system message at t=0
|
||||
if self.turn_count == 0:
|
||||
if self.use_subtask_experience:
|
||||
subtask_query_key = (
|
||||
"Task:\n"
|
||||
+ search_query
|
||||
+ "\n\nSubtask: "
|
||||
+ subtask
|
||||
+ "\nSubtask Instruction: "
|
||||
+ subtask_info
|
||||
)
|
||||
retrieved_similar_subtask, retrieved_subtask_experience = (
|
||||
self.knowledge_base.retrieve_episodic_experience(subtask_query_key)
|
||||
)
|
||||
|
||||
# Dirty fix to replace id with element description during subtask retrieval
|
||||
pattern = r"\(\d+"
|
||||
retrieved_subtask_experience = re.sub(
|
||||
pattern, "(element_description", retrieved_subtask_experience
|
||||
)
|
||||
retrieved_subtask_experience = retrieved_subtask_experience.replace(
|
||||
"_id", "_description"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"SIMILAR SUBTASK EXPERIENCE: %s",
|
||||
retrieved_similar_subtask
|
||||
+ "\n"
|
||||
+ retrieved_subtask_experience.strip(),
|
||||
)
|
||||
instruction += "\nYou may refer to some similar subtask experience if you think they are useful. {}".format(
|
||||
retrieved_similar_subtask + "\n" + retrieved_subtask_experience
|
||||
)
|
||||
|
||||
self.generator_agent.add_system_prompt(
|
||||
self.generator_agent.system_prompt.replace(
|
||||
"SUBTASK_DESCRIPTION", subtask
|
||||
)
|
||||
.replace("TASK_DESCRIPTION", instruction)
|
||||
.replace("FUTURE_TASKS", ", ".join([f.name for f in future_tasks]))
|
||||
.replace("DONE_TASKS", ",".join(d.name for d in done_task))
|
||||
)
|
||||
|
||||
# Reflection generation does not add its own response, it only gets the trajectory
|
||||
reflection = None
|
||||
if self.enable_reflection:
|
||||
# Load the initial subtask info
|
||||
if self.turn_count == 0:
|
||||
text_content = textwrap.dedent(f"""
|
||||
Subtask Description: {subtask}
|
||||
Subtask Information: {subtask_info}
|
||||
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:
|
||||
text_content = self.clean_worker_generation_for_reflection(
|
||||
self.planner_history[-1]
|
||||
)
|
||||
self.reflection_agent.add_message(
|
||||
text_content=text_content,
|
||||
image_content=obs["screenshot"],
|
||||
role="user",
|
||||
)
|
||||
reflection = call_llm_safe(self.reflection_agent)
|
||||
self.reflections.append(reflection)
|
||||
logger.info("REFLECTION: %s", reflection)
|
||||
|
||||
generator_message = (
|
||||
f"\nYou may use this reflection on the previous action and overall trajectory: {reflection}\n"
|
||||
if reflection and self.turn_count > 0
|
||||
else ""
|
||||
) + f"Text Buffer = [{','.join(agent.notes)}]."
|
||||
|
||||
# Only provide subinfo in the very first message to avoid over influence and redundancy
|
||||
if self.turn_count == 0:
|
||||
generator_message += f"Remember only complete the subtask: {subtask}\n"
|
||||
generator_message += f"You can use this extra information for completing the current subtask: {subtask_info}.\n"
|
||||
|
||||
# logger.info("GENERATOR MESSAGE: %s", generator_message)
|
||||
|
||||
self.generator_agent.add_message(
|
||||
generator_message, image_content=obs["screenshot"], role="user"
|
||||
)
|
||||
|
||||
plan = call_llm_safe(self.generator_agent)
|
||||
self.planner_history.append(plan)
|
||||
logger.info("PLAN: %s", plan)
|
||||
self.generator_agent.add_message(plan, role="assistant")
|
||||
|
||||
# Calculate input/output tokens and gpt-4o cost
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
cost = input_tokens * (0.0050 / 1000) + output_tokens * (0.0150 / 1000)
|
||||
self.cost_this_turn += cost
|
||||
logger.info("EXECTUOR COST: %s", self.cost_this_turn)
|
||||
|
||||
# Use the DescriptionBasedACI 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 = {
|
||||
"current_subtask": subtask,
|
||||
"current_subtask_info": subtask_info,
|
||||
"executor_plan": plan,
|
||||
"plan_code": plan_code,
|
||||
"reflection": reflection,
|
||||
"num_input_tokens_executor": input_tokens,
|
||||
"num_output_tokens_executor": output_tokens,
|
||||
}
|
||||
self.turn_count += 1
|
||||
|
||||
self.screenshot_inputs.append(obs["screenshot"])
|
||||
self.flush_messages()
|
||||
|
||||
return executor_info, [exec_code]
|
||||
|
||||
# Removes the previous action verification, and removes any extraneous grounded actions
|
||||
def clean_worker_generation_for_reflection(self, worker_generation: str) -> str:
|
||||
# Remove the previous action verification
|
||||
res = worker_generation[worker_generation.find("(Screenshot Analysis)") :]
|
||||
action = extract_first_agent_function(worker_generation)
|
||||
# Cut off extra grounded actions
|
||||
res = res[: res.find("(Grounded Action)")]
|
||||
res += f"(Grounded Action)\n```python\n{action}\n```\n"
|
||||
return res
|
||||
Reference in New Issue
Block a user