From f7acffd159d035fe78f8099cb1093fd4cba908ec Mon Sep 17 00:00:00 2001 From: David Piccinini Date: Sun, 17 May 2026 03:03:08 -0300 Subject: [PATCH] feat: indexed UI tree with SoM annotations + improved element resolution - Flat indexed UI element list with numeric IDs and center coordinates - Screenshot annotations: numbered colored boxes (Set of Marks) - click_element resolves by index from cached element list - Cursor crosshair + focused element highlight in screenshots Co-Authored-By: Claude Opus 4.6 --- easybot_windows_agent/main.py | 184 ++++++++++++++++++++++++++-------- 1 file changed, 144 insertions(+), 40 deletions(-) diff --git a/easybot_windows_agent/main.py b/easybot_windows_agent/main.py index 5811d67..11567c1 100644 --- a/easybot_windows_agent/main.py +++ b/easybot_windows_agent/main.py @@ -36,6 +36,40 @@ RECONNECT_DELAY_BASE = 5 RECONNECT_DELAY_MAX = 120 TOOL_TIMEOUT = 60 # max seconds for a single tool call +# ── UI Element Cache (Set of Marks) ───────────────────────────────────────── + +_last_ui_elements: list[dict] = [] + +# Control types considered interactive for SoM indexing +_INTERACTIVE_TYPES = { + "ButtonControl", "EditControl", "CheckBoxControl", "RadioButtonControl", + "ListItemControl", "MenuItemControl", "ComboBoxControl", "TextControl", + "HyperlinkControl", "TabItemControl", "TreeItemControl", "DataItemControl", + "SliderControl", "MenuBarItemControl", "SpinnerControl", "SplitButtonControl", + "ToggleButtonControl", +} + +# Color palette for SoM annotations (by control type category) +_SOM_COLORS: dict[str, str] = { + "ButtonControl": "#FF6B6B", + "EditControl": "#4ECDC4", + "CheckBoxControl": "#45B7D1", + "RadioButtonControl": "#45B7D1", + "ListItemControl": "#96CEB4", + "MenuItemControl": "#FFEAA7", + "ComboBoxControl": "#DDA0DD", + "TextControl": "#98D8C8", + "HyperlinkControl": "#87CEEB", + "TabItemControl": "#F7DC6F", + "TreeItemControl": "#82E0AA", + "DataItemControl": "#AED6F1", + "SliderControl": "#F1948A", + "MenuBarItemControl": "#FDEBD0", + "SpinnerControl": "#D2B4DE", + "SplitButtonControl": "#F5B041", + "ToggleButtonControl": "#FF6B6B", +} + # ── Tool Handlers ──────────────────────────────────────────────────────────── @@ -74,11 +108,12 @@ async def tool_run_command(args: dict[str, Any]) -> dict[str, Any]: async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]: - """Get the accessibility tree of a window or the desktop.""" + """Get indexed interactive UI elements with coordinates (flat list for SoM).""" + global _last_ui_elements import uiautomation as auto window_title = args.get("window") - max_depth = args.get("depth", 3) + max_depth = args.get("depth", 8) try: if window_title: @@ -93,46 +128,50 @@ async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]: if not root: root = auto.GetRootControl() - def walk(ctrl: Any, depth: int, max_d: int) -> list[dict]: - if depth > max_d: - return [] - items = [] + elements: list[dict] = [] + + def collect(ctrl: Any, depth: int) -> None: + if depth > max_depth: + return try: - name = ctrl.Name or "" ctrl_type = ctrl.ControlTypeName or "" - auto_id = ctrl.AutomationId or "" - rect = ctrl.BoundingRectangle - value = "" - try: - value = ctrl.GetValuePattern().Value if ctrl.GetValuePattern() else "" - except Exception: - pass - - item: dict[str, Any] = { - "type": ctrl_type, - "name": name, - } - if auto_id: - item["automation_id"] = auto_id - if value: - item["value"] = value - if rect and rect.width() > 0: - item["rect"] = {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()} - - children = [] - if depth < max_d: - for child in ctrl.GetChildren(): - children.extend(walk(child, depth + 1, max_d)) - if children: - item["children"] = children - - items.append(item) + if ctrl_type in _INTERACTIVE_TYPES: + name = ctrl.Name or "" + auto_id = ctrl.AutomationId or "" + rect = ctrl.BoundingRectangle + if rect and rect.width() > 0 and rect.height() > 0: + cx = rect.left + rect.width() // 2 + cy = rect.top + rect.height() // 2 + elem: dict[str, Any] = { + "id": len(elements), + "type": ctrl_type, + "name": name, + "coords": [cx, cy], + "rect": {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()}, + } + if auto_id: + elem["automation_id"] = auto_id + # Try to get value + try: + vp = ctrl.GetValuePattern() + if vp: + val = vp.Value + if val: + elem["value"] = val + except Exception: + pass + elements.append(elem) + except Exception: + pass + try: + for child in ctrl.GetChildren(): + collect(child, depth + 1) except Exception: pass - return items - tree = walk(root, 0, max_depth) - return {"window": root.Name or "Desktop", "tree": tree} + collect(root, 0) + _last_ui_elements = elements + return {"window": root.Name or "Desktop", "element_count": len(elements), "elements": elements} except Exception as e: return {"error": str(e)} @@ -201,7 +240,8 @@ def _get_ui_feedback() -> dict[str, Any]: async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]: - """Click a UI element by name, automation_id, or coordinates.""" + """Click a UI element by index (from UI tree), name, automation_id, or coordinates.""" + index = args.get("index") name = args.get("name") auto_id = args.get("automation_id") x = args.get("x") @@ -210,6 +250,16 @@ async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]: double = args.get("double", False) try: + # Resolve index to coordinates from cached UI tree + if index is not None: + if not _last_ui_elements: + return {"error": "No UI tree cached. Call get_ui_tree first."} + if index < 0 or index >= len(_last_ui_elements): + return {"error": f"Index {index} out of range (0-{len(_last_ui_elements) - 1})"} + elem = _last_ui_elements[index] + x, y = elem["coords"] + logger.debug("Index %d resolved to (%d, %d) — %s", index, x, y, elem.get("name", "")) + if x is not None and y is not None: import pyautogui if double: @@ -295,11 +345,13 @@ async def tool_press_keys(args: dict[str, Any]) -> dict[str, Any]: async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]: - """Take a screenshot with UI annotations: cursor marker + focused element highlight.""" - from PIL import ImageGrab, ImageDraw + """Take a screenshot with Set of Marks (SoM) annotations from the UI tree.""" + from PIL import ImageGrab, ImageDraw, ImageFont try: region = args.get("region") + annotate = args.get("annotate", True) + if region: bbox = (region["x"], region["y"], region["x"] + region["w"], region["y"] + region["h"]) img = ImageGrab.grab(bbox=bbox) @@ -343,6 +395,58 @@ async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]: except Exception: pass + # 3. Set of Marks (SoM): draw indexed labels on interactive elements + if annotate and _last_ui_elements: + # Try to load a small font for labels + try: + font = ImageFont.truetype("arial.ttf", 11) + except Exception: + try: + font = ImageFont.truetype("C:\\Windows\\Fonts\\arial.ttf", 11) + except Exception: + font = ImageFont.load_default() + + region_ox = region["x"] if region else 0 + region_oy = region["y"] if region else 0 + + for elem in _last_ui_elements: + try: + r = elem["rect"] + ex = r["x"] - region_ox + ey = r["y"] - region_oy + ew = r["w"] + eh = r["h"] + + # Skip elements outside visible area + if ex + ew < 0 or ey + eh < 0 or ex > img.width or ey > img.height: + continue + + color = _SOM_COLORS.get(elem["type"], "#AAAAAA") + idx = elem["id"] + + # Draw bounding box (thin outline) + draw.rectangle([ex, ey, ex + ew, ey + eh], outline=color, width=1) + + # Draw index label (small colored background with white text) + label_text = str(idx) + # Get text size + text_bbox = draw.textbbox((0, 0), label_text, font=font) + tw = text_bbox[2] - text_bbox[0] + 4 + th = text_bbox[3] - text_bbox[1] + 2 + + # Position label at top-left of element + lx = ex + ly = ey - th - 1 + if ly < 0: + ly = ey # put inside if no room above + + # Draw label background + draw.rectangle([lx, ly, lx + tw, ly + th], fill=color) + # Draw label text + draw.text((lx + 2, ly + 1), label_text, fill="white", font=font) + except Exception: + continue + # Resize for efficiency (max 1280px wide) if img.width > 1280: ratio = 1280 / img.width