chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,242 @@
|
||||
<h1 align="center">
|
||||
<img src="../../images/agent_s.png" alt="Logo" style="vertical-align:middle" width="60"> Agent S:
|
||||
<small>Using Computers Like a Human</small>
|
||||
</h1>
|
||||
|
||||
<p align="center">
|
||||
🌐 <a href="https://www.simular.ai/agent-s">[Website]</a>
|
||||
📄 <a href="https://arxiv.org/abs/2410.08164">[Paper]</a>
|
||||
🎥 <a href="https://www.youtube.com/watch?v=OBDE3Knte0g">[Video]</a>
|
||||
🗨️ <a href="https://discord.gg/E2XfsK9fPV">[Discord]</a>
|
||||
</p>
|
||||
|
||||
## 🥳 Updates
|
||||
- [x] **2025/01/22**: The [Agent S paper](https://arxiv.org/abs/2410.08164) is accepted to ICLR 2025!
|
||||
- [x] **2025/01/21**: Released v0.1.2 of [gui-agents](https://github.com/simular-ai/Agent-S) library, with support for Linux and Windows!
|
||||
- [x] **2024/12/05**: Released v0.1.0 of [gui-agents](https://github.com/simular-ai/Agent-S) library, allowing you to use Agent-S for Mac, OSWorld, and WindowsAgentArena with ease!
|
||||
- [x] **2024/10/10**: Released [Agent S paper](https://arxiv.org/abs/2410.08164) and codebase!
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [💡 Introduction](#-introduction)
|
||||
2. [🎯 Current Results](#-current-results)
|
||||
3. [🛠️ Installation](#%EF%B8%8F-installation)
|
||||
4. [🚀 Usage](#-usage)
|
||||
5. [🙌 Contributors](#-contributors)
|
||||
6. [💬 Citation](#-citation)
|
||||
|
||||
## 💡 Introduction
|
||||
|
||||
<p align="center">
|
||||
<img src="../../images/teaser.png" width="800">
|
||||
</p>
|
||||
|
||||
Welcome to **Agent S**, an open-source framework designed to enable autonomous interaction with computers through Agent-Computer Interface. Our mission is to build intelligent GUI agents that can learn from past experiences and perform complex tasks autonomously on your computer.
|
||||
|
||||
Whether you're interested in AI, automation, or contributing to cutting-edge agent-based systems, we're excited to have you here!
|
||||
|
||||
## 🎯 Current Results
|
||||
|
||||
<p align="center">
|
||||
<img src="../../images/results.png" width="600">
|
||||
<br>
|
||||
Results of Successful Rate (%) on the OSWorld full test set of all 369 test examples using Image + Accessibility Tree input.
|
||||
</p>
|
||||
|
||||
|
||||
## 🛠️ Installation & Setup
|
||||
|
||||
> ❗**Warning**❗: If you are on a Linux machine, creating a `conda` environment will interfere with `pyatspi`. As of now, there's no clean solution for this issue. Proceed through the installation without using `conda` or any virtual environment.
|
||||
|
||||
Clone the repository:
|
||||
```
|
||||
git clone https://github.com/simular-ai/Agent-S.git
|
||||
```
|
||||
|
||||
Install the gui-agents package:
|
||||
```
|
||||
pip install gui-agents
|
||||
```
|
||||
|
||||
Set your LLM API Keys and other environment variables. You can do this by adding the following line to your .bashrc (Linux), or .zshrc (MacOS) file.
|
||||
|
||||
```
|
||||
export OPENAI_API_KEY=<YOUR_API_KEY>
|
||||
```
|
||||
|
||||
Alternatively, you can set the environment variable in your Python script:
|
||||
|
||||
```
|
||||
import os
|
||||
os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>"
|
||||
```
|
||||
|
||||
We also support Azure OpenAI, Anthropic, and vLLM inference. For more information refer to [../../models.md](models.md).
|
||||
|
||||
### Setup Retrieval from Web using Perplexica
|
||||
Agent S works best with web-knowledge retrieval. To enable this feature, you need to setup Perplexica:
|
||||
|
||||
1. Ensure Docker Desktop is installed and running on your system.
|
||||
|
||||
2. Navigate to the directory containing the project files.
|
||||
|
||||
```bash
|
||||
cd Perplexica
|
||||
git submodule update --init
|
||||
```
|
||||
|
||||
3. Rename the `sample.config.toml` file to `config.toml`. For Docker setups, you need only fill in the following fields:
|
||||
|
||||
- `OPENAI`: Your OpenAI API key. **You only need to fill this if you wish to use OpenAI's models**.
|
||||
- `OLLAMA`: Your Ollama API URL. You should enter it as `http://host.docker.internal:PORT_NUMBER`. If you installed Ollama on port 11434, use `http://host.docker.internal:11434`. For other ports, adjust accordingly. **You need to fill this if you wish to use Ollama's models instead of OpenAI's**.
|
||||
- `GROQ`: Your Groq API key. **You only need to fill this if you wish to use Groq's hosted models**.
|
||||
- `ANTHROPIC`: Your Anthropic API key. **You only need to fill this if you wish to use Anthropic models**.
|
||||
|
||||
**Note**: You can change these after starting Perplexica from the settings dialog.
|
||||
|
||||
- `SIMILARITY_MEASURE`: The similarity measure to use (This is filled by default; you can leave it as is if you are unsure about it.)
|
||||
|
||||
4. Ensure you are in the directory containing the `docker-compose.yaml` file and execute:
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
5. Next, export your Perplexica URL. This URL is used to interact with the Perplexica API backend. The port is given by the `config.toml` in your Perplexica directory.
|
||||
|
||||
```bash
|
||||
export PERPLEXICA_URL=http://localhost:{port}/api/search
|
||||
```
|
||||
|
||||
6. Our implementation of Agent S incorporates the Perplexica API to integrate a search engine capability, which allows for a more convenient and responsive user experience. If you want to tailor the API to your settings and specific requirements, you may modify the URL and the message of request parameters in `agent_s/query_perplexica.py`. For a comprehensive guide on configuring the Perplexica API, please refer to [Perplexica Search API Documentation](https://github.com/ItzCrazyKns/Perplexica/blob/master/docs/API/SEARCH.md)
|
||||
|
||||
For a more detailed setup and usage guide, please refer to the [Perplexica Repository](https://github.com/ItzCrazyKns/Perplexica.git).
|
||||
|
||||
### Setup Paddle-OCR Server
|
||||
|
||||
Switch to a new terminal where you will run Agent S. Set the OCR_SERVER_ADDRESS environment variable as shown below. For a better experience, add the following line directly to your .bashrc (Linux), or .zshrc (MacOS) file.
|
||||
|
||||
```
|
||||
export OCR_SERVER_ADDRESS=http://localhost:8000/ocr/
|
||||
```
|
||||
|
||||
Run the ocr_server.py file code to use OCR-based bounding boxes.
|
||||
|
||||
```
|
||||
cd Agent-S
|
||||
python gui_agents/utils/ocr_server.py
|
||||
```
|
||||
|
||||
You can change the server address by editing the address in [gui_agents/s1/utils/ocr_server.py](utils/ocr_server.py) file.
|
||||
|
||||
|
||||
> ❗**Warning**❗: The agent will directly run python code to control your computer. Please use with care.
|
||||
|
||||
## 🚀 Usage
|
||||
|
||||
### CLI
|
||||
|
||||
Run agent_s on your computer using:
|
||||
```
|
||||
agent_s1 --model gpt-4o
|
||||
```
|
||||
This will show a user query prompt where you can enter your query and interact with Agent S. You can use any model from the list of supported models in [models.md](../../models.md).
|
||||
|
||||
### `gui_agents` SDK
|
||||
|
||||
To deploy Agent S on MacOS or Windows:
|
||||
|
||||
```
|
||||
import pyautogui
|
||||
import io
|
||||
from gui_agents.core.AgentS import GraphSearchAgent
|
||||
import platform
|
||||
|
||||
if platform.system() == "Darwin":
|
||||
from gui_agents.aci.MacOSACI import MacOSACI, UIElement
|
||||
grounding_agent = MacOSACI()
|
||||
elif platform.system() == "Windows":
|
||||
from gui_agents.aci.WindowsOSACI import WindowsACI, UIElement
|
||||
grounding_agent = WindowsACI()
|
||||
elif platform.system() == "Linux":
|
||||
from gui_agents.aci.LinuxOSACI import LinuxACI, UIElement
|
||||
grounding_agent = LinuxACI()
|
||||
else:
|
||||
raise ValueError("Unsupported platform")
|
||||
|
||||
engine_params = {
|
||||
"engine_type": "openai",
|
||||
"model": "gpt-4o",
|
||||
}
|
||||
|
||||
agent = GraphSearchAgent(
|
||||
engine_params,
|
||||
grounding_agent,
|
||||
platform="ubuntu", # "macos", "windows"
|
||||
action_space="pyautogui",
|
||||
observation_type="mixed",
|
||||
search_engine="Perplexica"
|
||||
)
|
||||
|
||||
# Get screenshot.
|
||||
screenshot = pyautogui.screenshot()
|
||||
buffered = io.BytesIO()
|
||||
screenshot.save(buffered, format="PNG")
|
||||
screenshot_bytes = buffered.getvalue()
|
||||
|
||||
# Get accessibility tree.
|
||||
acc_tree = UIElement.systemWideElement()
|
||||
|
||||
obs = {
|
||||
"screenshot": screenshot_bytes,
|
||||
"accessibility_tree": acc_tree,
|
||||
}
|
||||
|
||||
instruction = "Close VS Code"
|
||||
info, action = agent.predict(instruction=instruction, observation=obs)
|
||||
|
||||
exec(action[0])
|
||||
```
|
||||
|
||||
Refer to `cli_app.py` for more details on how the inference loop works.
|
||||
|
||||
#### Downloading the Knowledege Base
|
||||
|
||||
Agent S2 uses a knowledge base that continually updates with new knowledge during inference. The knowledge base is initially downloaded when initializing `GraphSearchAgent`. The knowledge base is stored as assets under our [GitHub Releases](https://github.com/simular-ai/Agent-S/releases). The `GraphSearchAgent` initialization will only download the knowledge base for your specified platform and agent version (e.g s1, s2). If you'd like to download the knowledge base programmatically, you can use the following code:
|
||||
|
||||
```
|
||||
download_kb_data(
|
||||
version="s2",
|
||||
release_tag="v0.2.2",
|
||||
download_dir="kb_data",
|
||||
platform="linux" # "darwin", "windows"
|
||||
)
|
||||
```
|
||||
|
||||
This will download Agent S2's knowledge base for Linux from release tag `v0.2.2` to the `kb_data` directory. Refer to our [GitHub Releases](https://github.com/simular-ai/Agent-S/releases) or release tags that include the knowledge bases.
|
||||
|
||||
### OSWorld
|
||||
|
||||
To deploy Agent S in OSWorld, follow the [OSWorld Deployment instructions](OSWorld.md).
|
||||
|
||||
### WindowsAgentArena
|
||||
|
||||
To deploy Agent S in WindowsAgentArena, follow the [WindowsAgentArena Deployment instructions](WindowsAgentArena.md).
|
||||
|
||||
## 🙌 Contributors
|
||||
|
||||
We’re grateful to all the [amazing people](https://github.com/simular-ai/Agent-S/graphs/contributors) who have contributed to this project. Thank you! 🙏
|
||||
|
||||
## 💬 Citation
|
||||
```
|
||||
@misc{agashe2024agentsopenagentic,
|
||||
title={Agent S: An Open Agentic Framework that Uses Computers Like a Human},
|
||||
author={Saaket Agashe and Jiuzhou Han and Shuyu Gan and Jiachen Yang and Ang Li and Xin Eric Wang},
|
||||
year={2024},
|
||||
eprint={2410.08164},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.AI},
|
||||
url={https://arxiv.org/abs/2410.08164},
|
||||
}
|
||||
```
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
## Deploying Agent-S in WindowsAgentArena
|
||||
> ⚠️ **Warning**: The refactored code has not be fully tested on WindowsAgentArena. To reproduce the results on WindowsAgentArena, please use commit 496a9fa of this repository.
|
||||
|
||||
1. To use the Agent S with WindowsAgentArena, follows the setup instructions at: https://github.com/microsoft/WindowsAgentArena.git. **Please use the development mode while preparing the image and running the client as instructed in https://github.com/microsoft/WindowsAgentArena/blob/main/docs/Development-Tips.md.**
|
||||
|
||||
2. To deploy our agent in the WindowsAgentArena, copy the agent_s folder in this repository to `WindowsAgentArena/src/win-arena-container/client/mm_agents`.
|
||||
|
||||
3. Change the name of the GraphSearchAgent.py file to agent.py to conform to the WindowsAgentArena Setup.
|
||||
|
||||
4. Copy the ocr_server.py file to client/folder `WindowsAgentArena/src/win-arena-container/client` folder
|
||||
|
||||
```
|
||||
cd WindowsAgentArena/src/win-arena-container/client
|
||||
cp mm_agents/agent_s/ocr_server.py .
|
||||
```
|
||||
|
||||
5. Update the `start_client.sh` file in `WindowsAgentArena/src/win-arena-container` by adding the following line before Running the agent on line 75.
|
||||
|
||||
```
|
||||
python ocr_server.py &
|
||||
```
|
||||
|
||||
6. In the `src/win-arena-container/client/run.py` file import Agent S
|
||||
```
|
||||
from mm_agents.agent_s.agent import GraphSearchAgent
|
||||
```
|
||||
|
||||
7. In the `src/win-arena-container/client/run.py` file, instantiate Agent S by adding the following lines after line 187 where the if condition for NAVI agent ends
|
||||
|
||||
```python
|
||||
elif cfg_args["agent_name"] == "agent_s":
|
||||
if cfg_args["som_origin"] in ["a11y"]:
|
||||
som_config = None
|
||||
elif cfg_args["som_origin"] in ["oss", "mixed-oss"]:
|
||||
som_config = {
|
||||
"pipeline": ["webparse", "groundingdino", "ocr"],
|
||||
"groundingdino": {
|
||||
"prompts": ["icon", "image"]
|
||||
},
|
||||
"ocr": {
|
||||
"class_name": "TesseractOCR"
|
||||
},
|
||||
"webparse": {
|
||||
"cdp_url": f"http://{args.emulator_ip}:9222"
|
||||
}
|
||||
}
|
||||
if args.model.startswith("claude"):
|
||||
engine_type = "anthropic"
|
||||
elif args.model.startswith("gpt"):
|
||||
engine_type = "openai"
|
||||
else:
|
||||
engine_type = "vllm"
|
||||
|
||||
engine_params = {
|
||||
"engine_type": engine_type,
|
||||
"model": args.model,
|
||||
}
|
||||
agent = GraphSearchAgent(
|
||||
engine_params=engine_params,
|
||||
experiment_type='windowsAgentArena',
|
||||
temperature=args.temperature
|
||||
)
|
||||
```
|
||||
|
||||
8. Run Agent S on WindowsAgentArena by changing the following parameters in the `scripts/run-local.sh` file
|
||||
|
||||
```
|
||||
agent="agent_s"
|
||||
model="gpt-4o"
|
||||
```
|
||||
@@ -0,0 +1,36 @@
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
def agent_action(func):
|
||||
func.is_agent_action = True
|
||||
return func
|
||||
|
||||
|
||||
class ACI:
|
||||
def __init__(self, top_app_only: bool = True, ocr: bool = False):
|
||||
self.top_app_only = top_app_only
|
||||
self.ocr = ocr
|
||||
self.index_out_of_range_flag = False
|
||||
self.notes: List[str] = []
|
||||
self.clipboard = ""
|
||||
self.nodes: List[Any] = []
|
||||
|
||||
def get_active_apps(self, obs: Dict) -> List[str]:
|
||||
pass
|
||||
|
||||
def get_top_app(self):
|
||||
pass
|
||||
|
||||
def preserve_nodes(self, tree: Any, exclude_roles: set = None) -> List[Dict]:
|
||||
pass
|
||||
|
||||
def linearize_and_annotate_tree(
|
||||
self, obs: Dict, show_all_elements: bool = False
|
||||
) -> str:
|
||||
pass
|
||||
|
||||
def find_element(self, element_id: int) -> Dict:
|
||||
pass
|
||||
@@ -0,0 +1,846 @@
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from typing import Dict, List, Optional, Tuple, Any, Sequence
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI
|
||||
from gui_agents.s1.utils.common_utils import box_iou
|
||||
|
||||
import platform
|
||||
|
||||
if platform.system() == "Linux":
|
||||
import pyatspi
|
||||
from pyatspi import Accessible, StateType, STATE_SHOWING
|
||||
from pyatspi import Action as ATAction
|
||||
from pyatspi import Component # , Document
|
||||
from pyatspi import Text as ATText
|
||||
from pyatspi import Value as ATValue
|
||||
|
||||
from pyatspi import Accessible, StateType
|
||||
from lxml.etree import _Element
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
import lxml.etree
|
||||
import concurrent.futures
|
||||
|
||||
_accessibility_ns_map_ubuntu = {
|
||||
"st": "https://accessibility.ubuntu.example.org/ns/state",
|
||||
"attr": "https://accessibility.ubuntu.example.org/ns/attributes",
|
||||
"cp": "https://accessibility.ubuntu.example.org/ns/component",
|
||||
"doc": "https://accessibility.ubuntu.example.org/ns/document",
|
||||
"docattr": "https://accessibility.ubuntu.example.org/ns/document/attributes",
|
||||
"txt": "https://accessibility.ubuntu.example.org/ns/text",
|
||||
"val": "https://accessibility.ubuntu.example.org/ns/value",
|
||||
"act": "https://accessibility.ubuntu.example.org/ns/action",
|
||||
}
|
||||
|
||||
MAX_DEPTH = 50
|
||||
MAX_WIDTH = 1024
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
# Agent action decorator
|
||||
def agent_action(func):
|
||||
func.is_agent_action = True
|
||||
return func
|
||||
|
||||
|
||||
class LinuxACI(ACI):
|
||||
def __init__(self, top_app=None, vm_version="new", top_app_only=True, ocr=True):
|
||||
self.active_apps = set()
|
||||
self.top_app = top_app
|
||||
self.top_app_only = (
|
||||
top_app_only # Only include top app in the accessibility tree
|
||||
)
|
||||
self.ocr = ocr
|
||||
self.index_out_of_range_flag = False
|
||||
self.app_setup_code = 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'])
|
||||
"""
|
||||
|
||||
self.top_active_app = None
|
||||
self.notes = []
|
||||
self.clipboard = ""
|
||||
|
||||
# TODO: this is terrible, fix this
|
||||
global state_ns, component_ns, attributes_ns, value_ns
|
||||
if vm_version == "old":
|
||||
|
||||
state_ns = "uri:deskat:state.at-spi.gnome.org"
|
||||
component_ns = "uri:deskat:component.at-spi.gnome.org"
|
||||
else:
|
||||
attributes_ns = "https://accessibility.windows.example.org/ns/attributes"
|
||||
state_ns = "https://accessibility.ubuntu.example.org/ns/state"
|
||||
component_ns = "https://accessibility.ubuntu.example.org/ns/component"
|
||||
value_ns = "https://accessibility.ubuntu.example.org/ns/value"
|
||||
|
||||
def get_active_apps(self, obs: Dict) -> List[str]:
|
||||
tree = ET.ElementTree(ET.fromstring(obs["accessibility_tree"]))
|
||||
apps = []
|
||||
exclude_list = ["gjs", "gnome-shell"]
|
||||
for node in tree.iter():
|
||||
# Keep applications and only those which have children
|
||||
if (
|
||||
node.tag.endswith("application")
|
||||
and list(node)
|
||||
and node.attrib.get("name", "") not in exclude_list
|
||||
):
|
||||
apps.append(node.attrib.get("name", "").replace("\\", ""))
|
||||
return apps
|
||||
|
||||
def check_new_apps(self, old_apps, new_apps):
|
||||
return new_apps - old_apps
|
||||
|
||||
def get_top_app(self, obs):
|
||||
return self.top_app
|
||||
|
||||
def find_active_applications(self, tree):
|
||||
# names of applications to keep TODO: soffice is a single application with all the isntances like impress, calc etc. being frames this will need to be dealt with separately
|
||||
to_keep = ["gnome-shell"]
|
||||
apps_with_active_tag = []
|
||||
for application in list(tree.getroot()):
|
||||
app_name = application.attrib.get("name")
|
||||
for frame in application:
|
||||
is_active = frame.attrib.get("{{{:}}}active".format(state_ns), "false")
|
||||
if is_active == "true":
|
||||
apps_with_active_tag.append(app_name)
|
||||
if apps_with_active_tag:
|
||||
to_keep.append(apps_with_active_tag[-1])
|
||||
return to_keep
|
||||
|
||||
def filter_active_app(self, tree):
|
||||
for application in list(tree.getroot()):
|
||||
app_name = application.attrib.get("name")
|
||||
for frame in application:
|
||||
is_active = frame.attrib.get("{{{:}}}active".format(state_ns), "false")
|
||||
if is_active == "true":
|
||||
return app_name
|
||||
return None
|
||||
|
||||
def filter_nodes(self, tree, show_all=False):
|
||||
# created and populate a preserved nodes list which filters out unnecessary elements and keeps only those elements which are currently showing on the screen
|
||||
# TODO: include offscreen elements and then scroll to them before clicking
|
||||
preserved_nodes = []
|
||||
exclude_tags = ["panel", "window", "filler", "frame", "separator", "scroll-bar"]
|
||||
|
||||
for node in tree.iter():
|
||||
if node.tag not in exclude_tags:
|
||||
if show_all:
|
||||
if node.attrib.get(f"{{{state_ns}}}visible") == "true":
|
||||
coords: Tuple[int, int] = eval(
|
||||
node.get(
|
||||
"{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
)
|
||||
)
|
||||
if coords[0] >= 0 and coords[1] >= 0:
|
||||
preserved_nodes.append(node)
|
||||
# if show_all is false, only show elements that are currently showing on screen
|
||||
else:
|
||||
if node.attrib.get(f"{{{state_ns}}}showing") == "true":
|
||||
coords: Tuple[int, int] = eval(
|
||||
node.get(
|
||||
"{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
)
|
||||
)
|
||||
|
||||
if coords[0] >= 0 and coords[1] >= 0:
|
||||
preserved_nodes.append(node)
|
||||
|
||||
return preserved_nodes
|
||||
|
||||
def linearize_tree(self, preserved_nodes):
|
||||
# TODO: Run an ablation to check if class and desc
|
||||
# linearized_accessibility_tree = ["id\ttag\tname\ttext\tclass\tdescription"]
|
||||
linearized_accessibility_tree = ["id\ttag\tname\ttext"]
|
||||
for idx, node in enumerate(preserved_nodes):
|
||||
if node.text:
|
||||
text = (
|
||||
node.text
|
||||
if '"' not in node.text
|
||||
else '"{:}"'.format(node.text.replace('"', '""'))
|
||||
)
|
||||
else:
|
||||
text = '""'
|
||||
|
||||
linearized_accessibility_tree.append(
|
||||
"{:}\t{:}\t{:}\t{:}".format(
|
||||
idx,
|
||||
node.tag,
|
||||
node.get("name", ""),
|
||||
text,
|
||||
# node.get("{{{:}}}class".format(attributes_ns), ""),
|
||||
# node.get("{{{:}}}description".format(attributes_ns), ""),
|
||||
)
|
||||
)
|
||||
|
||||
# returning list of linearized elements
|
||||
return linearized_accessibility_tree
|
||||
|
||||
def extract_elements_from_screenshot(self, screenshot) -> Dict:
|
||||
"""Uses paddle-ocr to extract elements with text from the screenshot. The elements will be added to the linearized accessibility tree downstream"""
|
||||
|
||||
# Convert screenshot to PIL image
|
||||
def send_image_to_ocr(screenshot) -> Dict:
|
||||
|
||||
url = os.environ.get("OCR_SERVER_ADDRESS", "")
|
||||
if url == "":
|
||||
raise Exception("OCR SERVER ADDRESS NOT SET")
|
||||
encoded_screenshot = base64.b64encode(screenshot).decode("utf-8")
|
||||
data = {"img_bytes": encoded_screenshot}
|
||||
print("Getting OCR response")
|
||||
ocr_start = time.time()
|
||||
response = requests.post(url, json=data)
|
||||
print("Got OCR response in", time.time() - ocr_start)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
return {
|
||||
"error": f"Request failed with status code {response.status_code}",
|
||||
"results": [],
|
||||
}
|
||||
|
||||
return send_image_to_ocr(screenshot)["results"]
|
||||
|
||||
def add_ocr_elements(
|
||||
self, screenshot, linearized_accessibility_tree, preserved_nodes
|
||||
):
|
||||
# Get the bounding boxes of the elements in the linearized accessibility tree
|
||||
tree_bboxes = []
|
||||
for node in preserved_nodes:
|
||||
coordinates: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}size".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
tree_bboxes.append(
|
||||
[
|
||||
coordinates[0],
|
||||
coordinates[1],
|
||||
coordinates[0] + sizes[0],
|
||||
coordinates[1] + sizes[1],
|
||||
]
|
||||
)
|
||||
|
||||
# Use OCR to found boxes that might be missing from the accessibility tree
|
||||
try:
|
||||
ocr_bboxes = self.extract_elements_from_screenshot(screenshot)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
ocr_bboxes = []
|
||||
else:
|
||||
# Check for intersection over union between the existing atree bounding boxes and the ocr bounding boxes, if ocr bounding boxes are new add them to the linearized accesibility tree
|
||||
if (
|
||||
len(ocr_bboxes) > 0
|
||||
): # Only check IOUs and add if there are any bounding boxes returned by the ocr module
|
||||
preserved_nodes_index = len(preserved_nodes)
|
||||
for ind, (i, content, box) in enumerate(ocr_bboxes):
|
||||
# x1, y1, x2, y2 = int(box.get('left', 0)), int(box['top']), int(), int(box['bottom'])
|
||||
(
|
||||
x1,
|
||||
y1,
|
||||
x2,
|
||||
y2,
|
||||
) = (
|
||||
int(box.get("left", 0)),
|
||||
int(box.get("top", 0)),
|
||||
int(box.get("right", 0)),
|
||||
int(box.get("bottom", 0)),
|
||||
)
|
||||
iou = box_iou(
|
||||
np.array(tree_bboxes, dtype=np.float32),
|
||||
np.array([[x1, y1, x2, y2]], dtype=np.float32),
|
||||
).flatten()
|
||||
|
||||
if max(iou) < 0.1:
|
||||
# Add the element to the linearized accessibility tree
|
||||
# TODO: ocr detected elements should be classified for their tag, currently set to push button for the agent to think they are interactable
|
||||
linearized_accessibility_tree.append(
|
||||
f"{preserved_nodes_index}\tpush-button\t\t{content}\t\t"
|
||||
)
|
||||
|
||||
# add to preserved node with the component_ns prefix node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
node = ET.Element(
|
||||
"ocr_node",
|
||||
attrib={
|
||||
"text": content,
|
||||
"{{{}}}screencoord".format(
|
||||
component_ns
|
||||
): "({},{})".format(x1, y1),
|
||||
"{{{}}}size".format(component_ns): "({},{})".format(
|
||||
x2 - x1, y2 - y1
|
||||
),
|
||||
},
|
||||
)
|
||||
preserved_nodes.append(node)
|
||||
preserved_nodes_index += 1
|
||||
|
||||
return linearized_accessibility_tree, preserved_nodes
|
||||
|
||||
def linearize_and_annotate_tree(self, obs, show_all=False):
|
||||
accessibility_tree = obs["accessibility_tree"]
|
||||
screenshot = obs["screenshot"]
|
||||
|
||||
# convert the accessibility tree from a string representation to an xml tree
|
||||
tree = ET.ElementTree(ET.fromstring(accessibility_tree))
|
||||
|
||||
# Get the applications to keep based on the active applications
|
||||
to_keep = self.find_active_applications(tree)
|
||||
self.top_app = to_keep[-1]
|
||||
|
||||
# Remove applications which are not included in the to_keep list
|
||||
if not show_all:
|
||||
for application in list(tree.getroot()):
|
||||
if application.attrib.get("name", "") not in to_keep:
|
||||
tree.getroot().remove(application)
|
||||
|
||||
# Save tree for debugging
|
||||
with open("tree_raw.xml", "wb") as file:
|
||||
tree.write(file, encoding="utf-8", xml_declaration=True)
|
||||
|
||||
# Filter out filler elements and overlapping elements
|
||||
preserved_nodes = self.filter_nodes(tree, show_all)
|
||||
|
||||
assert len(preserved_nodes) > 0
|
||||
|
||||
# Linearize the tree as tsv
|
||||
linearized_accessibility_tree = self.linearize_tree(preserved_nodes)
|
||||
|
||||
# Add OCR elements to the linearized accessibility tree to account for elements that are not in the accessibility tree
|
||||
if self.ocr:
|
||||
linearized_accessibility_tree, preserved_nodes = self.add_ocr_elements(
|
||||
screenshot, linearized_accessibility_tree, preserved_nodes
|
||||
)
|
||||
|
||||
# Convert accessibility tree to a string
|
||||
linearized_accessibility_tree = "\n".join(linearized_accessibility_tree)
|
||||
|
||||
# TODO: side-effect, set in separate functions
|
||||
self.nodes = preserved_nodes
|
||||
|
||||
return linearized_accessibility_tree
|
||||
|
||||
def find_element(self, element_id):
|
||||
try:
|
||||
selected_element = self.nodes[int(element_id)]
|
||||
except:
|
||||
print("The index of the selected element was out of range.")
|
||||
selected_element = self.nodes[0]
|
||||
self.index_out_of_range_flag = True
|
||||
return selected_element
|
||||
|
||||
@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] = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}size".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
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
|
||||
"""
|
||||
return self.app_setup_code.replace("APP_NAME", app_code)
|
||||
|
||||
@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:
|
||||
# Use the provided element_id or default to None
|
||||
node = self.find_element(element_id) if element_id is not None else None
|
||||
except:
|
||||
node = None
|
||||
|
||||
if node is not None:
|
||||
# If a node is found, retrieve its coordinates and size
|
||||
coordinates = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
# Start typing at the center of the element
|
||||
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, 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 = eval(
|
||||
node1.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes1 = eval(node1.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
coordinates2 = eval(
|
||||
node2.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes2 = eval(node2.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x1 = coordinates1[0] + sizes1[0] // 2
|
||||
y1 = coordinates1[1] + sizes1[1] // 2
|
||||
|
||||
x2 = coordinates2[0] + sizes2[0] // 2
|
||||
y2 = coordinates2[1] + sizes2[1] // 2
|
||||
|
||||
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 scroll(self, element_id: int, clicks: int):
|
||||
"""Scroll the element in the specified direction
|
||||
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)
|
||||
# print(node.attrib)
|
||||
coordinates = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
return (
|
||||
f"import pyautogui; pyautogui.moveTo({x}, {y}); pyautogui.scroll({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"""
|
||||
return """DONE"""
|
||||
|
||||
@agent_action
|
||||
def fail(self):
|
||||
"""End the current task with a failure"""
|
||||
return """FAIL"""
|
||||
|
||||
|
||||
def _create_atspi_node(
|
||||
node: Accessible, depth: int = 0, flag: Optional[str] = None
|
||||
) -> _Element:
|
||||
node_name = node.name
|
||||
attribute_dict: Dict[str, Any] = {"name": node_name}
|
||||
|
||||
# States
|
||||
states: List[StateType] = node.getState().get_states()
|
||||
for st in states:
|
||||
state_name: str = StateType._enum_lookup[st]
|
||||
state_name: str = state_name.split("_", maxsplit=1)[1].lower()
|
||||
if len(state_name) == 0:
|
||||
continue
|
||||
attribute_dict[
|
||||
"{{{:}}}{:}".format(_accessibility_ns_map_ubuntu["st"], state_name)
|
||||
] = "true"
|
||||
|
||||
# Attributes
|
||||
attributes: Dict[str, str] = node.get_attributes()
|
||||
for attribute_name, attribute_value in attributes.items():
|
||||
if len(attribute_name) == 0:
|
||||
continue
|
||||
attribute_dict[
|
||||
"{{{:}}}{:}".format(_accessibility_ns_map_ubuntu["attr"], attribute_name)
|
||||
] = attribute_value
|
||||
|
||||
# Component
|
||||
if (
|
||||
attribute_dict.get(
|
||||
"{{{:}}}visible".format(_accessibility_ns_map_ubuntu["st"]), "false"
|
||||
)
|
||||
== "true"
|
||||
and attribute_dict.get(
|
||||
"{{{:}}}showing".format(_accessibility_ns_map_ubuntu["st"]), "false"
|
||||
)
|
||||
== "true"
|
||||
):
|
||||
try:
|
||||
component: Component = node.queryComponent()
|
||||
except NotImplementedError:
|
||||
pass
|
||||
else:
|
||||
bbox: Sequence[int] = component.getExtents(pyatspi.XY_SCREEN)
|
||||
attribute_dict[
|
||||
"{{{:}}}screencoord".format(_accessibility_ns_map_ubuntu["cp"])
|
||||
] = str(tuple(bbox[0:2]))
|
||||
attribute_dict["{{{:}}}size".format(_accessibility_ns_map_ubuntu["cp"])] = (
|
||||
str(tuple(bbox[2:]))
|
||||
)
|
||||
|
||||
text = ""
|
||||
# Text
|
||||
try:
|
||||
text_obj: ATText = node.queryText()
|
||||
# only text shown on current screen is available
|
||||
# attribute_dict["txt:text"] = text_obj.getText(0, text_obj.characterCount)
|
||||
text: str = text_obj.getText(0, text_obj.characterCount)
|
||||
# if flag=="thunderbird":
|
||||
# appeared in thunderbird (uFFFC) (not only in thunderbird), "Object
|
||||
# Replacement Character" in Unicode, "used as placeholder in text for
|
||||
# an otherwise unspecified object; uFFFD is another "Replacement
|
||||
# Character", just in case
|
||||
text = text.replace("\ufffc", "").replace("\ufffd", "")
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
# Image, Selection, Value, Action
|
||||
try:
|
||||
node.queryImage()
|
||||
attribute_dict["image"] = "true"
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
try:
|
||||
node.querySelection()
|
||||
attribute_dict["selection"] = "true"
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
try:
|
||||
value: ATValue = node.queryValue()
|
||||
value_key = f"{{{_accessibility_ns_map_ubuntu['val']}}}"
|
||||
|
||||
for attr_name, attr_func in [
|
||||
("value", lambda: value.currentValue),
|
||||
("min", lambda: value.minimumValue),
|
||||
("max", lambda: value.maximumValue),
|
||||
("step", lambda: value.minimumIncrement),
|
||||
]:
|
||||
try:
|
||||
attribute_dict[f"{value_key}{attr_name}"] = str(attr_func())
|
||||
except:
|
||||
pass
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
try:
|
||||
action: ATAction = node.queryAction()
|
||||
for i in range(action.nActions):
|
||||
action_name: str = action.getName(i).replace(" ", "-")
|
||||
attribute_dict[
|
||||
"{{{:}}}{:}_desc".format(
|
||||
_accessibility_ns_map_ubuntu["act"], action_name
|
||||
)
|
||||
] = action.getDescription(i)
|
||||
attribute_dict[
|
||||
"{{{:}}}{:}_kb".format(_accessibility_ns_map_ubuntu["act"], action_name)
|
||||
] = action.getKeyBinding(i)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
|
||||
# Add from here if we need more attributes in the future...
|
||||
|
||||
raw_role_name: str = node.getRoleName().strip()
|
||||
node_role_name = (raw_role_name or "unknown").replace(" ", "-")
|
||||
|
||||
if not flag:
|
||||
if raw_role_name == "document spreadsheet":
|
||||
flag = "calc"
|
||||
if raw_role_name == "application" and node.name == "Thunderbird":
|
||||
flag = "thunderbird"
|
||||
|
||||
xml_node = lxml.etree.Element(
|
||||
node_role_name, attrib=attribute_dict, nsmap=_accessibility_ns_map_ubuntu
|
||||
)
|
||||
|
||||
if len(text) > 0:
|
||||
xml_node.text = text
|
||||
|
||||
if depth == MAX_DEPTH:
|
||||
logger.warning("Max depth reached")
|
||||
return xml_node
|
||||
|
||||
if flag == "calc" and node_role_name == "table":
|
||||
# Maximum column: 1024 if ver<=7.3 else 16384
|
||||
# Maximum row: 104 8576
|
||||
# Maximun sheet: 1 0000
|
||||
|
||||
global libreoffice_version_tuple
|
||||
MAXIMUN_COLUMN = 1024 if libreoffice_version_tuple < (7, 4) else 16384
|
||||
MAX_ROW = 104_8576
|
||||
|
||||
index_base = 0
|
||||
first_showing = False
|
||||
column_base = None
|
||||
for r in range(MAX_ROW):
|
||||
for clm in range(column_base or 0, MAXIMUN_COLUMN):
|
||||
child_node: Accessible = node[index_base + clm]
|
||||
showing: bool = child_node.getState().contains(STATE_SHOWING)
|
||||
if showing:
|
||||
child_node: _Element = _create_atspi_node(
|
||||
child_node, depth + 1, flag
|
||||
)
|
||||
if not first_showing:
|
||||
column_base = clm
|
||||
first_showing = True
|
||||
xml_node.append(child_node)
|
||||
elif first_showing and column_base is not None or clm >= 500:
|
||||
break
|
||||
if first_showing and clm == column_base or not first_showing and r >= 500:
|
||||
break
|
||||
index_base += MAXIMUN_COLUMN
|
||||
return xml_node
|
||||
else:
|
||||
try:
|
||||
for i, ch in enumerate(node):
|
||||
if i == MAX_WIDTH:
|
||||
logger.warning("Max width reached")
|
||||
break
|
||||
xml_node.append(_create_atspi_node(ch, depth + 1, flag))
|
||||
except:
|
||||
logger.warning(
|
||||
"Error occurred during children traversing. Has Ignored. Node: %s",
|
||||
lxml.etree.tostring(xml_node, encoding="unicode"),
|
||||
)
|
||||
return xml_node
|
||||
|
||||
|
||||
class UIElement(object):
|
||||
def __init__(self, node):
|
||||
self.node = node
|
||||
|
||||
def getAttributeNames(self):
|
||||
attributes = self.node.getAttributes()
|
||||
|
||||
@staticmethod
|
||||
def systemWideElement():
|
||||
# desktop = pyatspi.Registry.getDesktop(0)
|
||||
# for app in desktop:
|
||||
# for window in app:
|
||||
# if window.getState().contains(pyatspi.STATE_ACTIVE):
|
||||
# active_node = app
|
||||
# return UIElement(active_node)
|
||||
desktop: Accessible = pyatspi.Registry.getDesktop(0)
|
||||
xml_node = lxml.etree.Element(
|
||||
"desktop-frame", nsmap=_accessibility_ns_map_ubuntu
|
||||
)
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
futures = [
|
||||
executor.submit(_create_atspi_node, app_node, 1) for app_node in desktop
|
||||
]
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
xml_tree = future.result()
|
||||
xml_node.append(xml_tree)
|
||||
return lxml.etree.tostring(xml_node, encoding="unicode")
|
||||
|
||||
@property
|
||||
def states(self):
|
||||
state_names = []
|
||||
states: List[StateType] = self.node.getState().get_states()
|
||||
for st in states:
|
||||
state_name: str = StateType._enum_lookup[st]
|
||||
state_names.append(state_name)
|
||||
return state_names
|
||||
|
||||
@property
|
||||
def attributes(self):
|
||||
try:
|
||||
attributes: List[str] = self.node.getAttributes()
|
||||
attribute_dict = {}
|
||||
for attrbt in attributes:
|
||||
attribute_name: str
|
||||
attribute_value: str
|
||||
attribute_name, attribute_value = attrbt.split(":", maxsplit=1)
|
||||
attribute_dict[attribute_name] = attribute_value
|
||||
return attribute_dict
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
@property
|
||||
def component(self):
|
||||
try:
|
||||
component: Component = self.node.queryComponent()
|
||||
return component
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
@property
|
||||
def value(self):
|
||||
try:
|
||||
value: ATValue = self.node.queryValue()
|
||||
return value
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
try:
|
||||
text_obj: ATText = self.node.queryText()
|
||||
except NotImplementedError:
|
||||
return ""
|
||||
else:
|
||||
text: str = text_obj.getText(0, text_obj.characterCount)
|
||||
text = text.replace("\ufffc", "").replace("\ufffd", "")
|
||||
return text
|
||||
|
||||
@property
|
||||
def role(self):
|
||||
return self.node.getRoleName()
|
||||
|
||||
def children(self):
|
||||
"""Return list of children of the current node"""
|
||||
return list(self.node)
|
||||
|
||||
def __repr__(self):
|
||||
return "UIElement%s" % (self.node)
|
||||
@@ -0,0 +1,572 @@
|
||||
import base64
|
||||
import os
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import platform
|
||||
from gui_agents.s1.utils.common_utils import box_iou
|
||||
|
||||
if platform.system() == "Darwin":
|
||||
from AppKit import *
|
||||
from ApplicationServices import (
|
||||
AXUIElementCopyAttributeNames,
|
||||
AXUIElementCopyAttributeValue,
|
||||
AXUIElementCreateSystemWide,
|
||||
)
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI, agent_action
|
||||
|
||||
|
||||
def _normalize_key(key: str) -> str:
|
||||
"""Convert 'cmd' to 'command' for pyautogui compatibility"""
|
||||
return "command" if key == "cmd" else key
|
||||
|
||||
|
||||
def list_apps_in_directories(directories):
|
||||
apps = []
|
||||
for directory in directories:
|
||||
if os.path.exists(directory):
|
||||
directory_apps = [
|
||||
app for app in os.listdir(directory) if app.endswith(".app")
|
||||
]
|
||||
apps.extend(directory_apps)
|
||||
return apps
|
||||
|
||||
|
||||
class MacOSACI(ACI):
|
||||
def __init__(self, top_app_only: bool = True, ocr: bool = False):
|
||||
super().__init__(top_app_only=top_app_only, ocr=ocr)
|
||||
# Directories to search for applications in MacOS
|
||||
directories_to_search = ["/System/Applications", "/Applications"]
|
||||
self.all_apps = list_apps_in_directories(directories_to_search)
|
||||
|
||||
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 = []
|
||||
|
||||
# Inner function to recursively traverse the accessibility tree
|
||||
def traverse_and_preserve(element):
|
||||
role = element.attribute("AXRole")
|
||||
|
||||
if role not in exclude_roles:
|
||||
# TODO: get coordinate values directly from interface
|
||||
position = element.attribute("AXPosition")
|
||||
size = element.attribute("AXSize")
|
||||
if position and size:
|
||||
pos_parts = position.__repr__().split().copy()
|
||||
# Find the parts containing 'x:' and 'y:'
|
||||
x_part = next(part for part in pos_parts if part.startswith("x:"))
|
||||
y_part = next(part for part in pos_parts if part.startswith("y:"))
|
||||
|
||||
# Extract the numerical values after 'x:' and 'y:'
|
||||
x = float(x_part.split(":")[1])
|
||||
y = float(y_part.split(":")[1])
|
||||
|
||||
size_parts = size.__repr__().split().copy()
|
||||
# Find the parts containing 'Width:' and 'Height:'
|
||||
width_part = next(
|
||||
part for part in size_parts if part.startswith("w:")
|
||||
)
|
||||
height_part = next(
|
||||
part for part in size_parts if part.startswith("h:")
|
||||
)
|
||||
|
||||
# Extract the numerical values after 'Width:' and 'Height:'
|
||||
w = float(width_part.split(":")[1])
|
||||
h = float(height_part.split(":")[1])
|
||||
|
||||
if x >= 0 and y >= 0 and w > 0 and h > 0:
|
||||
preserved_nodes.append(
|
||||
{
|
||||
"position": (x, y),
|
||||
"size": (w, h),
|
||||
"title": str(element.attribute("AXTitle")),
|
||||
"text": str(element.attribute("AXDescription"))
|
||||
or str(element.attribute("AXValue")),
|
||||
"role": str(element.attribute("AXRole")),
|
||||
}
|
||||
)
|
||||
|
||||
children = element.children()
|
||||
if children:
|
||||
for child_ref in children:
|
||||
child_element = UIElement(child_ref)
|
||||
traverse_and_preserve(child_element)
|
||||
|
||||
# Start traversing from the given 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
|
||||
],
|
||||
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, 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}\tAXButton\t\t{content}\t\t"
|
||||
)
|
||||
|
||||
node = {
|
||||
"position": (x1, y1),
|
||||
"size": (x2 - x1, y2 - y1),
|
||||
"title": "",
|
||||
"text": content,
|
||||
"role": "AXButton",
|
||||
}
|
||||
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:
|
||||
accessibility_tree = obs["accessibility_tree"]
|
||||
screenshot = obs["screenshot"]
|
||||
self.top_app = (
|
||||
NSWorkspace.sharedWorkspace().frontmostApplication().localizedName()
|
||||
)
|
||||
tree = UIElement(accessibility_tree.attribute("AXFocusedApplication"))
|
||||
exclude_roles = ["AXGroup", "AXLayoutArea", "AXLayoutItem", "AXUnknown"]
|
||||
preserved_nodes = self.preserve_nodes(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:
|
||||
tree_elements, preserved_nodes = self.add_ocr_elements(
|
||||
screenshot, tree_elements, preserved_nodes, "AXButton"
|
||||
)
|
||||
|
||||
self.nodes = preserved_nodes
|
||||
return "\n".join(tree_elements)
|
||||
|
||||
def find_element(self, element_id: int) -> Dict:
|
||||
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
|
||||
"""
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_or_file_name)}); pyautogui.press('enter'); time.sleep(1.0)"
|
||||
|
||||
@agent_action
|
||||
def switch_applications(self, app_or_file_name):
|
||||
"""Switch to a different an application. Utility function to use instead of command+tab
|
||||
Args:
|
||||
app_or_file_name:str, the name of the application or file to switch to
|
||||
"""
|
||||
return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_or_file_name)}); pyautogui.press('enter'); time.sleep(1.0)"
|
||||
|
||||
@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 = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
command = "import pyautogui; "
|
||||
|
||||
# Normalize any 'cmd' to 'command'
|
||||
hold_keys = [_normalize_key(k) for k in hold_keys]
|
||||
|
||||
# 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 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 (return) key should be pressed after typing the text, otherwise assign it to False.
|
||||
"""
|
||||
try:
|
||||
# Use the provided element_id or default to None
|
||||
node = self.find_element(element_id) if element_id is not None else None
|
||||
except:
|
||||
node = None
|
||||
|
||||
if node is not None:
|
||||
# If a node is found, retrieve its coordinates and size
|
||||
coordinates = node["position"]
|
||||
sizes = node["size"]
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
# Start typing at the center of the element
|
||||
command = "import pyautogui; "
|
||||
command += f"pyautogui.click({x}, {y}); "
|
||||
|
||||
if overwrite:
|
||||
# Use 'command' instead of 'cmd'
|
||||
command += f"pyautogui.hotkey('command', 'a', interval=1); 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:
|
||||
# Use 'command' instead of 'cmd'
|
||||
command += f"pyautogui.hotkey('command', 'a', interval=1); 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"]
|
||||
|
||||
# Calculate the center of the element
|
||||
x1 = coordinates1[0] + sizes1[0] // 2
|
||||
y1 = coordinates1[1] + sizes1[1] // 2
|
||||
|
||||
x2 = coordinates2[0] + sizes2[0] // 2
|
||||
y2 = coordinates2[1] + sizes2[1] // 2
|
||||
|
||||
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 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)
|
||||
# print(node.attrib)
|
||||
coordinates = node["position"]
|
||||
sizes = node["size"]
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
return (
|
||||
f"import pyautogui; pyautogui.moveTo({x}, {y}); pyautogui.scroll({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. ['shift', 'c'])
|
||||
"""
|
||||
# Normalize any 'cmd' to 'command'
|
||||
keys = [_normalize_key(k) for k in keys]
|
||||
# add quotes around the keys
|
||||
keys = [f"'{key}'" for key in keys]
|
||||
return f"import pyautogui; pyautogui.hotkey({', '.join(keys)}, interval=1)"
|
||||
|
||||
@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
|
||||
"""
|
||||
# Normalize any 'cmd' to 'command' in both lists
|
||||
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
|
||||
"""
|
||||
return f"""import time; time.sleep({time})"""
|
||||
|
||||
@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"""
|
||||
|
||||
|
||||
class UIElement(object):
|
||||
|
||||
def __init__(self, ref=None):
|
||||
self.ref = ref
|
||||
|
||||
def getAttributeNames(self):
|
||||
error_code, attributeNames = AXUIElementCopyAttributeNames(self.ref, None)
|
||||
return list(attributeNames)
|
||||
|
||||
def attribute(self, key: str):
|
||||
error, value = AXUIElementCopyAttributeValue(self.ref, key, None)
|
||||
return value
|
||||
|
||||
def children(self):
|
||||
return self.attribute("AXChildren")
|
||||
|
||||
def systemWideElement():
|
||||
ref = AXUIElementCreateSystemWide()
|
||||
return UIElement(ref)
|
||||
|
||||
def role(self):
|
||||
return self.attribute("AXRole")
|
||||
|
||||
def position(self):
|
||||
pos = self.attribute("AXPosition")
|
||||
if pos is None:
|
||||
return None
|
||||
pos_parts = pos.__repr__().split().copy()
|
||||
# Find the parts containing 'x:' and 'y:'
|
||||
x_part = next(part for part in pos_parts if part.startswith("x:"))
|
||||
y_part = next(part for part in pos_parts if part.startswith("y:"))
|
||||
|
||||
# Extract the numerical values after 'x:' and 'y:'
|
||||
x = float(x_part.split(":")[1])
|
||||
y = float(y_part.split(":")[1])
|
||||
|
||||
return (x, y)
|
||||
|
||||
def size(self):
|
||||
size = self.attribute("AXSize")
|
||||
if size is None:
|
||||
return None
|
||||
size_parts = size.__repr__().split().copy()
|
||||
# Find the parts containing 'Width:' and 'Height:'
|
||||
width_part = next(part for part in size_parts if part.startswith("w:"))
|
||||
height_part = next(part for part in size_parts if part.startswith("h:"))
|
||||
|
||||
# Extract the numerical values after 'Width:' and 'Height:'
|
||||
w = float(width_part.split(":")[1])
|
||||
h = float(height_part.split(":")[1])
|
||||
return (w, h)
|
||||
|
||||
def isValid(self):
|
||||
if self.position() is not None and self.size() is not None:
|
||||
return True
|
||||
|
||||
def parse(self, element):
|
||||
position = element.position(element)
|
||||
size = element.size(element)
|
||||
return {
|
||||
"position": position,
|
||||
"size": size,
|
||||
"title": str(element.attribute("AXTitle")),
|
||||
"text": str(element.attribute("AXDescription"))
|
||||
or str(element.attribute("AXValue")),
|
||||
"role": str(element.attribute("AXRole")),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_current_applications(obs: Dict):
|
||||
# Get the shared workspace instance
|
||||
workspace = NSWorkspace.sharedWorkspace()
|
||||
|
||||
# Get a list of running applications
|
||||
running_apps = workspace.runningApplications()
|
||||
|
||||
# Iterate through the list and print each application's name
|
||||
current_apps = []
|
||||
for app in running_apps:
|
||||
if app.activationPolicy() == 0:
|
||||
app_name = app.localizedName()
|
||||
current_apps.append(app_name)
|
||||
|
||||
return current_apps
|
||||
|
||||
@staticmethod
|
||||
def list_apps_in_directories():
|
||||
directories_to_search = ["/System/Applications", "/Applications"]
|
||||
apps = []
|
||||
for directory in directories_to_search:
|
||||
if os.path.exists(directory):
|
||||
directory_apps = [
|
||||
app for app in os.listdir(directory) if app.endswith(".app")
|
||||
]
|
||||
apps.extend(directory_apps)
|
||||
return apps
|
||||
|
||||
@staticmethod
|
||||
def get_top_app(obs: Dict):
|
||||
return NSWorkspace.sharedWorkspace().frontmostApplication().localizedName()
|
||||
|
||||
def __repr__(self):
|
||||
return "UIElement%s" % (self.ref)
|
||||
@@ -0,0 +1,532 @@
|
||||
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})"
|
||||
@@ -0,0 +1,609 @@
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from typing import Dict, List, Tuple
|
||||
import numpy as np
|
||||
import requests
|
||||
from gui_agents.s1.utils.common_utils import box_iou
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
state_ns = "uri:deskat:state.at-spi.gnome.org"
|
||||
component_ns = "uri:deskat:component.at-spi.gnome.org"
|
||||
|
||||
|
||||
# Agent action decorator
|
||||
def agent_action(func):
|
||||
func.is_agent_action = True
|
||||
return func
|
||||
|
||||
|
||||
class GroundingAgent:
|
||||
def __init__(self, vm_version: str, top_app=None, top_app_only=True, ocr=True):
|
||||
self.active_apps = set()
|
||||
self.top_app = top_app
|
||||
self.top_app_only = (
|
||||
top_app_only # Only include top app in the accessibility tree
|
||||
)
|
||||
self.ocr = ocr
|
||||
self.index_out_of_range_flag = False
|
||||
self.app_setup_code = 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'])
|
||||
"""
|
||||
|
||||
self.top_active_app = None
|
||||
self.notes = []
|
||||
self.clipboard = ""
|
||||
|
||||
# TODO: this is terrible, fix this
|
||||
# global state_ns, component_ns, attributes_ns, value_ns
|
||||
# if vm_version == "old":
|
||||
# state_ns = "uri:deskat:state.at-spi.gnome.org"
|
||||
# component_ns = "uri:deskat:component.at-spi.gnome.org"
|
||||
# elif vm_version == 'win':
|
||||
# state_ns = "uri:deskat:state.at-spi.gnome.org"
|
||||
# component_ns = "uri:deskat:component.at-spi.gnome.org"
|
||||
# else:
|
||||
# attributes_ns = "https://accessibility.windows.example.org/ns/attributes"
|
||||
# state_ns = "https://accessibility.ubuntu.example.org/ns/state"
|
||||
# component_ns = "https://accessibility.ubuntu.example.org/ns/component"
|
||||
# value_ns = "https://accessibility.ubuntu.example.org/ns/value"
|
||||
|
||||
def get_current_applications(self, obs):
|
||||
tree = ET.ElementTree(ET.fromstring(obs["accessibility_tree"]))
|
||||
apps = []
|
||||
root = tree.getroot()
|
||||
for item in root:
|
||||
apps.append(item.get("name", "").replace("\\", ""))
|
||||
return apps
|
||||
|
||||
def check_new_apps(self, old_apps, new_apps):
|
||||
return new_apps - old_apps
|
||||
|
||||
def find_active_applications(self, tree):
|
||||
# names of applications to keep TODO: soffice is a single application with all the isntances like impress, calc etc. being frames this will need to be dealt with separately
|
||||
to_keep = ["Program Manager"]
|
||||
apps_with_active_tag = []
|
||||
for application in list(tree.getroot()):
|
||||
app_name = application.get("name")
|
||||
for frame in application:
|
||||
is_active = frame.attrib.get("{{{:}}}active".format(state_ns), "false")
|
||||
if is_active == "true":
|
||||
apps_with_active_tag.append(app_name)
|
||||
print(apps_with_active_tag)
|
||||
if apps_with_active_tag:
|
||||
to_keep.append(apps_with_active_tag[-1])
|
||||
return to_keep
|
||||
|
||||
def filter_active_app(self, tree):
|
||||
for application in list(tree.getroot()):
|
||||
app_name = application.attrib.get("name")
|
||||
for frame in application:
|
||||
is_active = frame.attrib.get("{{{:}}}active".format(state_ns), "false")
|
||||
if is_active == "true":
|
||||
return app_name
|
||||
return None
|
||||
|
||||
def filter_nodes(self, tree, show_all=False):
|
||||
# created and populate a preserved nodes list which filters out unnecessary elements and keeps only those elements which are currently showing on the screen
|
||||
# TODO: include offscreen elements and then scroll to them before clicking
|
||||
preserved_nodes = []
|
||||
exclude_tags = ["panel", "window", "filler", "frame", "separator", "scroll-bar"]
|
||||
|
||||
for node in tree.iter():
|
||||
if node.tag not in exclude_tags:
|
||||
if show_all:
|
||||
if node.attrib.get(f"{{{state_ns}}}enabled") == "true":
|
||||
coords: Tuple[int, int] = eval(
|
||||
node.get(
|
||||
"{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
)
|
||||
)
|
||||
if coords[0] >= 0 and coords[1] >= 0:
|
||||
preserved_nodes.append(node)
|
||||
# if show_all is false, only show elements that are currently showing on screen
|
||||
else:
|
||||
if node.attrib.get(f"{{{state_ns}}}visible") == "true":
|
||||
coords: Tuple[int, int] = eval(
|
||||
node.get(
|
||||
"{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
)
|
||||
)
|
||||
|
||||
if coords[0] >= 0 and coords[1] >= 0:
|
||||
preserved_nodes.append(node)
|
||||
return preserved_nodes
|
||||
|
||||
def linearize_tree(self, preserved_nodes):
|
||||
# TODO: Run an ablation to check if class and desc
|
||||
# linearized_accessibility_tree = ["id\ttag\tname\ttext\tclass\tdescription"]
|
||||
linearized_accessibility_tree = ["id\ttag\tname\ttext"]
|
||||
for idx, node in enumerate(preserved_nodes):
|
||||
if node.text:
|
||||
text = (
|
||||
node.text
|
||||
if '"' not in node.text
|
||||
else '"{:}"'.format(node.text.replace('"', '""'))
|
||||
)
|
||||
else:
|
||||
text = '""'
|
||||
|
||||
linearized_accessibility_tree.append(
|
||||
"{:}\t{:}\t{:}\t{:}".format(
|
||||
idx,
|
||||
node.tag,
|
||||
node.get("name", ""),
|
||||
text,
|
||||
# node.get("{{{:}}}class".format(attributes_ns), ""),
|
||||
# node.get("{{{:}}}description".format(attributes_ns), ""),
|
||||
)
|
||||
)
|
||||
|
||||
# returning list of linearized elements
|
||||
return linearized_accessibility_tree
|
||||
|
||||
def extract_elements_from_screenshot(self, screenshot) -> Dict:
|
||||
"""Uses paddle-ocr to extract elements with text from the screenshot. The elements will be added to the linearized accessibility tree downstream"""
|
||||
|
||||
# Convert screenshot to PIL image
|
||||
def send_image_to_ocr(screenshot) -> Dict:
|
||||
|
||||
# url = os.environ.get("OCR_SERVER_ADDRESS", "")
|
||||
url = "http://127.0.0.1:8083/ocr/"
|
||||
if url == "":
|
||||
raise Exception("OCR SERVER ADDRESS NOT SET")
|
||||
encoded_screenshot = base64.b64encode(screenshot).decode("utf-8")
|
||||
data = {"img_bytes": encoded_screenshot}
|
||||
response = requests.post(url, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
return {
|
||||
"error": f"Request failed with status code {response.status_code}",
|
||||
"results": [],
|
||||
}
|
||||
|
||||
return send_image_to_ocr(screenshot)["results"]
|
||||
|
||||
def add_ocr_elements(
|
||||
self, screenshot, linearized_accessibility_tree, preserved_nodes
|
||||
):
|
||||
# Get the bounding boxes of the elements in the linearized accessibility tree
|
||||
tree_bboxes = []
|
||||
for node in preserved_nodes:
|
||||
coordinates: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}size".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
tree_bboxes.append(
|
||||
[
|
||||
coordinates[0],
|
||||
coordinates[1],
|
||||
coordinates[0] + sizes[0],
|
||||
coordinates[1] + sizes[1],
|
||||
]
|
||||
)
|
||||
|
||||
# Use OCR to found boxes that might be missing from the accessibility tree
|
||||
try:
|
||||
ocr_bboxes = self.extract_elements_from_screenshot(screenshot)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
ocr_bboxes = []
|
||||
else:
|
||||
# Check for intersection over union between the existing atree bounding boxes and the ocr bounding boxes, if ocr bounding boxes are new add them to the linearized accesibility tree
|
||||
if (
|
||||
len(ocr_bboxes) > 0
|
||||
): # Only check IOUs and add if there are any bounding boxes returned by the ocr module
|
||||
preserved_nodes_index = len(preserved_nodes)
|
||||
for ind, (i, content, box) in enumerate(ocr_bboxes):
|
||||
# x1, y1, x2, y2 = int(box.get('left', 0)), int(box['top']), int(), int(box['bottom'])
|
||||
(
|
||||
x1,
|
||||
y1,
|
||||
x2,
|
||||
y2,
|
||||
) = (
|
||||
int(box.get("left", 0)),
|
||||
int(box.get("top", 0)),
|
||||
int(box.get("right", 0)),
|
||||
int(box.get("bottom", 0)),
|
||||
)
|
||||
iou = box_iou(
|
||||
np.array(tree_bboxes, dtype=np.float32),
|
||||
np.array([[x1, y1, x2, y2]], dtype=np.float32),
|
||||
).flatten()
|
||||
|
||||
if max(iou) < 0.1:
|
||||
# Add the element to the linearized accessibility tree
|
||||
# TODO: ocr detected elements should be classified for their tag, currently set to push button for the agent to think they are interactable
|
||||
linearized_accessibility_tree.append(
|
||||
f"{preserved_nodes_index}\tpush-button\t\t{content}\t\t"
|
||||
)
|
||||
|
||||
# add to preserved node with the component_ns prefix node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)"
|
||||
node = ET.Element(
|
||||
"ocr_node",
|
||||
attrib={
|
||||
"text": content,
|
||||
"{{{}}}screencoord".format(
|
||||
component_ns
|
||||
): "({},{})".format(x1, y1),
|
||||
"{{{}}}size".format(component_ns): "({},{})".format(
|
||||
x2 - x1, y2 - y1
|
||||
),
|
||||
},
|
||||
)
|
||||
preserved_nodes.append(node)
|
||||
preserved_nodes_index += 1
|
||||
|
||||
return linearized_accessibility_tree, preserved_nodes
|
||||
|
||||
def linearize_and_annotate_tree(self, obs, show_all=False):
|
||||
accessibility_tree = obs["accessibility_tree"]
|
||||
screenshot = obs["screenshot"]
|
||||
|
||||
# convert the accessibility tree from a string representation to an xml tree
|
||||
tree = ET.ElementTree(ET.fromstring(accessibility_tree))
|
||||
|
||||
# Get the applications to keep based on the active applications
|
||||
to_keep = self.find_active_applications(tree)
|
||||
self.top_app = to_keep[-1]
|
||||
|
||||
# Remove applications which are not included in the to_keep list
|
||||
if not show_all:
|
||||
for application in list(tree.getroot()):
|
||||
if application.attrib.get("name", "") not in to_keep:
|
||||
tree.getroot().remove(application)
|
||||
|
||||
# Save tree for debugging
|
||||
# from datetime import datetime
|
||||
# with open(f"tree_raw_{datetime.now()}.xml", "wb") as file:
|
||||
# tree.write(file, encoding="utf-8", xml_declaration=True)
|
||||
|
||||
# Filter out filler elements and overlapping elements
|
||||
preserved_nodes = self.filter_nodes(tree, show_all)
|
||||
|
||||
assert len(preserved_nodes) > 0
|
||||
|
||||
# Linearize the tree as tsv
|
||||
linearized_accessibility_tree = self.linearize_tree(preserved_nodes)
|
||||
|
||||
# Add OCR elements to the linearized accessibility tree to account for elements that are not in the accessibility tree
|
||||
if self.ocr:
|
||||
linearized_accessibility_tree, preserved_nodes = self.add_ocr_elements(
|
||||
screenshot, linearized_accessibility_tree, preserved_nodes
|
||||
)
|
||||
|
||||
# Convert accessibility tree to a string
|
||||
linearized_accessibility_tree = "\n".join(linearized_accessibility_tree)
|
||||
|
||||
# TODO: side-effect, set in separate functions
|
||||
self.nodes = preserved_nodes
|
||||
|
||||
return linearized_accessibility_tree
|
||||
|
||||
def find_element(self, element_id):
|
||||
try:
|
||||
selected_element = self.nodes[int(element_id)]
|
||||
except:
|
||||
print("The index of the selected element was out of range.")
|
||||
selected_element = self.nodes[0]
|
||||
self.index_out_of_range_flag = True
|
||||
return selected_element
|
||||
|
||||
@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] = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}size".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
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_window(self):
|
||||
"""Switch to a different application that is already open"""
|
||||
# return self.app_setup_code.replace("APP_NAME", app_code)
|
||||
return f"import pyautogui; pyautogui.hotkey('alt', 'tab');"
|
||||
|
||||
@agent_action
|
||||
def type(
|
||||
self,
|
||||
text: str,
|
||||
element_id: int = None,
|
||||
overwrite: bool = False,
|
||||
enter: bool = False,
|
||||
):
|
||||
"""Type text into the element
|
||||
Args:
|
||||
text:str the text to type
|
||||
element_id:int ID of the element to type into. If not provided, typing will start at the current cursor location.
|
||||
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:
|
||||
# Use the provided element_id or default to None
|
||||
node = self.find_element(element_id) if element_id is not None else None
|
||||
except:
|
||||
node = None
|
||||
|
||||
if node is not None:
|
||||
# If a node is found, retrieve its coordinates and size
|
||||
coordinates = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
# Start typing at the center of the element
|
||||
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
|
||||
|
||||
# if overwrite:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.hotkey("ctrl", "a"); pyautogui.press("backspace"); pyautogui.typewrite({repr(text)})"""
|
||||
# else:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.hotkey("ctrl", "a"); pyautogui.press("backspace"); pyautogui.typewrite("{text}")"""
|
||||
|
||||
# @agent_action
|
||||
# def type_and_enter(self, element_id:int, text:str, overwrite: bool = True):
|
||||
# '''Type text into the element and press enter
|
||||
# Args:
|
||||
# element_id:int ID of the element to type into
|
||||
# text:str the text to type into the element
|
||||
# '''
|
||||
# try:
|
||||
# node = self.find_element(element_id)
|
||||
# except:
|
||||
# node = self.find_element(0)
|
||||
# # print(node.attrib)
|
||||
# coordinates = eval(
|
||||
# node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)"))
|
||||
# sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# # Calculate the center of the element
|
||||
# x = coordinates[0] + sizes[0] // 2
|
||||
# y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
# # Return pyautoguicode to type into the element
|
||||
# if overwrite:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.hotkey("ctrl", "a"); pyautogui.press("backspace"); pyautogui.typewrite({repr(text)}); pyautogui.press("enter")"""
|
||||
# else:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.typewrite({repr(text)}); pyautogui.press("enter")"""
|
||||
|
||||
# @agent_action
|
||||
# def copy_text(self, element_id:int):
|
||||
# '''Copy the selected text, use instead of ctrl+c
|
||||
# Args:
|
||||
# element_id:int ID of the element to copy text from
|
||||
# '''
|
||||
# try:
|
||||
# node = self.find_element(element_id)
|
||||
# except:
|
||||
# node = self.find_element(0)
|
||||
|
||||
# self.clipboard = node.text
|
||||
|
||||
# @agent_action
|
||||
# def paste_text(self, element_id:int, overwrite: bool = True):
|
||||
# '''Paste text from the clipboard into the element, use instead of ctrl+v
|
||||
# Args:
|
||||
# element_id:int ID of the element to copy text from
|
||||
# overwrite:bool a boolean value to determine if the text should be pasted over the existing text or appended to it
|
||||
# '''
|
||||
# try:
|
||||
# node = self.find_element(element_id)
|
||||
# except:
|
||||
# node = self.find_element(0)
|
||||
|
||||
# coordinates = eval(
|
||||
# node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)"))
|
||||
# sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# # Calculate the center of the element
|
||||
# x = coordinates[0] + sizes[0] // 2
|
||||
# y = coordinates[1] + sizes[1] // 2
|
||||
|
||||
# # Return pyautoguicode to paste into the element
|
||||
# if overwrite:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.typewrite("{self.clipboard}");"""
|
||||
# else:
|
||||
# return f"""import pyautogui; pyautogui.click({x}, {y}); pyautogui.hotkey("ctrl", "a"); pyautogui.press("backspace"); pyautogui.typewrite("{self.clipboard}");"""
|
||||
|
||||
@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, 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 = eval(
|
||||
node1.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes1 = eval(node1.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
coordinates2 = eval(
|
||||
node2.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes2 = eval(node2.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x1 = coordinates1[0] + sizes1[0] // 2
|
||||
y1 = coordinates1[1] + sizes1[1] // 2
|
||||
|
||||
x2 = coordinates2[0] + sizes2[0] // 2
|
||||
y2 = coordinates2[1] + sizes2[1] // 2
|
||||
|
||||
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 scroll(self, element_id: int, clicks: int):
|
||||
"""Scroll the element in the specified direction
|
||||
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)
|
||||
# print(node.attrib)
|
||||
coordinates = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes = eval(node.get("{{{:}}}size".format(component_ns), "(-1, -1)"))
|
||||
|
||||
# Calculate the center of the element
|
||||
x = coordinates[0] + sizes[0] // 2
|
||||
y = coordinates[1] + sizes[1] // 2
|
||||
return (
|
||||
f"import pyautogui; pyautogui.moveTo({x}, {y}); pyautogui.scroll({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"""
|
||||
return """DONE"""
|
||||
|
||||
@agent_action
|
||||
def fail(self):
|
||||
"""End the current task with a failure"""
|
||||
return """FAIL"""
|
||||
@@ -0,0 +1,281 @@
|
||||
import argparse
|
||||
import datetime
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pyautogui
|
||||
|
||||
from gui_agents.s1.core.AgentS import GraphSearchAgent, UIAgent
|
||||
|
||||
current_platform = platform.system().lower()
|
||||
|
||||
# Global flag to track pause state for debugging
|
||||
paused = False
|
||||
|
||||
|
||||
def get_char():
|
||||
"""Get a single character from stdin without pressing Enter"""
|
||||
try:
|
||||
# Import termios and tty on Unix-like systems
|
||||
if platform.system() in ["Darwin", "Linux"]:
|
||||
import termios
|
||||
import tty
|
||||
|
||||
fd = sys.stdin.fileno()
|
||||
old_settings = termios.tcgetattr(fd)
|
||||
try:
|
||||
tty.setraw(sys.stdin.fileno())
|
||||
ch = sys.stdin.read(1)
|
||||
finally:
|
||||
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
|
||||
return ch
|
||||
else:
|
||||
# Windows fallback
|
||||
import msvcrt
|
||||
|
||||
return msvcrt.getch().decode("utf-8", errors="ignore")
|
||||
except:
|
||||
return input() # Fallback for non-terminal environments
|
||||
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
"""Handle Ctrl+C signal for debugging during agent execution"""
|
||||
global paused
|
||||
|
||||
if not paused:
|
||||
print("\n\n🔸 Agent-S Workflow Paused 🔸")
|
||||
print("=" * 50)
|
||||
print("Options:")
|
||||
print(" • Press Ctrl+C again to quit")
|
||||
print(" • Press Esc to resume workflow")
|
||||
print("=" * 50)
|
||||
|
||||
paused = True
|
||||
|
||||
while paused:
|
||||
try:
|
||||
print("\n[PAUSED] Waiting for input... ", end="", flush=True)
|
||||
char = get_char()
|
||||
|
||||
if ord(char) == 3: # Ctrl+C
|
||||
print("\n\n🛑 Exiting Agent-S...")
|
||||
sys.exit(0)
|
||||
elif ord(char) == 27: # Esc
|
||||
print("\n\n▶️ Resuming Agent-S workflow...")
|
||||
paused = False
|
||||
break
|
||||
else:
|
||||
print(f"\n Unknown command: '{char}' (ord: {ord(char)})")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n🛑 Exiting Agent-S...")
|
||||
sys.exit(0)
|
||||
else:
|
||||
# Already paused, second Ctrl+C means quit
|
||||
print("\n\n🛑 Exiting Agent-S...")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
# Set up signal handler for Ctrl+C
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
if current_platform == "darwin":
|
||||
from gui_agents.s1.aci.MacOSACI import MacOSACI, UIElement
|
||||
elif current_platform == "linux":
|
||||
from gui_agents.s1.aci.LinuxOSACI import LinuxACI, UIElement
|
||||
elif current_platform == "windows":
|
||||
from gui_agents.s1.aci.WindowsOSACI import WindowsACI, UIElement
|
||||
else:
|
||||
raise ValueError(f"Unsupported platform: {current_platform}")
|
||||
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
|
||||
|
||||
log_dir = "logs"
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
file_handler = logging.FileHandler(
|
||||
os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8"
|
||||
)
|
||||
debug_handler = logging.FileHandler(
|
||||
os.path.join("logs", "debug-{:}.log".format(datetime_str)), encoding="utf-8"
|
||||
)
|
||||
stdout_handler = logging.StreamHandler(sys.stdout)
|
||||
sdebug_handler = logging.FileHandler(
|
||||
os.path.join("logs", "sdebug-{:}.log".format(datetime_str)), encoding="utf-8"
|
||||
)
|
||||
|
||||
file_handler.setLevel(logging.INFO)
|
||||
debug_handler.setLevel(logging.DEBUG)
|
||||
stdout_handler.setLevel(logging.INFO)
|
||||
sdebug_handler.setLevel(logging.DEBUG)
|
||||
|
||||
formatter = logging.Formatter(
|
||||
fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
|
||||
)
|
||||
file_handler.setFormatter(formatter)
|
||||
debug_handler.setFormatter(formatter)
|
||||
stdout_handler.setFormatter(formatter)
|
||||
sdebug_handler.setFormatter(formatter)
|
||||
|
||||
stdout_handler.addFilter(logging.Filter("desktopenv"))
|
||||
sdebug_handler.addFilter(logging.Filter("desktopenv"))
|
||||
|
||||
logger.addHandler(file_handler)
|
||||
logger.addHandler(debug_handler)
|
||||
logger.addHandler(stdout_handler)
|
||||
logger.addHandler(sdebug_handler)
|
||||
|
||||
platform_os = platform.system()
|
||||
|
||||
|
||||
def show_permission_dialog(code: str, action_description: str):
|
||||
"""Show a platform-specific permission dialog and return True if approved."""
|
||||
if platform.system() == "Darwin":
|
||||
result = os.system(
|
||||
f'osascript -e \'display dialog "Do you want to execute this action?\n\n{code} which will try to {action_description}" with title "Action Permission" buttons {{"Cancel", "OK"}} default button "OK" cancel button "Cancel"\''
|
||||
)
|
||||
return result == 0
|
||||
elif platform.system() == "Linux":
|
||||
result = os.system(
|
||||
f'zenity --question --title="Action Permission" --text="Do you want to execute this action?\n\n{code}" --width=400 --height=200'
|
||||
)
|
||||
return result == 0
|
||||
return False
|
||||
|
||||
|
||||
def run_agent(agent: UIAgent, instruction: str):
|
||||
global paused
|
||||
obs = {}
|
||||
traj = "Task:\n" + instruction
|
||||
subtask_traj = ""
|
||||
for step in range(15):
|
||||
# Check if we're in paused state and wait
|
||||
while paused:
|
||||
time.sleep(0.1)
|
||||
obs["accessibility_tree"] = UIElement.systemWideElement()
|
||||
|
||||
# Get screen shot using pyautogui.
|
||||
# Take a screenshot
|
||||
screenshot = pyautogui.screenshot()
|
||||
|
||||
# Save the screenshot to a BytesIO object
|
||||
buffered = io.BytesIO()
|
||||
screenshot.save(buffered, format="PNG")
|
||||
|
||||
# Get the byte value of the screenshot
|
||||
screenshot_bytes = buffered.getvalue()
|
||||
# Convert to base64 string.
|
||||
obs["screenshot"] = screenshot_bytes
|
||||
|
||||
# Check again for pause state before prediction
|
||||
while paused:
|
||||
time.sleep(0.1)
|
||||
|
||||
print(f"\n🔄 Step {step + 1}/15: Getting next action from agent...")
|
||||
|
||||
# Get next action code from the agent
|
||||
info, code = agent.predict(instruction=instruction, observation=obs)
|
||||
|
||||
if "done" in code[0].lower() or "fail" in code[0].lower():
|
||||
if platform.system() == "Darwin":
|
||||
os.system(
|
||||
f'osascript -e \'display dialog "Task Completed" with title "OpenACI Agent" buttons "OK" default button "OK"\''
|
||||
)
|
||||
elif platform.system() == "Linux":
|
||||
os.system(
|
||||
f'zenity --info --title="OpenACI Agent" --text="Task Completed" --width=200 --height=100'
|
||||
)
|
||||
|
||||
agent.update_narrative_memory(traj)
|
||||
break
|
||||
|
||||
if "next" in code[0].lower():
|
||||
continue
|
||||
|
||||
if "wait" in code[0].lower():
|
||||
print("⏳ Agent requested wait...")
|
||||
time.sleep(5)
|
||||
continue
|
||||
|
||||
else:
|
||||
time.sleep(1.0)
|
||||
print("EXECUTING CODE:", code[0])
|
||||
|
||||
# Check for pause state before execution
|
||||
while paused:
|
||||
time.sleep(0.1)
|
||||
|
||||
# Ask for permission before executing
|
||||
exec(code[0])
|
||||
time.sleep(1.0)
|
||||
|
||||
# Update task and subtask trajectories and optionally the episodic memory
|
||||
traj += (
|
||||
"\n\nReflection:\n"
|
||||
+ str(info["reflection"])
|
||||
+ "\n\n----------------------\n\nPlan:\n"
|
||||
+ info["executor_plan"]
|
||||
)
|
||||
subtask_traj = agent.update_episodic_memory(info, subtask_traj)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run GraphSearchAgent with specified model."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="gpt-4o-mini",
|
||||
help="Specify the model to use (e.g., gpt-4o)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if current_platform == "Darwin":
|
||||
grounding_agent = MacOSACI()
|
||||
elif current_platform == "Windows":
|
||||
grounding_agent = WindowsACI()
|
||||
elif current_platform == "Linux":
|
||||
grounding_agent = LinuxACI()
|
||||
else:
|
||||
raise ValueError("Unsupported platform")
|
||||
|
||||
while True:
|
||||
query = input("Query: ")
|
||||
if "gpt" in args.model:
|
||||
engine_type = "openai"
|
||||
elif "claude" in args.model:
|
||||
engine_type = "anthropic"
|
||||
engine_params = {
|
||||
"engine_type": engine_type,
|
||||
"model": args.model,
|
||||
}
|
||||
|
||||
agent = GraphSearchAgent(
|
||||
engine_params,
|
||||
grounding_agent,
|
||||
platform=current_platform,
|
||||
action_space="pyautogui",
|
||||
observation_type="mixed",
|
||||
)
|
||||
|
||||
agent.reset()
|
||||
|
||||
# Run the agent on your own device
|
||||
run_agent(agent, query)
|
||||
|
||||
response = input("Would you like to provide another query? (y/n): ")
|
||||
if response.lower() != "y":
|
||||
break
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,382 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI
|
||||
from gui_agents.s1.core.Manager import Manager
|
||||
from gui_agents.s1.core.Worker import Worker
|
||||
from gui_agents.s1.utils.common_utils import Node
|
||||
from gui_agents.utils import download_kb_data
|
||||
|
||||
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 GraphSearchAgent(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 = "pyatuogui",
|
||||
observation_type: str = "mixed",
|
||||
search_engine: Optional[str] = None,
|
||||
memory_root_path: str = os.getcwd(),
|
||||
memory_folder_name: str = "kb_s1",
|
||||
kb_release_tag: str = "v0.2.2",
|
||||
):
|
||||
"""Initialize GraphSearchAgent
|
||||
|
||||
Args:
|
||||
engine_params: Configuration parameters for the LLM engine
|
||||
grounding_agent: Instance of ACI class for UI interaction
|
||||
platform: Operating system platform (macos, ubuntu)
|
||||
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)
|
||||
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".
|
||||
"""
|
||||
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.
|
||||
print("Downloading knowledge base initial Agent-S knowledge...")
|
||||
self.local_kb_path = os.path.join(
|
||||
self.memory_root_path, self.memory_folder_name
|
||||
)
|
||||
|
||||
if not os.path.exists(self.local_kb_path):
|
||||
download_kb_data(
|
||||
version="s1",
|
||||
release_tag=kb_release_tag,
|
||||
download_dir=self.local_kb_path,
|
||||
platform=self.platform,
|
||||
)
|
||||
print(
|
||||
f"Successfully completed download of knowledge base for version s1, 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."
|
||||
)
|
||||
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset agent state and initialize components"""
|
||||
# Initialize core components
|
||||
self.planner = Manager(
|
||||
self.engine_params,
|
||||
self.grounding_agent,
|
||||
platform=self.platform,
|
||||
search_engine=self.engine,
|
||||
local_kb_path=self.local_kb_path,
|
||||
)
|
||||
self.executor = Worker(
|
||||
self.engine_params,
|
||||
self.grounding_agent,
|
||||
platform=self.platform,
|
||||
local_kb_path=self.local_kb_path,
|
||||
)
|
||||
|
||||
# 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_feedback: str = ""
|
||||
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]]:
|
||||
"""Predict next UI action sequence
|
||||
|
||||
Args:
|
||||
instruction: Natural language instruction
|
||||
observation: Current UI state observation Dictionary {"accessibility_tree": str, "screenshot": bytes}
|
||||
info: Dictionary containing additional information.
|
||||
|
||||
Returns:
|
||||
Tuple of (agent info dict, list of actions)
|
||||
"""
|
||||
# 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, then true again after a failed plan
|
||||
if self.requires_replan:
|
||||
logger.info("(RE)PLANNING...")
|
||||
# failure feedback is the reason for the failure of the previous plan
|
||||
planner_info, self.subtasks = self.planner.get_action_queue(
|
||||
instruction=instruction,
|
||||
observation=observation,
|
||||
failure_feedback=self.failure_feedback,
|
||||
)
|
||||
|
||||
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...")
|
||||
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
|
||||
if "FAIL" in actions:
|
||||
self.requires_replan = True
|
||||
# set the failure feedback to the evaluator feedback
|
||||
self.failure_feedback = f"Completed subtasks: {self.completed_tasks}. The subtask {self.current_subtask} cannot be completed. Please try another approach. {executor_info['plan_code']}. Please replan."
|
||||
self.needs_next_subtask = True
|
||||
|
||||
# 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
|
||||
|
||||
elif "DONE" in actions:
|
||||
self.requires_replan = False
|
||||
self.completed_tasks.append(self.current_subtask)
|
||||
self.needs_next_subtask = True
|
||||
if self.subtasks:
|
||||
self.should_send_action = False
|
||||
self.subtask_status = "Done"
|
||||
|
||||
self.reset_executor_state()
|
||||
|
||||
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,18 @@
|
||||
from typing import Dict, Optional
|
||||
|
||||
from gui_agents.s1.mllm.MultimodalAgent import LMMAgent
|
||||
|
||||
|
||||
class BaseModule:
|
||||
def __init__(self, engine_params: Dict, platform: str):
|
||||
self.engine_params = engine_params
|
||||
self.platform = platform
|
||||
|
||||
def _create_agent(
|
||||
self, system_prompt: str = None, engine_params: Optional[Dict] = None
|
||||
) -> LMMAgent:
|
||||
"""Create a new LMMAgent instance"""
|
||||
agent = LMMAgent(engine_params or self.engine_params)
|
||||
if system_prompt:
|
||||
agent.add_system_prompt(system_prompt)
|
||||
return agent
|
||||
@@ -0,0 +1,250 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Dict, Tuple
|
||||
|
||||
import numpy as np
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
|
||||
from gui_agents.s1.core.BaseModule import BaseModule
|
||||
from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s1.mllm.MultimodalEngine import OpenAIEmbeddingEngine
|
||||
from gui_agents.s1.utils.common_utils import (
|
||||
load_embeddings,
|
||||
load_knowledge_base,
|
||||
save_embeddings,
|
||||
)
|
||||
from gui_agents.s1.utils.query_perplexica import query_to_perplexica
|
||||
|
||||
|
||||
class KnowledgeBase(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
local_kb_path: str,
|
||||
platform: str,
|
||||
engine_params: Dict,
|
||||
use_image_for_search: bool = False,
|
||||
):
|
||||
super().__init__(engine_params, platform)
|
||||
|
||||
self.local_kb_path = local_kb_path
|
||||
|
||||
# initialize embedding engine
|
||||
# TODO: Support other embedding engines
|
||||
self.embedding_engine = OpenAIEmbeddingEngine(
|
||||
api_key=(
|
||||
engine_params["api_key"]
|
||||
if "api_key" in engine_params
|
||||
else os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
)
|
||||
|
||||
# Initialize paths for different memory types
|
||||
self.episodic_memory_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "episodic_memory.json"
|
||||
)
|
||||
self.narrative_memory_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "narrative_memory.json"
|
||||
)
|
||||
self.embeddings_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "embeddings.pkl"
|
||||
)
|
||||
|
||||
self.rag_module_system_prompt = PROCEDURAL_MEMORY.RAG_AGENT.replace(
|
||||
"CURRENT_OS", self.platform
|
||||
)
|
||||
|
||||
# All three agent share a generic RAG prompt that ask agent to provide information for UI automation in CURRENT_OS
|
||||
self.query_formulator = self._create_agent(self.rag_module_system_prompt)
|
||||
self.llm_search_agent = self._create_agent(self.rag_module_system_prompt)
|
||||
self.knowledge_fusion_agent = self._create_agent(self.rag_module_system_prompt)
|
||||
|
||||
self.use_image_for_search = use_image_for_search
|
||||
|
||||
def retrieve_knowledge(
|
||||
self, instruction: str, search_query: str, search_engine: str = "llm"
|
||||
) -> Tuple[str, str]:
|
||||
"""Retrieve knowledge using search engine
|
||||
Args:
|
||||
instruction (str): task instruction
|
||||
observation (Dict): current observation
|
||||
search_engine (str): search engine to use"""
|
||||
|
||||
# Use search engine to retrieve knowledge based on the formulated query
|
||||
search_results = self._search(instruction, search_query, search_engine)
|
||||
|
||||
return search_query, search_results
|
||||
|
||||
def formulate_query(self, instruction: str, observation: Dict) -> str:
|
||||
"""Formulate search query based on instruction and current state"""
|
||||
query_path = os.path.join(
|
||||
self.local_kb_path, self.platform, "formulate_query.json"
|
||||
)
|
||||
try:
|
||||
with open(query_path, "r") as f:
|
||||
formulate_query = json.load(f)
|
||||
except:
|
||||
formulate_query = {}
|
||||
|
||||
if instruction in formulate_query:
|
||||
return formulate_query[instruction]
|
||||
|
||||
self.query_formulator.add_message(
|
||||
f"The task is: {instruction}\n"
|
||||
f"Accessibility tree of the current desktop UI state: {observation['linearized_accessibility_tree']}\n"
|
||||
"To use google search to get some useful information, first carefully analyze "
|
||||
"the accessibility tree of the current desktop UI state, then given the task "
|
||||
"instruction, formulate a question that can be used to search on the Internet "
|
||||
"for information in helping with the task execution.\n"
|
||||
"The question should not be too general or too specific. Please ONLY provide "
|
||||
"the question.\nQuestion:",
|
||||
image_content=(
|
||||
observation["screenshot"]
|
||||
if self.use_image_for_search and "screenshot" in observation
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
search_query = self.query_formulator.get_response().strip().replace('"', "")
|
||||
print("search query: ", search_query)
|
||||
formulate_query[instruction] = search_query
|
||||
with open(query_path, "w") as f:
|
||||
json.dump(formulate_query, f, indent=2)
|
||||
|
||||
return search_query
|
||||
|
||||
def _search(self, instruction: str, search_query: str, search_engine: str) -> str:
|
||||
"""Execute search using specified engine"""
|
||||
|
||||
# Default to perplexica rag knowledge to see if the query exists
|
||||
file = os.path.join(
|
||||
self.local_kb_path, self.platform, f"{search_engine}_rag_knowledge.json"
|
||||
)
|
||||
|
||||
try:
|
||||
with open(file, "r") as f:
|
||||
exist_search_results = json.load(f)
|
||||
except:
|
||||
exist_search_results = {}
|
||||
|
||||
if instruction in exist_search_results:
|
||||
return exist_search_results[instruction]
|
||||
if search_engine.lower() == "llm":
|
||||
# Use LLM's internal knowledge like a search engine
|
||||
self.llm_search_agent.add_message(search_query)
|
||||
search_results = self.llm_search_agent.get_response()
|
||||
elif search_engine.lower() == "perplexica":
|
||||
# Use perplexica to search for the query
|
||||
search_results = query_to_perplexica(search_query)
|
||||
else:
|
||||
raise ValueError(f"Unsupported search engine: {search_engine}")
|
||||
|
||||
exist_search_results[instruction] = search_results.strip()
|
||||
with open(
|
||||
os.path.join(
|
||||
self.local_kb_path,
|
||||
self.platform,
|
||||
f"{search_engine}_rag_knowledge.json",
|
||||
),
|
||||
"w",
|
||||
) as f:
|
||||
json.dump(exist_search_results, f, indent=2)
|
||||
|
||||
return search_results
|
||||
|
||||
def retrieve_narrative_experience(self, instruction: str) -> Tuple[str, str]:
|
||||
"""Retrieve narrative experience using embeddings"""
|
||||
knowledge_base = load_knowledge_base(self.narrative_memory_path)
|
||||
if not knowledge_base:
|
||||
return "None", "None"
|
||||
|
||||
embeddings = load_embeddings(self.embeddings_path)
|
||||
|
||||
# Get or create instruction embedding
|
||||
instruction_embedding = embeddings.get(instruction)
|
||||
|
||||
if instruction_embedding is None:
|
||||
instruction_embedding = self.embedding_engine.get_embeddings(instruction)
|
||||
embeddings[instruction] = instruction_embedding
|
||||
|
||||
# Get or create embeddings for knowledge base entries
|
||||
candidate_embeddings = []
|
||||
for key in knowledge_base:
|
||||
candidate_embedding = embeddings.get(key)
|
||||
if candidate_embedding is None:
|
||||
candidate_embedding = self.embedding_engine.get_embeddings(key)
|
||||
embeddings[key] = candidate_embedding
|
||||
|
||||
candidate_embeddings.append(candidate_embedding)
|
||||
|
||||
save_embeddings(self.embeddings_path, embeddings)
|
||||
|
||||
similarities = cosine_similarity(
|
||||
instruction_embedding, np.vstack(candidate_embeddings)
|
||||
)[0]
|
||||
sorted_indices = np.argsort(similarities)[::-1]
|
||||
|
||||
keys = list(knowledge_base.keys())
|
||||
idx = 1 if keys[sorted_indices[0]] == instruction else 0
|
||||
return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]
|
||||
|
||||
def retrieve_episodic_experience(self, instruction: str) -> Tuple[str, str]:
|
||||
"""Retrieve similar task experience using embeddings"""
|
||||
knowledge_base = load_knowledge_base(self.episodic_memory_path)
|
||||
if not knowledge_base:
|
||||
return "None", "None"
|
||||
|
||||
embeddings = load_embeddings(self.embeddings_path)
|
||||
|
||||
# Get or create instruction embedding
|
||||
instruction_embedding = embeddings.get(instruction)
|
||||
|
||||
if instruction_embedding is None:
|
||||
instruction_embedding = self.embedding_engine.get_embeddings(instruction)
|
||||
embeddings[instruction] = instruction_embedding
|
||||
|
||||
# Get or create embeddings for knowledge base entries
|
||||
candidate_embeddings = []
|
||||
for key in knowledge_base:
|
||||
candidate_embedding = embeddings.get(key)
|
||||
if candidate_embedding is None:
|
||||
candidate_embedding = self.embedding_engine.get_embeddings(key)
|
||||
embeddings[key] = candidate_embedding
|
||||
|
||||
candidate_embeddings.append(candidate_embedding)
|
||||
|
||||
save_embeddings(self.embeddings_path, embeddings)
|
||||
|
||||
similarities = cosine_similarity(
|
||||
instruction_embedding, np.vstack(candidate_embeddings)
|
||||
)[0]
|
||||
sorted_indices = np.argsort(similarities)[::-1]
|
||||
|
||||
keys = list(knowledge_base.keys())
|
||||
idx = 1 if keys[sorted_indices[0]] == instruction else 0
|
||||
return keys[sorted_indices[idx]], knowledge_base[keys[sorted_indices[idx]]]
|
||||
|
||||
def knowledge_fusion(
|
||||
self,
|
||||
observation: Dict,
|
||||
instruction: str,
|
||||
web_knowledge: str,
|
||||
similar_task: str,
|
||||
experience: str,
|
||||
) -> str:
|
||||
"""Combine web knowledge with similar task experience"""
|
||||
self.knowledge_fusion_agent.add_message(
|
||||
f"Task: {instruction}\n"
|
||||
f"Accessibility tree of the current desktop UI state: {observation['linearized_accessibility_tree']}\n"
|
||||
f"**Web search result**:\n{web_knowledge}\n\n"
|
||||
f"**Retrieved similar task experience**:\n"
|
||||
f"Similar task:{similar_task}\n{experience}\n\n"
|
||||
f"Based on the web search result and the retrieved similar task experience, "
|
||||
f"if you think the similar task experience is indeed useful to the main task, "
|
||||
f"integrate it with the web search result. Provide the final knowledge in a numbered list.",
|
||||
image_content=(
|
||||
observation["screenshot"]
|
||||
if self.use_image_for_search and "screenshot" in observation
|
||||
else None
|
||||
),
|
||||
)
|
||||
return self.knowledge_fusion_agent.get_response()
|
||||
@@ -0,0 +1,280 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI
|
||||
from gui_agents.s1.core.BaseModule import BaseModule
|
||||
from gui_agents.s1.core.Knowledge import KnowledgeBase
|
||||
from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s1.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,
|
||||
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 submodules of the Manager
|
||||
self.generator_agent = self._create_agent(PROCEDURAL_MEMORY.MANAGER_PROMPT)
|
||||
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.knowledge_base = KnowledgeBase(self.local_kb_path, platform, engine_params)
|
||||
|
||||
self.planner_history = []
|
||||
|
||||
self.turn_count = 0
|
||||
self.search_engine = search_engine
|
||||
self.multi_round = multi_round
|
||||
self.platform = platform
|
||||
|
||||
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)
|
||||
subtask_summarization = call_llm_safe(self.episode_summarization_agent)
|
||||
self.episode_summarization_agent.add_message(subtask_summarization)
|
||||
|
||||
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)
|
||||
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, failure_feedback: str = ""
|
||||
) -> Tuple[Dict, str]:
|
||||
agent = self.grounding_agent
|
||||
|
||||
self.active_apps = agent.get_active_apps(observation)
|
||||
|
||||
tree_input = agent.linearize_and_annotate_tree(observation)
|
||||
observation["linearized_accessibility_tree"] = tree_input
|
||||
|
||||
# Perform Retrieval only at the first planning step
|
||||
if self.turn_count == 0:
|
||||
|
||||
self.search_query = self.knowledge_base.formulate_query(
|
||||
instruction, observation
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
)
|
||||
|
||||
generator_message = (
|
||||
f"Accessibility Tree: {tree_input}\n"
|
||||
f"The clipboard contains: {agent.clipboard}."
|
||||
f"The current open applications are {agent.get_active_apps(observation)}"
|
||||
+ (
|
||||
f" Previous plan failed at step: {failure_feedback}"
|
||||
if failure_feedback
|
||||
else ""
|
||||
)
|
||||
)
|
||||
|
||||
self.generator_agent.add_message(
|
||||
generator_message, image_content=observation.get("screenshot", None)
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
self.planner_history.append(plan)
|
||||
|
||||
self.turn_count += 1
|
||||
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
|
||||
# Set Cost based on GPT-4o
|
||||
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]:
|
||||
# Add initial instruction and plan to the agent's message history
|
||||
self.dag_translator_agent.add_message(
|
||||
f"Instruction: {instruction}\nPlan: {plan}"
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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,
|
||||
failure_feedback: str = None,
|
||||
):
|
||||
"""Generate the action list based on the instruction
|
||||
instruction:str: Instruction for the task
|
||||
"""
|
||||
# Generate the high level plan
|
||||
planner_info, plan = self._generate_step_by_step_plan(
|
||||
observation, instruction, failure_feedback
|
||||
)
|
||||
|
||||
# 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,203 @@
|
||||
import inspect
|
||||
import textwrap
|
||||
|
||||
|
||||
class PROCEDURAL_MEMORY:
|
||||
@staticmethod
|
||||
def construct_worker_procedural_memory(agent_class):
|
||||
procedural_memory = textwrap.dedent(f"""\
|
||||
You are an expert in graphical user interfaces and Python code. You are responsible for executing the current subtask: `SUBTASK_DESCRIPTION` of the larger goal: `TASK_DESCRIPTION`.
|
||||
IMPORTANT: ** The subtasks: ['DONE_TASKS'] have already been done. The future subtasks ['FUTURE_TASKS'] will be done in the future by me. You must only perform the current subtask: `SUBTASK_DESCRIPTION`. Do not try to do future subtasks. **
|
||||
You are working in CURRENT_OS. You must only complete the subtask provided and not the larger goal.
|
||||
You are provided with:
|
||||
1. A simplified accessibility tree of the UI at the current time step.
|
||||
2. A screenshot of the current time step.
|
||||
3. The history of your previous interactions with the UI.
|
||||
4. Access to the following class and methods to interact with the UI:
|
||||
class Agent:
|
||||
""")
|
||||
|
||||
for attr_name in dir(agent_class):
|
||||
attr = getattr(agent_class, attr_name)
|
||||
if callable(attr) and hasattr(attr, "is_agent_action"):
|
||||
# Use inspect to get the full function signature
|
||||
signature = inspect.signature(attr)
|
||||
procedural_memory += f"""
|
||||
def {attr_name}{signature}:
|
||||
'''{attr.__doc__}'''
|
||||
"""
|
||||
|
||||
procedural_memory += textwrap.dedent("""
|
||||
Your response should be formatted like this:
|
||||
(Previous action verification)
|
||||
Carefully analyze based on the screenshot and the accessibility tree if the previous action was successful. If the previous action was not successful, provide a reason for the failure.
|
||||
|
||||
(Screenshot Analysis)
|
||||
Closely examine and describe the current state of the desktop along with the currently open applications.
|
||||
|
||||
(Next Action)
|
||||
Based on the current screenshot, the accessibility tree and the history of your previous interaction with the UI, decide on the next action in natural language to accomplish the given task.
|
||||
|
||||
(Grounded Action)
|
||||
Translate the next action into code using the provided API methods. Format the code like this:
|
||||
```python
|
||||
agent.click(123, 1, "left")
|
||||
```
|
||||
Note for the code:
|
||||
1. Only perform one action at a time.
|
||||
2. Do not put anything other than python code in the block. You can only use one function call at a time. Do not put more than one function call in the block.
|
||||
3. You must use only the available methods provided above to interact with the UI, do not invent new methods.
|
||||
3. Only return one code block every time. There must be a single line of code in the code block.
|
||||
4. Please only use the available methods provided above to interact with the UI.
|
||||
5. If you think the task is already completed, you can return `agent.done()` in the code block.
|
||||
6. If you think the task cannot be completed, you can return `agent.fail()` in the code block.
|
||||
7. Do not do anything other than the exact specified task. Return with `agent.done()` immediately after the task is completed or `agent.fail()` if it cannot be completed.
|
||||
8. Whenever possible use hot-keys or typing rather than mouse clicks.
|
||||
9. My computer's password is 'password', feel free to use it when you need sudo rights
|
||||
""")
|
||||
return procedural_memory.strip()
|
||||
|
||||
# MANAGER_PROMPT = """You are a planning agent for solving GUI navigation tasks. You will be provided the initial configuration of a system including accessibility, screenshot and other information. You need to solve the following task: TASK_DESCRIPTION. You will describe in as much detail as possible the steps required to complete the task by a GUI agent. Please do not include any verification steps in your plan that is not your responsibility. IMPORTANT: Your plan should be as concize as possible and should not include any unnecessary steps. Do not fine-tune, or embellish anything or cause any side effects. Generate the plan that can be accomplished in the shortest time. Please take the current state into account when generating the plan. Please provide the plan in a step-by-step format and make sure you do not include anything that's already done in the GUI in your plan."""
|
||||
|
||||
# TODO: exploring this prompt
|
||||
MANAGER_PROMPT = """You are a planning agent for solving GUI navigation tasks. You will be provided the initial configuration of a system including accessibility, screenshot and other information. You need to solve the following task: TASK_DESCRIPTION. You will describe in as much detail as possible the steps required to complete the task by a GUI agent. Please do not include any verification steps in your plan that is not your responsibility. IMPORTANT: Your plan should be as concize as possible and should not include any unnecessary steps. Do not fine-tune, or embellish anything or cause any side effects. Generate the plan that can be accomplished in the shortest time. Please take the current state into account when generating the plan. Please provide the plan in a step-by-step format and make sure you do not include anything that's already done in the GUI in your plan. You don't need to arrange the steps in order just list out everything that needs to be done. You may follow a dependency structure. Note that the execution agent that will complete your plan can't actually see everything thats visible to you."""
|
||||
|
||||
# NOTE: below prompt results in suboptimal initial plans
|
||||
# MANAGER_PROMPT = """You are an expert planning agent for GUI tasks. You will be provided with an initial state of the system including accessibility, screenshot and other information and the final state represented by the task: TASK_DESCRIPTION. Tell me everything that needs to be done in order to reach the goal state. You don't need to arrange the steps in order just list out everything that needs to be done. You may follow a dependency structure."""
|
||||
|
||||
# USED IN OSWORLD EXPERIMENTS
|
||||
RAG_AGENT_OSWORLD = """
|
||||
Given a desktop computer task instruction, you are an agent which should provide useful information as requested, to help another agent follow the instruction and perform the task.
|
||||
The domain of the desktop computer task is from [CURRENT_OS, VLC, LibreOffice, Chrome, Thunderbird, VS Code, GIMP].
|
||||
The task is: TASK_DESCRIPTION
|
||||
The simplified accessibility tree of the current computer UI is: ACCESSIBLITY_TREE
|
||||
"""
|
||||
|
||||
RAG_AGENT = """
|
||||
Given a desktop computer task instruction, you are an agent which should provide useful information as requested, to help another agent follow the instruction and perform the task in CURRENT_OS.
|
||||
"""
|
||||
|
||||
# TODO: confirm this prompt
|
||||
REFLECTION_ON_TRAJECTORY = """
|
||||
You are a reflection agent designed to assist in task execution by analyzing a trajectory of task execution until this time step and providing feedback for the next step prediction.
|
||||
You have access to the Task Description and Current Trajectory, and the image for each step. The most recent image is what happened after the latest action in the trajectory.
|
||||
You should ONLY provide informative reflection feedback (potential mitigation alternatives) based on your expertise for the planning agent when you observe the abnormal trajectory (e.g., contain consecutive failures).
|
||||
Otherwise, let the agent continue to proceed as planned.
|
||||
Make sure to avoid providing any information about specific planning or actions and avoid generating repeated reflection feedbacks.
|
||||
Assume the grounded action is correct, do not judge about it.
|
||||
"""
|
||||
|
||||
TASK_SUMMARIZATION_PROMPT = """
|
||||
You are a summarization agent designed to analyze a trajectory of desktop task execution.
|
||||
You have access to the Task Description and Whole Trajectory including plan, verification and reflection at each step.
|
||||
Your summarized information will be referred to by another agent when performing the tasks.
|
||||
You should follow the below instructions:
|
||||
1. If the task is successfully executed, you should summarize the successful plan based on the whole trajectory to finish the task.
|
||||
2. Otherwise, provide the reasons why the task is failed and potential suggestions that may avoid this failure.
|
||||
|
||||
**ATTENTION**
|
||||
1. Only extract the correct plan and do not provide redundant steps.
|
||||
2. Do not contain grounded actions in the plan.
|
||||
3. If there are the successfully used hot-keys, make sure to include them in the plan.
|
||||
4. The suggestions are for another agent not human, so they must be doable through the agent's action.
|
||||
5. Don't generate high-level suggestions (e.g., Implement Error Handling).
|
||||
"""
|
||||
|
||||
# DAG_TRANSLATOR_PROMPT = """You are a plan to Dependency Graph conversion agent. You will be provided a plan and you will generate a directed acyclic graph in the specified format for the plan. Each node in your graph should contain two fields name and subinfo. name is a one line description of each subtask. subinfo is all available information about executing that subtask available in the step by step plan. Please do not remove or edit any information out of the subinfo. The graph must be a directed acyclic graph. The graph must be connected. Do not include any repeated or optional steps in the graph, any extra info must go in the subinfo.
|
||||
# """
|
||||
|
||||
DAG_TRANSLATOR_PROMPT = """You are a plan to Dependency Graph conversion agent. Your task is to analyze a given plan and generate a structured JSON output representing the plan and its corresponding directed acyclic graph (DAG).
|
||||
|
||||
The output should be a valid JSON object wrapped in <json></json> tags, with the following structure:
|
||||
|
||||
<json>
|
||||
{
|
||||
"dag": {
|
||||
"nodes": [
|
||||
{
|
||||
"name": "Short name or brief description of the step",
|
||||
"info": "Detailed information about executing this step"
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
[
|
||||
{"name": "Name of the source node", "info": "Info of the source node"},
|
||||
{"name": "Name of the target node", "info": "Info of the target node"}
|
||||
]
|
||||
]
|
||||
}
|
||||
}
|
||||
</json>
|
||||
|
||||
Guidelines:
|
||||
1. The "plan" field should contain the entire original plan as a string.
|
||||
2. In the "dag" object:
|
||||
a. Each node in the "nodes" array should contain 'name' and 'info' fields.
|
||||
b. 'name' should be a concise, one-line description of the subtask.
|
||||
c. 'info' should contain all available information about executing that subtask from the original plan. Do not remove or edit any information from the 'info' field.
|
||||
3. The "edges" array should represent the connections between nodes, showing the order and dependencies of the steps.
|
||||
4. The graph must be a directed acyclic graph (DAG) and must be connected.
|
||||
5. Do not include repeated or optional steps in the graph. Any extra information should be incorporated into the 'info' field of the relevant node.
|
||||
|
||||
Analyze the given plan and provide the output in this JSON format within the <json></json> tags. Ensure the JSON is valid and properly escaped.
|
||||
"""
|
||||
|
||||
SUBTASK_SUMMARIZATION_PROMPT = """
|
||||
You are a summarization agent designed to analyze a trajectory of desktop task execution.
|
||||
You will summarize the correct plan and grounded actions based on the whole trajectory of a subtask, ensuring the summarized plan contains only correct and necessary steps.
|
||||
|
||||
**ATTENTION**
|
||||
1. Summarize the correct plan and its corresponding grounded actions. Carefully filter out any repeated or incorrect steps based on the verification output in the trajectory. Only include the necessary steps for successfully completing the subtask.
|
||||
2. ID Replacement in Grounded Actions:
|
||||
When summarizing grounded actions, replace all actual IDs with placeholders element1_id, element2_id, etc., while maintaining the total number of parameters.
|
||||
Ensure the placeholders (element1_id, element2_id, …) follow the order of appearance in the grounded actions.
|
||||
3. Only generate grounded actions that are explicitly present in the trajectory. Do not introduce any grounded actions that do not exist in the trajectory.
|
||||
4. For each step in the plan, provide a corresponding grounded action. Use the exact format:
|
||||
Action: [Description of the correct action]
|
||||
Grounded Action: [Grounded actions with element_id replacement]
|
||||
5. Exclude any other details that are not necessary for completing the task.
|
||||
"""
|
||||
|
||||
STATE_EVALUATOR_SYSTEM_PROMPT = """
|
||||
You are an impartial evaluator to evaluate the completeness of the given desktop computer task, you are also an expert of accessibility tree, os environment and python programming.
|
||||
The task is: TASK_DESCRIPTION, it is executed by a digital agent who can perform the task without knowing whether the task requirements are met.
|
||||
As an evaluator, your task is to judge whether the task is finished and meets the task requirement.
|
||||
You have access to the:
|
||||
1. Task instruction.
|
||||
2. The whole actions performed by the digital agent.
|
||||
3. The accessibility tree at the first step and the last step.
|
||||
4. The screenshot at the first step and the last step.
|
||||
|
||||
You are able to proceed your judgment process in the following ways based on the task instruction:
|
||||
1. By comparing the difference in the accessibility trees of the UI, you should judge whether the task is complete given the task instruction.
|
||||
2. If you cannot judge based on the observations, you can evalaute it by writing and running a python script to do a further examination. For example, you can use the 'subprocess' module to run the external command in a terminal to check whether an application has been installed.
|
||||
You can also call the file system API to do the file check, etc. You can also try to interactive with the environment via other methods or interface you are familiared with.
|
||||
|
||||
**IMPORTANT**
|
||||
1. If no python script is needed, you should provide your analysis and put the judgment at the end of the response in this format: Judgment: Yes/No
|
||||
2. Otherwise, you should format your response into two parts as shown below:
|
||||
```python
|
||||
# your code script here
|
||||
```
|
||||
|
||||
**ATTENTION**
|
||||
1. You should only use scripts when you have to.
|
||||
2. When you generate code script, only return one code block every time, the code block should contain the whole script you want to run. You must guarantee that the script is comprehensive and executable, make sure to print out the scripts' results for subsequent judgement.
|
||||
Additionally, the comment of the code is **PROHIBITED**
|
||||
3. You should strictly follow the response format mentioned above.
|
||||
|
||||
**SUBSEQUENCE**
|
||||
If you have generated the python script, I will execute it and return the corresponding result to you (Started with "The output after executing the script is:..."). Then you should judge whether the task has been completed or not comprehensively based on the script and its result,
|
||||
the task information, and the comparison of accessibility trees and screenshots. Provide your analysis and put the judgment at the end of the response in this format: Judgment: Yes/No
|
||||
"""
|
||||
|
||||
OBS_EVALUATOR_SYSTEM_PROMPT = """
|
||||
You are an impartial evaluator to evaluate the completeness of the given desktop computer task.
|
||||
The task is: TASK_DESCRIPTION, it is executed by a digital agent who can perform the task without knowing whether the task requirements are met.
|
||||
As an evaluator, your task is to judge whether the task is finished and meets the task requirement.
|
||||
You have access to the task instruction, the whole actions performed by the digital agent, the accessibility tree of the UI and screenshot at the first time step and the last time step.
|
||||
By comparing the difference in the accessibility trees of the UI, you should judge whether the task is complete given the task instruction.
|
||||
Provide your analysis and put the judgment at the end of the response in this format:
|
||||
Judgment: Yes/No
|
||||
Only say Yes or No in the Judgment section. Do not provide any other information in the Judgment section.
|
||||
"""
|
||||
@@ -0,0 +1,256 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Dict, List, Tuple
|
||||
import platform
|
||||
|
||||
from gui_agents.s1.aci.ACI import ACI
|
||||
from gui_agents.s1.core.BaseModule import BaseModule
|
||||
from gui_agents.s1.core.Knowledge import KnowledgeBase
|
||||
from gui_agents.s1.core.ProceduralMemory import PROCEDURAL_MEMORY
|
||||
from gui_agents.s1.utils import common_utils
|
||||
from gui_agents.s1.utils.common_utils import Node, calculate_tokens, call_llm_safe
|
||||
|
||||
logger = logging.getLogger("desktopenv.agent")
|
||||
|
||||
|
||||
class Worker(BaseModule):
|
||||
def __init__(
|
||||
self,
|
||||
engine_params: Dict,
|
||||
grounding_agent: ACI,
|
||||
local_kb_path: str,
|
||||
platform: str = platform.system().lower(),
|
||||
search_engine: str = "perplexica",
|
||||
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
|
||||
search_engine: str
|
||||
The search engine to use
|
||||
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.enable_reflection = enable_reflection
|
||||
self.search_engine = search_engine
|
||||
self.use_subtask_experience = use_subtask_experience
|
||||
self.reset()
|
||||
|
||||
def flush_messages(self, n):
|
||||
# After every max_trajectory_length trajectories, remove messages from the start except the system prompt
|
||||
for agent in [self.generator_agent]:
|
||||
if len(agent.messages) > 2 * n + 1:
|
||||
# Remove the user message and assistant message, both are 1 because the elements will move back after 1 pop
|
||||
agent.remove_message_at(1)
|
||||
agent.remove_message_at(1)
|
||||
|
||||
def reset(self):
|
||||
self.generator_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.construct_worker_procedural_memory(
|
||||
type(self.grounding_agent)
|
||||
).replace("CURRENT_OS", self.platform)
|
||||
)
|
||||
self.reflection_agent = self._create_agent(
|
||||
PROCEDURAL_MEMORY.REFLECTION_ON_TRAJECTORY
|
||||
)
|
||||
|
||||
self.knowledge_base = KnowledgeBase(
|
||||
local_kb_path=self.local_kb_path,
|
||||
platform=self.platform,
|
||||
engine_params=self.engine_params,
|
||||
)
|
||||
|
||||
self.turn_count = 0
|
||||
self.planner_history = []
|
||||
self.reflections = []
|
||||
self.cost_this_turn = 0
|
||||
self.tree_inputs = []
|
||||
self.screenshot_inputs = []
|
||||
|
||||
# TODO: Experimental
|
||||
def remove_ids_from_history(self):
|
||||
for message in self.generator_agent.messages:
|
||||
if message["role"] == "user":
|
||||
for content in message["content"]:
|
||||
if content["type"] == "text":
|
||||
# Regex pattern to match lines that start with a number followed by spaces and remove the number
|
||||
pattern = r"^\d+\s+"
|
||||
|
||||
# Apply the regex substitution on each line
|
||||
processed_lines = [
|
||||
re.sub(pattern, "", line)
|
||||
for line in content["text"].splitlines()
|
||||
]
|
||||
|
||||
# Join the processed lines back into a single string
|
||||
result = "\n".join(processed_lines)
|
||||
|
||||
result = result.replace("id\t", "")
|
||||
|
||||
# replace message content
|
||||
content["text"] = result
|
||||
|
||||
def generate_next_action(
|
||||
self,
|
||||
instruction: str,
|
||||
search_query: str,
|
||||
subtask: str,
|
||||
subtask_info: str,
|
||||
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
|
||||
|
||||
self.active_apps = agent.get_active_apps(obs)
|
||||
|
||||
# Get RAG knowledge, only update system message at t=0
|
||||
if self.turn_count == 0:
|
||||
# TODO: uncomment and fix for subtask level RAG
|
||||
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)
|
||||
)
|
||||
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))
|
||||
)
|
||||
|
||||
# Clear older messages - we keep full context. if you want to keep only the last n messages, you can use the flush_messages function
|
||||
# self.flush_messages(3) # flushes generator messages
|
||||
|
||||
# Reflection generation
|
||||
reflection = None
|
||||
if self.enable_reflection and self.turn_count > 0:
|
||||
# TODO: reuse planner history
|
||||
self.reflection_agent.add_message(
|
||||
"Task Description: "
|
||||
+ subtask
|
||||
+ " Instruction: "
|
||||
+ subtask_info
|
||||
+ "\n"
|
||||
+ "Current Trajectory: "
|
||||
+ "\n\n".join(self.planner_history)
|
||||
+ "\n"
|
||||
)
|
||||
reflection = call_llm_safe(self.reflection_agent)
|
||||
self.reflections.append(reflection)
|
||||
self.reflection_agent.add_message(reflection)
|
||||
|
||||
logger.info("REFLECTION: %s", reflection)
|
||||
|
||||
# Plan Generation
|
||||
tree_input = agent.linearize_and_annotate_tree(obs)
|
||||
|
||||
self.remove_ids_from_history()
|
||||
|
||||
# Bash terminal message.
|
||||
generator_message = (
|
||||
(
|
||||
f"\nYou may use the reflection on the previous trajectory: {reflection}\n"
|
||||
if reflection
|
||||
else ""
|
||||
)
|
||||
+ f"Accessibility Tree: {tree_input}\n"
|
||||
f"Text Buffer = [{','.join(agent.notes)}]. "
|
||||
f"The current open applications are {agent.get_active_apps(obs)} and the active app is {agent.get_top_app(obs)}.\n"
|
||||
)
|
||||
|
||||
print("ACTIVE APP IS: ", agent.get_top_app(obs))
|
||||
# Only provide subinfo in the very first message to avoid over influence and redundancy
|
||||
if self.turn_count == 0:
|
||||
generator_message += f"Remeber 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"]
|
||||
)
|
||||
|
||||
plan = call_llm_safe(self.generator_agent)
|
||||
self.planner_history.append(plan)
|
||||
logger.info("PLAN: %s", plan)
|
||||
|
||||
self.generator_agent.add_message(plan)
|
||||
|
||||
# Calculate input and output tokens
|
||||
input_tokens, output_tokens = calculate_tokens(self.generator_agent.messages)
|
||||
|
||||
# Set Cost based on GPT-4o
|
||||
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)
|
||||
|
||||
# Extract code block from the plan
|
||||
plan_code = common_utils.parse_single_code_from_string(
|
||||
plan.split("Grounded Action")[-1]
|
||||
)
|
||||
plan_code = common_utils.sanitize_code(plan_code)
|
||||
plan_code = common_utils.extract_first_agent_function(plan_code)
|
||||
exec_code = eval(plan_code)
|
||||
|
||||
# If agent selects an element that was out of range, it should not be executed just send a WAIT command.
|
||||
# TODO: should provide this as code feedback to the agent?
|
||||
if agent.index_out_of_range_flag:
|
||||
plan_code = "agent.wait(1.0)"
|
||||
exec_code = eval(plan_code)
|
||||
agent.index_out_of_range_flag = False
|
||||
|
||||
executor_info = {
|
||||
"current_subtask": subtask,
|
||||
"current_subtask_info": subtask_info,
|
||||
"executor_plan": plan,
|
||||
"linearized_accessibility_tree": tree_input,
|
||||
"plan_code": plan_code,
|
||||
"reflection": reflection,
|
||||
"num_input_tokens_executor": input_tokens,
|
||||
"num_output_tokens_executor": output_tokens,
|
||||
"executor_cost": cost,
|
||||
}
|
||||
self.turn_count += 1
|
||||
|
||||
self.tree_inputs.append(tree_input)
|
||||
self.screenshot_inputs.append(obs["screenshot"])
|
||||
|
||||
return executor_info, [exec_code]
|
||||
@@ -0,0 +1,263 @@
|
||||
# Author: Saaket Agashe
|
||||
# Date: 2021-09-15
|
||||
# License: MIT
|
||||
|
||||
import base64
|
||||
import re
|
||||
|
||||
from gui_agents.s1.mllm.MultimodalEngine import (
|
||||
LMMEngineAnthropic,
|
||||
LMMEngineAzureOpenAI,
|
||||
LMMEngineOpenAI,
|
||||
LMMEnginevLLM,
|
||||
)
|
||||
|
||||
data_type_map = {
|
||||
"openai": {"image_url": "image_url"},
|
||||
"anthropic": {"image_url": "image"},
|
||||
}
|
||||
|
||||
|
||||
class LMMAgent:
|
||||
def __init__(self, engine_params=None, system_prompt=None, engine=None):
|
||||
if engine is None:
|
||||
if engine_params is not None:
|
||||
engine_type = engine_params.get("engine_type")
|
||||
if engine_type == "openai":
|
||||
self.engine = LMMEngineOpenAI(**engine_params)
|
||||
elif engine_type == "anthropic":
|
||||
self.engine = LMMEngineAnthropic(**engine_params)
|
||||
elif engine_type == "azure":
|
||||
self.engine = LMMEngineAzureOpenAI(**engine_params)
|
||||
elif engine_type == "vllm":
|
||||
self.engine = LMMEnginevLLM(**engine_params)
|
||||
else:
|
||||
raise ValueError("engine_type must be either 'openai' or 'azure'")
|
||||
else:
|
||||
raise ValueError("engine_params must be provided")
|
||||
else:
|
||||
self.engine = engine
|
||||
|
||||
self.messages = [] # Empty messages
|
||||
|
||||
if system_prompt:
|
||||
self.add_system_prompt(system_prompt)
|
||||
else:
|
||||
self.add_system_prompt("You are a helpful assistant.")
|
||||
|
||||
def encode_image(self, image_content):
|
||||
# if image_content is a path to an image file, check type of the image_content to verify
|
||||
if isinstance(image_content, str):
|
||||
with open(image_content, "rb") as image_file:
|
||||
return base64.b64encode(image_file.read()).decode("utf-8")
|
||||
else:
|
||||
return base64.b64encode(image_content).decode("utf-8")
|
||||
|
||||
def reset(
|
||||
self,
|
||||
):
|
||||
|
||||
self.messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": self.system_prompt}],
|
||||
}
|
||||
]
|
||||
|
||||
def add_system_prompt(self, system_prompt):
|
||||
self.system_prompt = system_prompt
|
||||
if len(self.messages) > 0:
|
||||
self.messages[0] = {
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": self.system_prompt}],
|
||||
}
|
||||
else:
|
||||
self.messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": self.system_prompt}],
|
||||
}
|
||||
)
|
||||
|
||||
def remove_message_at(self, index):
|
||||
"""Remove a message at a given index"""
|
||||
if index < len(self.messages):
|
||||
self.messages.pop(index)
|
||||
|
||||
def replace_message_at(
|
||||
self, index, text_content, image_content=None, image_detail="high"
|
||||
):
|
||||
"""Replace a message at a given index"""
|
||||
if index < len(self.messages):
|
||||
self.messages[index] = {
|
||||
"role": self.messages[index]["role"],
|
||||
"content": [{"type": "text", "text": text_content}],
|
||||
}
|
||||
if image_content:
|
||||
base64_image = self.encode_image(image_content)
|
||||
self.messages[index]["content"].append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{base64_image}",
|
||||
"detail": image_detail,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
def add_message(
|
||||
self, text_content, image_content=None, role=None, image_detail="high"
|
||||
):
|
||||
"""Add a new message to the list of messages"""
|
||||
|
||||
# API-style inference from OpenAI and AzureOpenAI
|
||||
if isinstance(self.engine, (LMMEngineOpenAI, LMMEngineAzureOpenAI)):
|
||||
# infer role from previous message
|
||||
if role != "user":
|
||||
if self.messages[-1]["role"] == "system":
|
||||
role = "user"
|
||||
elif self.messages[-1]["role"] == "user":
|
||||
role = "assistant"
|
||||
elif self.messages[-1]["role"] == "assistant":
|
||||
role = "user"
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": [{"type": "text", "text": text_content}],
|
||||
}
|
||||
|
||||
if image_content:
|
||||
# Check if image_content is a list or a single image
|
||||
if isinstance(image_content, list):
|
||||
# If image_content is a list of images, loop through each image
|
||||
for image in image_content:
|
||||
base64_image = self.encode_image(image)
|
||||
message["content"].append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{base64_image}",
|
||||
"detail": image_detail,
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
# If image_content is a single image, handle it directly
|
||||
base64_image = self.encode_image(image_content)
|
||||
message["content"].append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{base64_image}",
|
||||
"detail": image_detail,
|
||||
},
|
||||
}
|
||||
)
|
||||
self.messages.append(message)
|
||||
|
||||
# For API-style inference from Anthropic
|
||||
elif isinstance(self.engine, LMMEngineAnthropic):
|
||||
# infer role from previous message
|
||||
if role != "user":
|
||||
if self.messages[-1]["role"] == "system":
|
||||
role = "user"
|
||||
elif self.messages[-1]["role"] == "user":
|
||||
role = "assistant"
|
||||
elif self.messages[-1]["role"] == "assistant":
|
||||
role = "user"
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": [{"type": "text", "text": text_content}],
|
||||
}
|
||||
|
||||
if image_content:
|
||||
# Check if image_content is a list or a single image
|
||||
if isinstance(image_content, list):
|
||||
# If image_content is a list of images, loop through each image
|
||||
for image in image_content:
|
||||
base64_image = self.encode_image(image)
|
||||
message["content"].append(
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": base64_image,
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
# If image_content is a single image, handle it directly
|
||||
base64_image = self.encode_image(image_content)
|
||||
message["content"].append(
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": base64_image,
|
||||
},
|
||||
}
|
||||
)
|
||||
self.messages.append(message)
|
||||
|
||||
# Locally hosted vLLM model inference
|
||||
elif isinstance(self.engine, LMMEnginevLLM):
|
||||
# infer role from previous message
|
||||
if role != "user":
|
||||
if self.messages[-1]["role"] == "system":
|
||||
role = "user"
|
||||
elif self.messages[-1]["role"] == "user":
|
||||
role = "assistant"
|
||||
elif self.messages[-1]["role"] == "assistant":
|
||||
role = "user"
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": [{"type": "text", "text": text_content}],
|
||||
}
|
||||
|
||||
if image_content:
|
||||
# Check if image_content is a list or a single image
|
||||
if isinstance(image_content, list):
|
||||
# If image_content is a list of images, loop through each image
|
||||
for image in image_content:
|
||||
base64_image = self.encode_image(image)
|
||||
message["content"].append(
|
||||
{
|
||||
"type": "image",
|
||||
"image": f"data:image;base64,{base64_image}",
|
||||
}
|
||||
)
|
||||
else:
|
||||
# If image_content is a single image, handle it directly
|
||||
base64_image = self.encode_image(image_content)
|
||||
message["content"].append(
|
||||
{"type": "image", "image": f"data:image;base64,{base64_image}"}
|
||||
)
|
||||
self.messages.append(message)
|
||||
|
||||
def get_response(
|
||||
self,
|
||||
user_message=None,
|
||||
image=None,
|
||||
messages=None,
|
||||
temperature=0.0,
|
||||
max_new_tokens=None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Generate the next response based on previous messages"""
|
||||
if messages is None:
|
||||
messages = self.messages
|
||||
if user_message:
|
||||
messages.append(
|
||||
{"role": "user", "content": [{"type": "text", "text": user_message}]}
|
||||
)
|
||||
|
||||
return self.engine.generate(
|
||||
messages,
|
||||
temperature=temperature,
|
||||
max_new_tokens=max_new_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,269 @@
|
||||
# Author: Saaket Agashe
|
||||
# Date: 2021-09-15
|
||||
# License: MIT
|
||||
|
||||
import os
|
||||
import re
|
||||
from io import BytesIO
|
||||
|
||||
import backoff
|
||||
import numpy as np
|
||||
import openai
|
||||
import requests
|
||||
from anthropic import Anthropic
|
||||
from openai import APIConnectionError, APIError, AzureOpenAI, OpenAI, RateLimitError
|
||||
from PIL import Image
|
||||
|
||||
# TODO: Import only if module exists, else ignore
|
||||
# from llava.model.builder import load_pretrained_model
|
||||
# from llava.mm_utils import (
|
||||
# process_images,
|
||||
# tokenizer_image_token,
|
||||
# get_model_name_from_path,
|
||||
# KeywordsStoppingCriteria,
|
||||
# )
|
||||
# from llava.constants import (
|
||||
# IMAGE_TOKEN_INDEX,
|
||||
# DEFAULT_IMAGE_TOKEN,
|
||||
# DEFAULT_IM_START_TOKEN,
|
||||
# DEFAULT_IM_END_TOKEN,
|
||||
# IMAGE_PLACEHOLDER,
|
||||
# )
|
||||
# from llava.conversation import conv_templates, SeparatorStyle
|
||||
|
||||
|
||||
# from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
|
||||
def image_parser(args):
|
||||
out = args.image_file.split(args.sep)
|
||||
return out
|
||||
|
||||
|
||||
def load_image(image_file):
|
||||
if image_file.startswith("http") or image_file.startswith("https"):
|
||||
response = requests.get(image_file)
|
||||
image = Image.open(BytesIO(response.content)).convert("RGB")
|
||||
else:
|
||||
image = Image.open(image_file).convert("RGB")
|
||||
return image
|
||||
|
||||
|
||||
def load_images(image_files):
|
||||
out = []
|
||||
for image_file in image_files:
|
||||
image = load_image(image_file)
|
||||
out.append(image)
|
||||
return out
|
||||
|
||||
|
||||
class LMMEngine:
|
||||
pass
|
||||
|
||||
|
||||
class LMMEngineOpenAI(LMMEngine):
|
||||
def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
|
||||
assert model is not None, "model must be provided"
|
||||
self.model = model
|
||||
|
||||
api_key = api_key or os.getenv("OPENAI_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
|
||||
|
||||
self.llm_client = OpenAI(api_key=self.api_key)
|
||||
|
||||
@backoff.on_exception(
|
||||
backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
|
||||
)
|
||||
def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
|
||||
"""Generate the next message based on previous messages"""
|
||||
return (
|
||||
self.llm_client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=max_new_tokens if max_new_tokens else 4096,
|
||||
temperature=temperature,
|
||||
**kwargs,
|
||||
)
|
||||
.choices[0]
|
||||
.message.content
|
||||
)
|
||||
|
||||
|
||||
class LMMEngineAnthropic(LMMEngine):
|
||||
def __init__(self, api_key=None, model=None, **kwargs):
|
||||
assert model is not None, "model must be provided"
|
||||
self.model = model
|
||||
|
||||
api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"An API Key needs to be provided in either the api_key parameter or as an environment variable named ANTHROPIC_API_KEY"
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
|
||||
self.llm_client = Anthropic(api_key=self.api_key)
|
||||
|
||||
@backoff.on_exception(
|
||||
backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
|
||||
)
|
||||
def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
|
||||
"""Generate the next message based on previous messages"""
|
||||
return (
|
||||
self.llm_client.messages.create(
|
||||
system=messages[0]["content"][0]["text"],
|
||||
model=self.model,
|
||||
messages=messages[1:],
|
||||
max_tokens=max_new_tokens if max_new_tokens else 4096,
|
||||
temperature=temperature,
|
||||
**kwargs,
|
||||
)
|
||||
.content[0]
|
||||
.text
|
||||
)
|
||||
|
||||
|
||||
class OpenAIEmbeddingEngine(LMMEngine):
|
||||
def __init__(
|
||||
self,
|
||||
api_key=None,
|
||||
rate_limit: int = -1,
|
||||
display_cost: bool = True,
|
||||
):
|
||||
"""Init an OpenAI Embedding engine
|
||||
|
||||
Args:
|
||||
api_key (_type_, optional): Auth key from OpenAI. Defaults to None.
|
||||
rate_limit (int, optional): Max number of requests per minute. Defaults to -1.
|
||||
display_cost (bool, optional): Display cost of API call. Defaults to True.
|
||||
"""
|
||||
self.model = "text-embedding-3-small"
|
||||
self.cost_per_thousand_tokens = 0.00002
|
||||
|
||||
api_key = api_key or os.getenv("OPENAI_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
|
||||
)
|
||||
self.api_key = api_key
|
||||
self.display_cost = display_cost
|
||||
self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
|
||||
|
||||
@backoff.on_exception(
|
||||
backoff.expo,
|
||||
(
|
||||
APIError,
|
||||
RateLimitError,
|
||||
APIConnectionError,
|
||||
),
|
||||
)
|
||||
def get_embeddings(self, text: str) -> np.ndarray:
|
||||
client = OpenAI(api_key=self.api_key)
|
||||
response = client.embeddings.create(model=self.model, input=text)
|
||||
if self.display_cost:
|
||||
total_tokens = response.usage.total_tokens
|
||||
cost = self.cost_per_thousand_tokens * total_tokens / 1000
|
||||
# print(f"Total cost for this embedding API call: {cost}")
|
||||
return np.array([data.embedding for data in response.data])
|
||||
|
||||
|
||||
class LMMEngineAzureOpenAI(LMMEngine):
|
||||
def __init__(
|
||||
self,
|
||||
api_key=None,
|
||||
azure_endpoint=None,
|
||||
model=None,
|
||||
api_version=None,
|
||||
rate_limit=-1,
|
||||
**kwargs
|
||||
):
|
||||
assert model is not None, "model must be provided"
|
||||
self.model = model
|
||||
|
||||
assert api_version is not None, "api_version must be provided"
|
||||
self.api_version = api_version
|
||||
|
||||
api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
|
||||
)
|
||||
|
||||
self.api_key = api_key
|
||||
|
||||
azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_API_BASE")
|
||||
if azure_endpoint is None:
|
||||
raise ValueError(
|
||||
"An Azure API endpoint needs to be provided in either the azure_endpoint parameter or as an environment variable named AZURE_OPENAI_API_BASE"
|
||||
)
|
||||
|
||||
self.azure_endpoint = azure_endpoint
|
||||
self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
|
||||
|
||||
self.llm_client = AzureOpenAI(
|
||||
azure_endpoint=self.azure_endpoint,
|
||||
api_key=self.api_key,
|
||||
api_version=self.api_version,
|
||||
)
|
||||
self.cost = 0.0
|
||||
|
||||
# @backoff.on_exception(backoff.expo, (APIConnectionError, APIError, RateLimitError), max_tries=10)
|
||||
def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
|
||||
"""Generate the next message based on previous messages"""
|
||||
completion = self.llm_client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=max_new_tokens if max_new_tokens else 4096,
|
||||
temperature=temperature,
|
||||
**kwargs,
|
||||
)
|
||||
total_tokens = completion.usage.total_tokens
|
||||
self.cost += 0.02 * ((total_tokens + 500) / 1000)
|
||||
return completion.choices[0].message.content
|
||||
|
||||
|
||||
class LMMEnginevLLM(LMMEngine):
|
||||
def __init__(
|
||||
self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
|
||||
):
|
||||
assert model is not None, "model must be provided"
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
|
||||
self.base_url = base_url or os.getenv("vLLM_ENDPOINT_URL")
|
||||
if self.base_url is None:
|
||||
raise ValueError(
|
||||
"An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named vLLM_ENDPOINT_URL"
|
||||
)
|
||||
|
||||
self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
|
||||
|
||||
self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)
|
||||
|
||||
# @backoff.on_exception(backoff.expo, (APIConnectionError, APIError, RateLimitError), max_tries=10)
|
||||
# TODO: Default params chosen for the Qwen model
|
||||
def generate(
|
||||
self,
|
||||
messages,
|
||||
temperature=0.0,
|
||||
top_p=0.8,
|
||||
repetition_penalty=1.05,
|
||||
max_new_tokens=512,
|
||||
**kwargs
|
||||
):
|
||||
"""Generate the next message based on previous messages"""
|
||||
completion = self.llm_client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=max_new_tokens if max_new_tokens else 4096,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
extra_body={"repetition_penalty": repetition_penalty},
|
||||
)
|
||||
return completion.choices[0].message.content
|
||||
@@ -0,0 +1,863 @@
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import tempfile
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from io import BytesIO
|
||||
from typing import Dict, List, Tuple, Union
|
||||
from xml.etree.ElementTree import Element
|
||||
|
||||
import numpy as np
|
||||
import tiktoken
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
|
||||
def find_leaf_nodes(xlm_file_str):
|
||||
if not xlm_file_str:
|
||||
return []
|
||||
|
||||
root = ET.fromstring(xlm_file_str)
|
||||
|
||||
# Recursive function to traverse the XML tree and collect leaf nodes
|
||||
def collect_leaf_nodes(node, leaf_nodes):
|
||||
# If the node has no children, it is a leaf node, add it to the list
|
||||
if not list(node):
|
||||
leaf_nodes.append(node)
|
||||
# If the node has children, recurse on each child
|
||||
for child in node:
|
||||
collect_leaf_nodes(child, leaf_nodes)
|
||||
|
||||
# List to hold all leaf nodes
|
||||
leaf_nodes = []
|
||||
collect_leaf_nodes(root, leaf_nodes)
|
||||
return leaf_nodes
|
||||
|
||||
|
||||
state_ns = "uri:deskat:state.at-spi.gnome.org"
|
||||
component_ns = "uri:deskat:component.at-spi.gnome.org"
|
||||
|
||||
|
||||
class Node(BaseModel):
|
||||
name: str
|
||||
info: str
|
||||
|
||||
|
||||
class Dag(BaseModel):
|
||||
nodes: List[Node]
|
||||
edges: List[List[Node]]
|
||||
|
||||
|
||||
NUM_IMAGE_TOKEN = 1105 # Value set of screen of size 1920x1080 for openai vision
|
||||
|
||||
|
||||
def call_llm_safe(agent) -> Union[str, Dag]:
|
||||
# Retry if fails
|
||||
max_retries = 3 # Set the maximum number of retries
|
||||
attempt = 0
|
||||
response = ""
|
||||
while attempt < max_retries:
|
||||
try:
|
||||
response = agent.get_response()
|
||||
break # If successful, break out of the loop
|
||||
except Exception as e:
|
||||
attempt += 1
|
||||
print(f"Attempt {attempt} failed: {e}")
|
||||
if attempt == max_retries:
|
||||
print("Max retries reached. Handling failure.")
|
||||
time.sleep(1.0)
|
||||
return response
|
||||
|
||||
|
||||
def calculate_tokens(messages, num_image_token=NUM_IMAGE_TOKEN) -> Tuple[int, int]:
|
||||
|
||||
num_input_images = 0
|
||||
output_message = messages[-1]
|
||||
|
||||
input_message = messages[:-1]
|
||||
|
||||
input_string = """"""
|
||||
for message in input_message:
|
||||
input_string += message["content"][0]["text"] + "\n"
|
||||
if len(message["content"]) > 1:
|
||||
num_input_images += 1
|
||||
|
||||
input_text_tokens = get_input_token_length(input_string)
|
||||
|
||||
input_image_tokens = num_image_token * num_input_images
|
||||
|
||||
output_tokens = get_input_token_length(output_message["content"][0]["text"])
|
||||
|
||||
return (input_text_tokens + input_image_tokens), output_tokens
|
||||
|
||||
|
||||
def judge_node(node: Element, platform="ubuntu", check_image=False) -> bool:
|
||||
keeps: bool = (
|
||||
node.tag.startswith("document")
|
||||
or node.tag.endswith("item")
|
||||
or node.tag.endswith("button")
|
||||
or node.tag.endswith("heading")
|
||||
or node.tag.endswith("label")
|
||||
or node.tag.endswith("scrollbar")
|
||||
or node.tag.endswith("searchbox")
|
||||
or node.tag.endswith("textbox")
|
||||
or node.tag.endswith("link")
|
||||
or node.tag.endswith("tabelement")
|
||||
or node.tag.endswith("textfield")
|
||||
or node.tag.endswith("textarea")
|
||||
or node.tag.endswith("menu")
|
||||
or node.tag.endswith("menu-item")
|
||||
or node.tag
|
||||
in {
|
||||
"alert",
|
||||
"canvas",
|
||||
"check-box",
|
||||
"combo-box",
|
||||
"entry",
|
||||
"icon",
|
||||
"image",
|
||||
"paragraph",
|
||||
"scroll-bar",
|
||||
"section",
|
||||
"slider",
|
||||
"static",
|
||||
"table-cell",
|
||||
"terminal",
|
||||
"text",
|
||||
"netuiribbontab",
|
||||
"start",
|
||||
"trayclockwclass",
|
||||
"traydummysearchcontrol",
|
||||
"uiimage",
|
||||
"uiproperty",
|
||||
"uiribboncommandbar",
|
||||
}
|
||||
)
|
||||
|
||||
keeps = (
|
||||
keeps
|
||||
and (
|
||||
platform == "ubuntu"
|
||||
and node.get("{{{:}}}showing".format(state_ns), "false") == "true"
|
||||
and node.get("{{{:}}}visible".format(state_ns), "false") == "true"
|
||||
or platform == "windows"
|
||||
and node.get("{{{:}}}visible".format(state_ns), "false") == "true"
|
||||
)
|
||||
and (
|
||||
node.get("name", "") != ""
|
||||
or node.text is not None
|
||||
and len(node.text) > 0
|
||||
or check_image
|
||||
and node.get("image", "false") == "true"
|
||||
)
|
||||
)
|
||||
# and (node.get("{{{:}}}enabled".format(state_ns), "false") == "true" \
|
||||
# or node.get("{{{:}}}editable".format(state_ns), "false") == "true" \
|
||||
# or node.get("{{{:}}}expandable".format(state_ns), "false") == "true" \
|
||||
# or node.get("{{{:}}}checkable".format(state_ns), "false") == "true"
|
||||
# ) \
|
||||
|
||||
coordinates: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}screencoord".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
sizes: Tuple[int, int] = eval(
|
||||
node.get("{{{:}}}size".format(component_ns), "(-1, -1)")
|
||||
)
|
||||
keeps = (
|
||||
keeps
|
||||
and coordinates[0] >= 0
|
||||
and coordinates[1] >= 0
|
||||
and sizes[0] > 0
|
||||
and sizes[1] > 0
|
||||
)
|
||||
return keeps
|
||||
|
||||
|
||||
def filter_nodes(root: Element, platform="ubuntu", check_image=False):
|
||||
filtered_nodes = []
|
||||
all_nodes = []
|
||||
for node in root.iter():
|
||||
all_nodes.append(node)
|
||||
|
||||
for node in root.iter():
|
||||
if judge_node(node, platform, check_image):
|
||||
filtered_nodes.append(node)
|
||||
|
||||
return filtered_nodes
|
||||
|
||||
|
||||
def draw_bounding_boxes(nodes, image_file_content, down_sampling_ratio=1.0):
|
||||
# Load the screenshot image
|
||||
image_stream = io.BytesIO(image_file_content)
|
||||
image = Image.open(image_stream)
|
||||
if float(down_sampling_ratio) != 1.0:
|
||||
image = image.resize(
|
||||
(
|
||||
int(image.size[0] * down_sampling_ratio),
|
||||
int(image.size[1] * down_sampling_ratio),
|
||||
)
|
||||
)
|
||||
draw = ImageDraw.Draw(image)
|
||||
marks = []
|
||||
drew_nodes = []
|
||||
text_informations: List[str] = ["index\ttag\tname\ttext"]
|
||||
|
||||
try:
|
||||
# Adjust the path to the font file you have or use a default one
|
||||
font = ImageFont.truetype("arial.ttf", 15)
|
||||
except IOError:
|
||||
# Fallback to a basic font if the specified font can't be loaded
|
||||
font = ImageFont.load_default()
|
||||
|
||||
index = 1
|
||||
|
||||
# Loop over all the visible nodes and draw their bounding boxes
|
||||
for _node in nodes:
|
||||
coords_str = _node.attrib.get(
|
||||
"{uri:deskat:component.at-spi.gnome.org}screencoord"
|
||||
)
|
||||
size_str = _node.attrib.get("{uri:deskat:component.at-spi.gnome.org}size")
|
||||
|
||||
if coords_str and size_str:
|
||||
try:
|
||||
# Parse the coordinates and size from the strings
|
||||
coords = tuple(map(int, coords_str.strip("()").split(", ")))
|
||||
size = tuple(map(int, size_str.strip("()").split(", ")))
|
||||
|
||||
import copy
|
||||
|
||||
original_coords = copy.deepcopy(coords)
|
||||
original_size = copy.deepcopy(size)
|
||||
|
||||
if float(down_sampling_ratio) != 1.0:
|
||||
# Downsample the coordinates and size
|
||||
coords = tuple(int(coord * down_sampling_ratio) for coord in coords)
|
||||
size = tuple(int(s * down_sampling_ratio) for s in size)
|
||||
|
||||
# Check for negative sizes
|
||||
if size[0] <= 0 or size[1] <= 0:
|
||||
raise ValueError(f"Size must be positive, got: {size}")
|
||||
|
||||
# Calculate the bottom-right corner of the bounding box
|
||||
bottom_right = (coords[0] + size[0], coords[1] + size[1])
|
||||
|
||||
# Check that bottom_right > coords (x1 >= x0, y1 >= y0)
|
||||
if bottom_right[0] < coords[0] or bottom_right[1] < coords[1]:
|
||||
raise ValueError(
|
||||
f"Invalid coordinates or size, coords: {coords}, size: {size}"
|
||||
)
|
||||
|
||||
# Check if the area only contains one color
|
||||
cropped_image = image.crop((*coords, *bottom_right))
|
||||
if len(set(list(cropped_image.getdata()))) == 1:
|
||||
continue
|
||||
|
||||
# Draw rectangle on image
|
||||
draw.rectangle([coords, bottom_right], outline="red", width=1)
|
||||
|
||||
# Draw index number at the bottom left of the bounding box with black background
|
||||
text_position = (
|
||||
coords[0],
|
||||
bottom_right[1],
|
||||
) # Adjust Y to be above the bottom right
|
||||
text_bbox: Tuple[int, int, int, int] = draw.textbbox(
|
||||
text_position, str(index), font=font, anchor="lb"
|
||||
)
|
||||
# offset: int = bottom_right[1]-text_bbox[3]
|
||||
# text_bbox = (text_bbox[0], text_bbox[1]+offset, text_bbox[2], text_bbox[3]+offset)
|
||||
|
||||
# draw.rectangle([text_position, (text_position[0] + 25, text_position[1] + 18)], fill='black')
|
||||
draw.rectangle(text_bbox, fill="black")
|
||||
draw.text(
|
||||
text_position, str(index), font=font, anchor="lb", fill="white"
|
||||
)
|
||||
|
||||
# each mark is an x, y, w, h tuple
|
||||
marks.append(
|
||||
[
|
||||
original_coords[0],
|
||||
original_coords[1],
|
||||
original_size[0],
|
||||
original_size[1],
|
||||
]
|
||||
)
|
||||
drew_nodes.append(_node)
|
||||
|
||||
if _node.text:
|
||||
node_text = (
|
||||
_node.text
|
||||
if '"' not in _node.text
|
||||
else '"{:}"'.format(_node.text.replace('"', '""'))
|
||||
)
|
||||
elif _node.get(
|
||||
"{uri:deskat:uia.windows.microsoft.org}class", ""
|
||||
).endswith("EditWrapper") and _node.get(
|
||||
"{uri:deskat:value.at-spi.gnome.org}value"
|
||||
):
|
||||
node_text: str = _node.get(
|
||||
"{uri:deskat:value.at-spi.gnome.org}value"
|
||||
)
|
||||
node_text = (
|
||||
node_text
|
||||
if '"' not in node_text
|
||||
else '"{:}"'.format(node_text.replace('"', '""'))
|
||||
)
|
||||
else:
|
||||
node_text = '""'
|
||||
text_information: str = "{:d}\t{:}\t{:}\t{:}".format(
|
||||
index, _node.tag, _node.get("name", ""), node_text
|
||||
)
|
||||
text_informations.append(text_information)
|
||||
|
||||
index += 1
|
||||
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
output_image_stream = io.BytesIO()
|
||||
image.save(output_image_stream, format="PNG")
|
||||
image_content = output_image_stream.getvalue()
|
||||
|
||||
return marks, drew_nodes, "\n".join(text_informations), image_content
|
||||
|
||||
|
||||
def print_nodes_with_indent(nodes, indent=0):
|
||||
for node in nodes:
|
||||
print(" " * indent, node.tag, node.attrib)
|
||||
print_nodes_with_indent(node, indent + 2)
|
||||
|
||||
|
||||
# Code based on https://github.com/xlang-ai/OSWorld/blob/main/mm_agents/agent.py
|
||||
|
||||
|
||||
def encode_image(image_content):
|
||||
return base64.b64encode(image_content).decode("utf-8")
|
||||
|
||||
|
||||
def encoded_img_to_pil_img(data_str):
|
||||
base64_str = data_str.replace("data:image/png;base64,", "")
|
||||
image_data = base64.b64decode(base64_str)
|
||||
image = Image.open(BytesIO(image_data))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def save_to_tmp_img_file(data_str):
|
||||
base64_str = data_str.replace("data:image/png;base64,", "")
|
||||
image_data = base64.b64decode(base64_str)
|
||||
image = Image.open(BytesIO(image_data))
|
||||
|
||||
tmp_img_path = os.path.join(tempfile.mkdtemp(), "tmp_img.png")
|
||||
image.save(tmp_img_path)
|
||||
|
||||
return tmp_img_path
|
||||
|
||||
|
||||
def linearize_accessibility_tree(accessibility_tree, platform="ubuntu", tag=False):
|
||||
# leaf_nodes = find_leaf_nodes(accessibility_tree)
|
||||
filtered_nodes = filter_nodes(ET.fromstring(accessibility_tree), platform)
|
||||
linearized_accessibility_tree = [
|
||||
"tag\tname\ttext\tposition (top-left x&y)\tsize (w&h)"
|
||||
]
|
||||
# Linearize the accessibility tree nodes into a table format
|
||||
|
||||
for node in filtered_nodes:
|
||||
# linearized_accessibility_tree += node.tag + "\t"
|
||||
# linearized_accessibility_tree += node.attrib.get('name') + "\t"
|
||||
if node.text:
|
||||
text = (
|
||||
node.text
|
||||
if '"' not in node.text
|
||||
else '"{:}"'.format(node.text.replace('"', '""'))
|
||||
)
|
||||
elif node.get("{uri:deskat:uia.windows.microsoft.org}class", "").endswith(
|
||||
"EditWrapper"
|
||||
) and node.get("{uri:deskat:value.at-spi.gnome.org}value"):
|
||||
text: str = node.get("{uri:deskat:value.at-spi.gnome.org}value")
|
||||
text = text if '"' not in text else '"{:}"'.format(text.replace('"', '""'))
|
||||
else:
|
||||
text = '""'
|
||||
# linearized_accessibility_tree += node.attrib.get(
|
||||
# , "") + "\t"
|
||||
# linearized_accessibility_tree += node.attrib.get('{uri:deskat:component.at-spi.gnome.org}size', "") + "\n"
|
||||
linearized_accessibility_tree.append(
|
||||
"{:}\t{:}\t{:}\t{:}\t{:}".format(
|
||||
node.tag,
|
||||
node.get("name", ""),
|
||||
text,
|
||||
node.get("{uri:deskat:component.at-spi.gnome.org}screencoord", ""),
|
||||
node.get("{uri:deskat:component.at-spi.gnome.org}size", ""),
|
||||
)
|
||||
)
|
||||
if tag:
|
||||
linearized_accessibility_tree = tag_accessibility_tree(
|
||||
linearized_accessibility_tree
|
||||
)
|
||||
|
||||
return "\n".join(linearized_accessibility_tree)
|
||||
|
||||
|
||||
def tag_accessibility_tree(linear_accessibility_tree):
|
||||
# Add 'id' to the first line
|
||||
linear_accessibility_tree[0] = "id\t" + linear_accessibility_tree[0]
|
||||
|
||||
# Start idx from 1 to correctly index into the list
|
||||
for idx in range(1, len(linear_accessibility_tree)):
|
||||
line = linear_accessibility_tree[idx]
|
||||
linear_accessibility_tree[idx] = f"[{str(idx)}]\t" + line
|
||||
|
||||
return linear_accessibility_tree
|
||||
|
||||
|
||||
def tag_screenshot(screenshot, accessibility_tree, platform="ubuntu"):
|
||||
nodes = filter_nodes(
|
||||
ET.fromstring(accessibility_tree), platform=platform, check_image=True
|
||||
)
|
||||
# Make tag screenshot
|
||||
marks, drew_nodes, element_list, tagged_screenshot = draw_bounding_boxes(
|
||||
nodes, screenshot
|
||||
)
|
||||
|
||||
return marks, drew_nodes, tagged_screenshot, element_list
|
||||
|
||||
|
||||
def parse_dag(text):
|
||||
pattern = r"<json>(.*?)</json>"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
json_str = match.group(1)
|
||||
try:
|
||||
json_data = json.loads(json_str)
|
||||
return Dag(**json_data["dag"])
|
||||
except json.JSONDecodeError:
|
||||
print("Error: Invalid JSON")
|
||||
return None
|
||||
except KeyError:
|
||||
print("Error: 'dag' key not found in JSON")
|
||||
return None
|
||||
except ValidationError as e:
|
||||
print(f"Error: Invalid data structure - {e}")
|
||||
return None
|
||||
else:
|
||||
print("Error: JSON not found")
|
||||
return None
|
||||
|
||||
|
||||
def parse_subinfo(subinfo_string):
|
||||
matches = re.findall(r"```json\s+(.*?)\s+```", subinfo_string, re.DOTALL)
|
||||
if matches:
|
||||
# Assuming there's only one match, parse the JSON string into a dictionary
|
||||
try:
|
||||
subinfo_dict = json.loads(matches[0])
|
||||
return subinfo_dict
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Failed to parse JSON: {e}")
|
||||
return {"error": e}
|
||||
else:
|
||||
return {
|
||||
"error": "Subinfo generated in incorrect format. Please use the correct format."
|
||||
}
|
||||
|
||||
|
||||
def parse_actions_from_string(input_string):
|
||||
if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
|
||||
return [input_string.strip()]
|
||||
# Search for a JSON string within the input string
|
||||
actions = []
|
||||
matches = re.findall(r"```json\s+(.*?)\s+```", input_string, re.DOTALL)
|
||||
if matches:
|
||||
# Assuming there's only one match, parse the JSON string into a dictionary
|
||||
try:
|
||||
for match in matches:
|
||||
action_dict = json.loads(match)
|
||||
actions.append(action_dict)
|
||||
return actions
|
||||
except json.JSONDecodeError as e:
|
||||
return f"Failed to parse JSON: {e}"
|
||||
else:
|
||||
matches = re.findall(r"```\s+(.*?)\s+```", input_string, re.DOTALL)
|
||||
if matches:
|
||||
# Assuming there's only one match, parse the JSON string into a dictionary
|
||||
try:
|
||||
for match in matches:
|
||||
action_dict = json.loads(match)
|
||||
actions.append(action_dict)
|
||||
return actions
|
||||
except json.JSONDecodeError as e:
|
||||
return f"Failed to parse JSON: {e}"
|
||||
else:
|
||||
try:
|
||||
action_dict = json.loads(input_string)
|
||||
return [action_dict]
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError("Invalid response format: " + input_string)
|
||||
|
||||
|
||||
def parse_fixed_action_from_string(input_string):
|
||||
pattern = r"```(?:\w+\s+)?(.*?)```"
|
||||
matches = re.findall(pattern, input_string)
|
||||
if matches:
|
||||
# Assuming there's only one match, parse the JSON string into a dictionary
|
||||
try:
|
||||
for match in matches:
|
||||
action = match
|
||||
return action
|
||||
except json.JSONDecodeError as e:
|
||||
return f"Failed to parse JSON: {e}"
|
||||
|
||||
return "agent.wait()"
|
||||
|
||||
|
||||
def parse_code_from_string(input_string):
|
||||
input_string = "\n".join(
|
||||
[line.strip() for line in input_string.split(";") if line.strip()]
|
||||
)
|
||||
if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
|
||||
return [input_string.strip()]
|
||||
|
||||
# This regular expression will match both ```code``` and ```python code```
|
||||
# and capture the `code` part. It uses a non-greedy match for the content inside.
|
||||
pattern = r"```(?:\w+\s+)?(.*?)```"
|
||||
# Find all non-overlapping matches in the string
|
||||
matches = re.findall(pattern, input_string, re.DOTALL)
|
||||
|
||||
# The regex above captures the content inside the triple backticks.
|
||||
# The `re.DOTALL` flag allows the dot `.` to match newline characters as well,
|
||||
# so the code inside backticks can span multiple lines.
|
||||
|
||||
# matches now contains all the captured code snippets
|
||||
|
||||
codes = []
|
||||
|
||||
for match in matches:
|
||||
match = match.strip()
|
||||
commands = [
|
||||
"WAIT",
|
||||
"DONE",
|
||||
"FAIL",
|
||||
] # fixme: updates this part when we have more commands
|
||||
|
||||
if match in commands:
|
||||
codes.append(match.strip())
|
||||
elif match.split("\n")[-1] in commands:
|
||||
if len(match.split("\n")) > 1:
|
||||
codes.append("\n".join(match.split("\n")[:-1]))
|
||||
codes.append(match.split("\n")[-1])
|
||||
else:
|
||||
codes.append(match)
|
||||
|
||||
return codes
|
||||
|
||||
|
||||
def parse_single_code_from_string(input_string):
|
||||
input_string = input_string.strip()
|
||||
if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
|
||||
return input_string.strip()
|
||||
|
||||
# This regular expression will match both ```code``` and ```python code```
|
||||
# and capture the `code` part. It uses a non-greedy match for the content inside.
|
||||
pattern = r"```(?:\w+\s+)?(.*?)```"
|
||||
# Find all non-overlapping matches in the string
|
||||
matches = re.findall(pattern, input_string, re.DOTALL)
|
||||
|
||||
# The regex above captures the content inside the triple backticks.
|
||||
# The `re.DOTALL` flag allows the dot `.` to match newline characters as well,
|
||||
# so the code inside backticks can span multiple lines.
|
||||
|
||||
# matches now contains all the captured code snippets
|
||||
|
||||
codes = []
|
||||
|
||||
for match in matches:
|
||||
match = match.strip()
|
||||
commands = [
|
||||
"WAIT",
|
||||
"DONE",
|
||||
"FAIL",
|
||||
] # fixme: updates this part when we have more commands
|
||||
|
||||
if match in commands:
|
||||
codes.append(match.strip())
|
||||
elif match.split("\n")[-1] in commands:
|
||||
if len(match.split("\n")) > 1:
|
||||
codes.append("\n".join(match.split("\n")[:-1]))
|
||||
codes.append(match.split("\n")[-1])
|
||||
else:
|
||||
codes.append(match)
|
||||
|
||||
return codes[0]
|
||||
|
||||
|
||||
def parse_action_from_fixed_code(action_string, linearized_accessibility_tree):
|
||||
|
||||
import re
|
||||
|
||||
def parse_action_from_agent_code(action_str):
|
||||
# First, extract the code block within triple backticks
|
||||
code_block_pattern = r"```(.*?)```"
|
||||
code_block_match = re.search(code_block_pattern, action_str, re.DOTALL)
|
||||
|
||||
if not code_block_match:
|
||||
raise ValueError("No code block found")
|
||||
|
||||
code_block = code_block_match.group(1).strip()
|
||||
|
||||
# Define a regex pattern to extract the action type and parameters
|
||||
action_pattern = r"agent\.(\w+)\((.*?)\)"
|
||||
match = re.match(action_pattern, code_block, re.IGNORECASE)
|
||||
|
||||
if match:
|
||||
action_type = match.group(1)
|
||||
params_str = match.group(2)
|
||||
|
||||
# Split the parameters by comma and strip any surrounding whitespace or quotes
|
||||
params = [
|
||||
param.strip().strip('"').strip("'") for param in params_str.split(",")
|
||||
]
|
||||
|
||||
# Convert numeric parameters to integers
|
||||
for i in range(len(params)):
|
||||
try:
|
||||
params[i] = int(params[i])
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
return action_type, params
|
||||
else:
|
||||
raise ValueError("Invalid action string format")
|
||||
|
||||
parsed_action = parse_action_from_agent_code(action_string)
|
||||
action_type, params = parsed_action
|
||||
code = ""
|
||||
|
||||
def get_position_from_tree(element_id):
|
||||
element = linearized_accessibility_tree[element_id]
|
||||
position_str, size_str = element.split("\t")[-2].replace("(", "").replace(
|
||||
")", ""
|
||||
), element.split("\t")[-1].replace("(", "").replace(")", "")
|
||||
top_x_str, top_y_str = position_str.split(",")
|
||||
top_x, top_y = int(top_x_str.strip()), int(top_y_str.strip())
|
||||
size_x_str, size_y_str = size_str.split(",")
|
||||
size_x, size_y = int(size_x_str.strip()), int(size_y_str.strip())
|
||||
centroid_x, centroid_y = top_x + size_x // 2, top_y + size_y // 2
|
||||
return centroid_x, centroid_y
|
||||
|
||||
if action_type == "left_click_element_by_id":
|
||||
element_id = int(params[0])
|
||||
centroid_x, centroid_y = get_position_from_tree(element_id)
|
||||
code = f"""position = ({centroid_x}, {centroid_y}); pyautogui.click(position)
|
||||
"""
|
||||
|
||||
elif action_type == "right_click_element_by_id":
|
||||
element_id = int(params[0])
|
||||
centroid_x, centroid_y = get_position_from_tree(element_id)
|
||||
code = f"""
|
||||
position = ({centroid_x}, {centroid_y}); pyautogui.click(position, button='right')
|
||||
"""
|
||||
|
||||
elif action_type == "hover_over_element_by_id":
|
||||
element_id = int(params[0])
|
||||
centroid_x, centroid_y = get_position_from_tree(element_id)
|
||||
code = (
|
||||
f"""position = ({centroid_x}, {centroid_y}); pyautogui.moveTo(position)"""
|
||||
)
|
||||
|
||||
elif action_type == "type_write_element_by_id":
|
||||
element_id = int(params[0])
|
||||
text = params[1]
|
||||
centroid_x, centroid_y = get_position_from_tree(element_id)
|
||||
code = f"""
|
||||
position = ({centroid_x}, {centroid_y}); pyautogui.click(position); time.sleep(0.75); pyautogui.typewrite("{text}")"""
|
||||
|
||||
elif action_type == "press_key_combinations":
|
||||
keys = params
|
||||
keys_str = '", "'.join(keys)
|
||||
code = f"""
|
||||
pyautogui.hotkey("{keys_str}")
|
||||
"""
|
||||
|
||||
elif action_type == "wait":
|
||||
code = """WAIT"""
|
||||
|
||||
elif action_type == "done":
|
||||
code = """DONE"""
|
||||
|
||||
elif action_type == "fail":
|
||||
code = "FAIL"
|
||||
|
||||
return [code.strip()]
|
||||
|
||||
|
||||
def parse_code_from_som_string(input_string, masks):
|
||||
# parse the output string by masks
|
||||
tag_vars = ""
|
||||
for i, mask in enumerate(masks):
|
||||
x, y, w, h = mask
|
||||
tag_vars += (
|
||||
"tag_"
|
||||
+ str(i + 1)
|
||||
+ "="
|
||||
+ "({}, {})".format(int(x + w // 2), int(y + h // 2))
|
||||
)
|
||||
tag_vars += "\n"
|
||||
|
||||
actions = parse_code_from_string(input_string)
|
||||
|
||||
for i, action in enumerate(actions):
|
||||
if action.strip() in ["WAIT", "DONE", "FAIL"]:
|
||||
pass
|
||||
else:
|
||||
action = tag_vars + action
|
||||
actions[i] = action
|
||||
|
||||
return actions
|
||||
|
||||
|
||||
def box_iou(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Fast vectorized IOU implementation using only NumPy
|
||||
boxes1: [N, 4] array of boxes
|
||||
boxes2: [M, 4] array of boxes
|
||||
Returns: [N, M] array of IOU values
|
||||
"""
|
||||
# Calculate areas of boxes1
|
||||
area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
|
||||
|
||||
# Calculate areas of boxes2
|
||||
area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
|
||||
|
||||
# Get intersections using broadcasting
|
||||
lt = np.maximum(boxes1[:, None, :2], boxes2[None, :, :2]) # [N,M,2]
|
||||
rb = np.minimum(boxes1[:, None, 2:], boxes2[None, :, 2:]) # [N,M,2]
|
||||
|
||||
# Calculate intersection areas
|
||||
wh = np.clip(rb - lt, 0, None) # [N,M,2]
|
||||
intersection = wh[:, :, 0] * wh[:, :, 1] # [N,M]
|
||||
|
||||
# Calculate union areas
|
||||
union = area1[:, None] + area2[None, :] - intersection
|
||||
|
||||
# Calculate IOU
|
||||
iou = np.where(union > 0, intersection / union, 0)
|
||||
return iou
|
||||
|
||||
|
||||
def calculate_iou(rect1, rect2):
|
||||
"""
|
||||
Calculate the Intersection over Union (IoU) of two rectangles using numpy.
|
||||
|
||||
Parameters:
|
||||
rect1, rect2: Tuples containing the coordinates of the rectangles in the form (x_min, y_min, x_max, y_max)
|
||||
|
||||
Returns:
|
||||
IoU: Intersection over Union value
|
||||
"""
|
||||
# Convert the coordinates to tensors
|
||||
box1 = np.array([rect1], dtype=np.float32)
|
||||
box2 = np.array([rect2], dtype=np.float32)
|
||||
|
||||
# Calculate IoU using numpy
|
||||
iou = box_iou(box1, box2)
|
||||
|
||||
return iou
|
||||
|
||||
|
||||
def text_cvt_orc_format_paddle(paddle_result):
|
||||
texts = []
|
||||
print("paddle_result: ", paddle_result)
|
||||
for i, line in enumerate(paddle_result[0]):
|
||||
points = np.array(line[0])
|
||||
print("points: ", points)
|
||||
location = {
|
||||
"left": int(min(points[:, 0])),
|
||||
"top": int(min(points[:, 1])),
|
||||
"right": int(max(points[:, 0])),
|
||||
"bottom": int(max(points[:, 1])),
|
||||
}
|
||||
print("location: ", location)
|
||||
content = line[1][0]
|
||||
texts.append((i, content, location))
|
||||
return texts
|
||||
|
||||
|
||||
def trim_accessibility_tree(linearized_accessibility_tree, max_tokens):
|
||||
enc = tiktoken.encoding_for_model("gpt-4")
|
||||
tokens = enc.encode(linearized_accessibility_tree)
|
||||
if len(tokens) > max_tokens:
|
||||
print("MAX TOKEN LENGTH OF ACCESSIBILITY TREE EXCEEDED")
|
||||
linearized_accessibility_tree = enc.decode(tokens[:max_tokens])
|
||||
linearized_accessibility_tree += "[...]\n"
|
||||
return linearized_accessibility_tree
|
||||
|
||||
|
||||
def get_input_token_length(input_string):
|
||||
enc = tiktoken.encoding_for_model("gpt-4")
|
||||
tokens = enc.encode(input_string)
|
||||
return len(tokens)
|
||||
|
||||
|
||||
def load_osworld_example(base_path: str, domain: str, id: int):
|
||||
example_path = f"{base_path}/{domain}"
|
||||
example_path = (
|
||||
f"/Users/saaketagashe/Documents/OSWorld/evaluation_examples/examples/{domain}"
|
||||
)
|
||||
examples = os.listdir(example_path)
|
||||
|
||||
with open(example_path + "/" + examples[id], "r") as f:
|
||||
example = json.load(f)
|
||||
|
||||
return example
|
||||
|
||||
|
||||
def sanitize_code(code):
|
||||
# This pattern captures the outermost double-quoted text
|
||||
if "\n" in code:
|
||||
pattern = r'(".*?")'
|
||||
# Find all matches in the text
|
||||
matches = re.findall(pattern, code, flags=re.DOTALL)
|
||||
if matches:
|
||||
# Replace the first occurrence only
|
||||
first_match = matches[0]
|
||||
code = code.replace(first_match, f'"""{first_match[1:-1]}"""', 1)
|
||||
return code
|
||||
|
||||
|
||||
def extract_first_agent_function(code_string):
|
||||
# Regular expression pattern to match 'agent' functions with any arguments, including nested parentheses
|
||||
pattern = r'agent\.[a-zA-Z_]+\((?:[^()\'"]|\'[^\']*\'|"[^"]*")*\)'
|
||||
|
||||
# Find all matches in the string
|
||||
matches = re.findall(pattern, code_string)
|
||||
|
||||
# Return the first match if found, otherwise return None
|
||||
return matches[0] if matches else None
|
||||
|
||||
|
||||
def load_knowledge_base(kb_path: str) -> Dict:
|
||||
try:
|
||||
with open(kb_path, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading knowledge base: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def load_embeddings(embeddings_path: str) -> Dict:
|
||||
try:
|
||||
with open(embeddings_path, "rb") as f:
|
||||
return pickle.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading embeddings: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def save_embeddings(embeddings_path: str, embeddings: Dict):
|
||||
try:
|
||||
with open(embeddings_path, "wb") as f:
|
||||
pickle.dump(embeddings, f)
|
||||
except Exception as e:
|
||||
print(f"Error saving embeddings: {e}")
|
||||
@@ -0,0 +1,58 @@
|
||||
import base64
|
||||
import gc
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
from fastapi import FastAPI
|
||||
from paddleocr import PaddleOCR
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI()
|
||||
ocr_module = PaddleOCR(use_angle_cls=True, lang="en")
|
||||
|
||||
|
||||
class ImageData(BaseModel):
|
||||
img_bytes: bytes
|
||||
|
||||
|
||||
def text_cvt_orc_format_paddle(paddle_result):
|
||||
texts = []
|
||||
print("paddle_result: ", paddle_result)
|
||||
for i, line in enumerate(paddle_result[0]):
|
||||
points = np.array(line[0])
|
||||
print("points: ", points)
|
||||
location = {
|
||||
"left": int(min(points[:, 0])),
|
||||
"top": int(min(points[:, 1])),
|
||||
"right": int(max(points[:, 0])),
|
||||
"bottom": int(max(points[:, 1])),
|
||||
}
|
||||
print("location: ", location)
|
||||
content = line[1][0]
|
||||
texts.append((i, content, location))
|
||||
return texts
|
||||
|
||||
|
||||
def ocr_results(screenshot):
|
||||
screenshot_img = Image.open(io.BytesIO(screenshot))
|
||||
result = ocr_module.ocr(np.array(screenshot_img), cls=True)
|
||||
return text_cvt_orc_format_paddle(result)
|
||||
|
||||
|
||||
@app.post("/ocr/")
|
||||
async def read_image(image_data: ImageData):
|
||||
image_bytes = base64.b64decode(image_data.img_bytes)
|
||||
results = ocr_results(image_bytes)
|
||||
|
||||
# Explicitly delete unused variables and run garbage collector
|
||||
del image_bytes
|
||||
gc.collect()
|
||||
|
||||
return {"results": results}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="127.0.0.1", port=8000)
|
||||
@@ -0,0 +1,33 @@
|
||||
import requests
|
||||
import toml
|
||||
import os
|
||||
|
||||
|
||||
def query_to_perplexica(query):
|
||||
# Retrieve the URL from an environment variable
|
||||
url = os.getenv("PERPLEXICA_URL")
|
||||
if not url:
|
||||
raise ValueError(
|
||||
"PERPLEXICA_URL environment variable not set. It may take the form: 'http://localhost:{port}/api/search'. The port number is set in the config.toml in the Perplexica directory."
|
||||
)
|
||||
|
||||
# Request Message
|
||||
message = {"focusMode": "webSearch", "query": query, "history": [["human", query]]}
|
||||
|
||||
response = requests.post(url, json=message)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()["message"]
|
||||
elif response.status_code == 400:
|
||||
raise ValueError(
|
||||
"The request is malformed or missing required fields, such as FocusModel or query"
|
||||
)
|
||||
else:
|
||||
raise ValueError("Internal Server Error")
|
||||
|
||||
|
||||
# Test Code
|
||||
if __name__ == "__main__":
|
||||
query = "What is Agent S?"
|
||||
response = query_to_perplexica(query)
|
||||
print(response)
|
||||
Reference in New Issue
Block a user