wehub-resource-sync
533 lines
18 KiB
Python
533 lines
18 KiB
Python
import base64
|
|
import os
|
|
import platform
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
import psutil
|
|
import requests
|
|
from gui_agents.s1.utils.common_utils import box_iou
|
|
|
|
if platform.system() == "Windows":
|
|
import pywinauto
|
|
from pywinauto import Desktop
|
|
import win32gui
|
|
import win32process
|
|
|
|
from gui_agents.s1.aci.ACI import ACI, agent_action
|
|
|
|
|
|
# Helper functions
|
|
def _normalize_key(key: str) -> str:
|
|
"""Convert 'ctrl' to 'control' for pyautogui compatibility"""
|
|
return "ctrl" if key == "control" else key
|
|
|
|
|
|
def list_apps_in_directories():
|
|
directories_to_search = [
|
|
os.environ.get("PROGRAMFILES", "C:\\Program Files"),
|
|
os.environ.get("PROGRAMFILES(X86)", "C:\\Program Files (x86)"),
|
|
]
|
|
apps = []
|
|
for directory in directories_to_search:
|
|
if os.path.exists(directory):
|
|
for root, dirs, files in os.walk(directory):
|
|
for file in files:
|
|
if file.endswith(".exe"):
|
|
apps.append(file)
|
|
return apps
|
|
|
|
|
|
# WindowsACI Class
|
|
class WindowsACI(ACI):
|
|
def __init__(self, top_app_only: bool = True, ocr: bool = False):
|
|
super().__init__(top_app_only=top_app_only, ocr=ocr)
|
|
self.nodes = []
|
|
self.all_apps = list_apps_in_directories()
|
|
|
|
def get_active_apps(self, obs: Dict) -> List[str]:
|
|
return UIElement.get_current_applications(obs)
|
|
|
|
def get_top_app(self, obs: Dict) -> str:
|
|
return UIElement.get_top_app(obs)
|
|
|
|
def preserve_nodes(self, tree, exclude_roles=None):
|
|
if exclude_roles is None:
|
|
exclude_roles = set()
|
|
|
|
preserved_nodes = []
|
|
|
|
def traverse_and_preserve(element):
|
|
role = element.role()
|
|
|
|
if role not in exclude_roles:
|
|
position = element.position()
|
|
size = element.size()
|
|
if position and size:
|
|
x, y = position
|
|
w, h = size
|
|
|
|
if x >= 0 and y >= 0 and w > 0 and h > 0:
|
|
preserved_nodes.append(
|
|
{
|
|
"position": (x, y),
|
|
"size": (w, h),
|
|
"title": element.title(),
|
|
"text": element.text(),
|
|
"role": role,
|
|
}
|
|
)
|
|
|
|
children = element.children()
|
|
if children:
|
|
for child_element in children:
|
|
traverse_and_preserve(child_element)
|
|
|
|
traverse_and_preserve(tree)
|
|
return preserved_nodes
|
|
|
|
def extract_elements_from_screenshot(self, screenshot: bytes) -> Dict[str, Any]:
|
|
url = os.environ.get("OCR_SERVER_ADDRESS")
|
|
if not url:
|
|
raise EnvironmentError("OCR SERVER ADDRESS NOT SET")
|
|
|
|
encoded_screenshot = base64.b64encode(screenshot).decode("utf-8")
|
|
response = requests.post(url, json={"img_bytes": encoded_screenshot})
|
|
|
|
if response.status_code != 200:
|
|
return {
|
|
"error": f"Request failed with status code {response.status_code}",
|
|
"results": [],
|
|
}
|
|
return response.json()
|
|
|
|
def add_ocr_elements(
|
|
self,
|
|
screenshot,
|
|
linearized_accessibility_tree: List[str],
|
|
preserved_nodes: List[Dict],
|
|
) -> Tuple[List[str], List[Dict]]:
|
|
"""
|
|
Add OCR-detected elements to the accessibility tree if they don't overlap with existing elements
|
|
Uses optimized NumPy implementation
|
|
"""
|
|
# Convert preserved nodes to numpy array of bounding boxes
|
|
if preserved_nodes:
|
|
tree_bboxes = np.array(
|
|
[
|
|
[
|
|
node["position"][0],
|
|
node["position"][1],
|
|
node["position"][0] + node["size"][0],
|
|
node["position"][1] + node["size"][1],
|
|
]
|
|
for node in preserved_nodes
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
else:
|
|
tree_bboxes = np.empty((0, 4), dtype=np.float32)
|
|
|
|
try:
|
|
ocr_bboxes = self.extract_elements_from_screenshot(screenshot)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
ocr_bboxes = []
|
|
else:
|
|
if ocr_bboxes:
|
|
preserved_nodes_index = len(preserved_nodes)
|
|
|
|
# Convert OCR boxes to numpy array
|
|
ocr_boxes_array = np.array(
|
|
[
|
|
[
|
|
int(box.get("left", 0)),
|
|
int(box.get("top", 0)),
|
|
int(box.get("right", 0)),
|
|
int(box.get("bottom", 0)),
|
|
]
|
|
for _, _, box in ocr_bboxes["results"]
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
# Calculate max IOUs efficiently
|
|
if len(tree_bboxes) > 0:
|
|
max_ious = box_iou(tree_bboxes, ocr_boxes_array).max(axis=0)
|
|
else:
|
|
max_ious = np.zeros(len(ocr_boxes_array))
|
|
|
|
# Process boxes with low IOU
|
|
for idx, ((_, content, box), max_iou) in enumerate(
|
|
zip(ocr_bboxes["results"], max_ious)
|
|
):
|
|
if max_iou < 0.1:
|
|
x1 = int(box.get("left", 0))
|
|
y1 = int(box.get("top", 0))
|
|
x2 = int(box.get("right", 0))
|
|
y2 = int(box.get("bottom", 0))
|
|
|
|
linearized_accessibility_tree.append(
|
|
f"{preserved_nodes_index}\tButton\t\t{content}\t\t"
|
|
)
|
|
|
|
node = {
|
|
"position": (x1, y1),
|
|
"size": (x2 - x1, y2 - y1),
|
|
"title": "",
|
|
"text": content,
|
|
"role": "Button",
|
|
}
|
|
preserved_nodes.append(node)
|
|
preserved_nodes_index += 1
|
|
|
|
return linearized_accessibility_tree, preserved_nodes
|
|
|
|
def linearize_and_annotate_tree(
|
|
self, obs: Dict, show_all_elements: bool = False
|
|
) -> str:
|
|
desktop = Desktop(backend="uia")
|
|
try:
|
|
tree = desktop.window(
|
|
handle=win32gui.GetForegroundWindow()
|
|
).wrapper_object()
|
|
except Exception as e:
|
|
print(f"Error accessing foreground window: {e}")
|
|
self.nodes = []
|
|
return ""
|
|
|
|
exclude_roles = ["Pane", "Group", "Unknown"]
|
|
preserved_nodes = self.preserve_nodes(UIElement(tree), exclude_roles).copy()
|
|
|
|
if not preserved_nodes and show_all_elements:
|
|
preserved_nodes = self.preserve_nodes(
|
|
UIElement(tree), exclude_roles=[]
|
|
).copy()
|
|
|
|
tree_elements = ["id\trole\ttitle\ttext"]
|
|
for idx, node in enumerate(preserved_nodes):
|
|
tree_elements.append(
|
|
f"{idx}\t{node['role']}\t{node['title']}\t{node['text']}"
|
|
)
|
|
|
|
if self.ocr:
|
|
screenshot = obs.get("screenshot", None)
|
|
if screenshot is not None:
|
|
# return tree_elements, preserved_nodes
|
|
tree_elements, preserved_nodes = self.add_ocr_elements(
|
|
screenshot, tree_elements, preserved_nodes
|
|
)
|
|
|
|
self.nodes = preserved_nodes
|
|
return "\n".join(tree_elements)
|
|
|
|
def find_element(self, element_id: int) -> Dict:
|
|
if not self.nodes:
|
|
print("No elements found in the accessibility tree.")
|
|
raise IndexError("No elements to select.")
|
|
try:
|
|
return self.nodes[element_id]
|
|
except IndexError:
|
|
print("The index of the selected element was out of range.")
|
|
self.index_out_of_range_flag = True
|
|
return self.nodes[0]
|
|
|
|
@agent_action
|
|
def open(self, app_or_file_name: str):
|
|
"""Open an application or file
|
|
Args:
|
|
app_or_file_name:str, the name of the application or file to open
|
|
"""
|
|
command = f"import pyautogui; import time; pyautogui.hotkey('win', 'r', interval=0.5); pyautogui.typewrite({repr(app_or_file_name)}); pyautogui.press('enter'); time.sleep(1.0)"
|
|
return command
|
|
|
|
@agent_action
|
|
def switch_applications(self, app_or_file_name):
|
|
"""Switch to a different application. Utility function to use instead of alt+tab
|
|
Args:
|
|
app_or_file_name:str, the name of the application or file to switch to
|
|
"""
|
|
command = f"import pyautogui; import time; pyautogui.hotkey('win', 'd', interval=0.5); pyautogui.typewrite({repr(app_or_file_name)}); pyautogui.press('enter'); time.sleep(1.0)"
|
|
return command
|
|
|
|
@agent_action
|
|
def click(
|
|
self,
|
|
element_id: int,
|
|
num_clicks: int = 1,
|
|
button_type: str = "left",
|
|
hold_keys: List = [],
|
|
):
|
|
"""Click on the element
|
|
Args:
|
|
element_id:int, ID of the element to click on
|
|
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
|
|
"""
|
|
node = self.find_element(element_id)
|
|
coordinates: Tuple[int, int] = node["position"]
|
|
sizes: Tuple[int, int] = node["size"]
|
|
|
|
# Calculate the center of the element
|
|
x = int(coordinates[0] + sizes[0] // 2)
|
|
y = int(coordinates[1] + sizes[1] // 2)
|
|
|
|
command = "import pyautogui; "
|
|
|
|
# Normalize any 'ctrl' to 'control'
|
|
hold_keys = [_normalize_key(k) for k in hold_keys]
|
|
|
|
for k in hold_keys:
|
|
command += f"pyautogui.keyDown({repr(k)}); "
|
|
command += f"""pyautogui.click({x}, {y}, clicks={num_clicks}, button={repr(button_type)}); """
|
|
for k in hold_keys:
|
|
command += f"pyautogui.keyUp({repr(k)}); "
|
|
return command
|
|
|
|
@agent_action
|
|
def type(
|
|
self,
|
|
element_id: int = None,
|
|
text: str = "",
|
|
overwrite: bool = False,
|
|
enter: bool = False,
|
|
):
|
|
"""Type text into the element
|
|
Args:
|
|
element_id:int ID of the element to type into. If not provided, typing will start at the current cursor location.
|
|
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.
|
|
"""
|
|
try:
|
|
node = self.find_element(element_id) if element_id is not None else None
|
|
except:
|
|
node = None
|
|
|
|
if node is not None:
|
|
coordinates = node["position"]
|
|
sizes = node["size"]
|
|
|
|
x = int(coordinates[0] + sizes[0] // 2)
|
|
y = int(coordinates[1] + sizes[1] // 2)
|
|
|
|
command = "import pyautogui; "
|
|
command += f"pyautogui.click({x}, {y}); "
|
|
|
|
if overwrite:
|
|
command += f"pyautogui.hotkey('ctrl', 'a', interval=0.5); pyautogui.press('backspace'); "
|
|
|
|
command += f"pyautogui.write({repr(text)}); "
|
|
|
|
if enter:
|
|
command += "pyautogui.press('enter'); "
|
|
else:
|
|
command = "import pyautogui; "
|
|
|
|
if overwrite:
|
|
command += f"pyautogui.hotkey('ctrl', 'a', interval=0.5); 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 for reuse during this task. Can be used for copy-pasting text, saving elements, etc. Use this instead of ctrl+c, ctrl+v.
|
|
Args:
|
|
text:List[str] the text to save to the knowledge
|
|
"""
|
|
self.notes.extend(text)
|
|
return """WAIT"""
|
|
|
|
@agent_action
|
|
def drag_and_drop(self, drag_from_id: int, drop_on_id: int, hold_keys: List = []):
|
|
"""Drag element1 and drop it on element2.
|
|
Args:
|
|
drag_from_id:int ID of element to drag
|
|
drop_on_id:int ID of element to drop on
|
|
hold_keys:List list of keys to hold while dragging
|
|
"""
|
|
node1 = self.find_element(drag_from_id)
|
|
node2 = self.find_element(drop_on_id)
|
|
coordinates1 = node1["position"]
|
|
sizes1 = node1["size"]
|
|
|
|
coordinates2 = node2["position"]
|
|
sizes2 = node2["size"]
|
|
|
|
x1 = int(coordinates1[0] + sizes1[0] // 2)
|
|
y1 = int(coordinates1[1] + sizes1[1] // 2)
|
|
|
|
x2 = int(coordinates2[0] + sizes2[0] // 2)
|
|
y2 = int(coordinates2[1] + sizes2[1] // 2)
|
|
|
|
command = "import pyautogui; "
|
|
|
|
command += f"pyautogui.moveTo({x1}, {y1}); "
|
|
for k in hold_keys:
|
|
command += f"pyautogui.keyDown({repr(k)}); "
|
|
command += f"pyautogui.dragTo({x2}, {y2}, duration=1.0); pyautogui.mouseUp(); "
|
|
for k in hold_keys:
|
|
command += f"pyautogui.keyUp({repr(k)}); "
|
|
|
|
return command
|
|
|
|
@agent_action
|
|
def scroll(self, element_id: int, clicks: int):
|
|
"""Scroll in the specified direction inside the specified element
|
|
Args:
|
|
element_id:int ID of the element to scroll in
|
|
clicks:int the number of clicks to scroll can be positive (up) or negative (down).
|
|
"""
|
|
try:
|
|
node = self.find_element(element_id)
|
|
except:
|
|
node = self.find_element(0)
|
|
|
|
coordinates = node["position"]
|
|
sizes = node["size"]
|
|
|
|
x = int(coordinates[0] + sizes[0] // 2)
|
|
y = int(coordinates[1] + sizes[1] // 2)
|
|
command = (
|
|
f"import pyautogui; pyautogui.moveTo({x}, {y}); pyautogui.scroll({clicks})"
|
|
)
|
|
return command
|
|
|
|
@agent_action
|
|
def hotkey(self, keys: List[str]):
|
|
"""Press a hotkey combination
|
|
Args:
|
|
keys:List[str] the keys to press in combination in a list format (e.g. ['shift', 'c'])
|
|
"""
|
|
keys = [_normalize_key(k) for k in keys]
|
|
keys = [f"'{key}'" for key in keys]
|
|
command = f"import pyautogui; pyautogui.hotkey({', '.join(keys)}, interval=0.5)"
|
|
return command
|
|
|
|
@agent_action
|
|
def hold_and_press(self, hold_keys: List[str], press_keys: List[str]):
|
|
"""Hold a list of keys and press a list of keys
|
|
Args:
|
|
hold_keys:List[str], list of keys to hold
|
|
press_keys:List[str], list of keys to press in a sequence
|
|
"""
|
|
hold_keys = [_normalize_key(k) for k in hold_keys]
|
|
press_keys = [_normalize_key(k) for k in press_keys]
|
|
|
|
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
|
|
"""
|
|
command = f"import time; time.sleep({time})"
|
|
return command
|
|
|
|
@agent_action
|
|
def done(self):
|
|
"""End the current task with a success"""
|
|
return """DONE"""
|
|
|
|
@agent_action
|
|
def fail(self):
|
|
"""End the current task with a failure"""
|
|
return """FAIL"""
|
|
|
|
|
|
# UIElement Class
|
|
class UIElement:
|
|
def __init__(self, element=None):
|
|
if isinstance(element, pywinauto.application.WindowSpecification):
|
|
self.element = element.wrapper_object()
|
|
else:
|
|
self.element = element # This should be a control wrapper
|
|
|
|
def get_attribute_names(self):
|
|
return list(self.element.element_info.get_properties().keys())
|
|
|
|
def attribute(self, key: str):
|
|
props = self.element.element_info.get_properties()
|
|
return props.get(key, None)
|
|
|
|
def children(self):
|
|
try:
|
|
return [UIElement(child) for child in self.element.children()]
|
|
except Exception as e:
|
|
print(f"Error accessing children: {e}")
|
|
return []
|
|
|
|
def role(self):
|
|
return self.element.element_info.control_type
|
|
|
|
def position(self):
|
|
rect = self.element.rectangle()
|
|
return (rect.left, rect.top)
|
|
|
|
def size(self):
|
|
rect = self.element.rectangle()
|
|
return (rect.width(), rect.height())
|
|
|
|
def title(self):
|
|
return self.element.element_info.name
|
|
|
|
def text(self):
|
|
return self.element.window_text()
|
|
|
|
def isValid(self):
|
|
return self.position() is not None and self.size() is not None
|
|
|
|
def parse(self):
|
|
position = self.position()
|
|
size = self.size()
|
|
return {
|
|
"position": position,
|
|
"size": size,
|
|
"title": self.title(),
|
|
"text": self.text(),
|
|
"role": self.role(),
|
|
}
|
|
|
|
@staticmethod
|
|
def get_current_applications(obs: Dict):
|
|
apps = []
|
|
for proc in psutil.process_iter(["pid", "name"]):
|
|
apps.append(proc.info["name"])
|
|
return apps
|
|
|
|
@staticmethod
|
|
def get_top_app(obs: Dict):
|
|
hwnd = win32gui.GetForegroundWindow()
|
|
_, pid = win32process.GetWindowThreadProcessId(hwnd)
|
|
for proc in psutil.process_iter(["pid", "name"]):
|
|
if proc.info["pid"] == pid:
|
|
return proc.info["name"]
|
|
return None
|
|
|
|
@staticmethod
|
|
def list_apps_in_directories():
|
|
return list_apps_in_directories()
|
|
|
|
@staticmethod
|
|
def systemWideElement():
|
|
desktop = Desktop(backend="uia")
|
|
return UIElement(desktop)
|
|
|
|
def __repr__(self):
|
|
return f"UIElement({self.element})"
|