Skip to content

Commit 4b06787

Browse files
lynchee-owolynchee-owo
committed
feat: use message listener in python script to get data from chrome extension
1 parent a49acb7 commit 4b06787

16 files changed

+182
-181
lines changed

‎benchmark.py

Lines changed: 73 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -8,131 +8,126 @@
88
from selenium.webdriver.chrome.options import Options
99
from selenium.webdriver.chrome.service import Service
1010
from webdriver_manager.chrome import ChromeDriverManager
11-
import threading
12-
from flask import Flask, jsonify, request
13-
from flask_cors import CORS
1411

15-
# Load the .env file
1612
load_dotenv()
1713
api_key = os.getenv('OPENAI_API_KEY')
1814

1915
# Hard-coded coordinates to open web-wand side panel
2016
extensions_pos = (1060, 110)
2117
web_wand_pos = (900, 280)
2218

23-
app = Flask(__name__)
24-
CORS(app)
25-
26-
task_status = "idle"
27-
task_history = []
28-
task_id = ""
29-
19+
# Place to store task execution results
3020
results_dir = 'results'
3121
os.makedirs(results_dir, exist_ok=True)
3222

33-
@app.route('/status', methods=['POST'])
34-
def status():
35-
global task_status
36-
task_status = request.json.get('status', 'idle')
37-
return jsonify({"message": "Got status"}), 200
23+
def setup_driver():
24+
chrome_options = Options()
25+
# Load the unpacked webwand chrome extension
26+
chrome_options.add_argument("--load-extension=./dist")
27+
service = Service(ChromeDriverManager().install())
28+
driver = webdriver.Chrome(service=service, options=chrome_options)
29+
return driver
30+
31+
def dispatch_event(driver, event_name, event):
32+
script = f"""
33+
var event = new CustomEvent('{event_name}', {{ detail: {json.dumps(event)} }});
34+
document.dispatchEvent(event);
35+
"""
36+
driver.execute_script(script)
3837

39-
@app.route('/history', methods=['POST'])
40-
def history():
41-
global task_history
42-
global task_id
43-
task_history = request.json.get('history', [])
38+
def add_task_listener(driver, task_id):
39+
# Async script to add event listeners for taskStatus, taskHistory, and screenshot events
40+
script = """
41+
var callback = arguments[0];
42+
var keepListening = true;
43+
var historyData = null;
44+
45+
var statusListener = function (e) {
46+
if (e.detail.status !== 'running' && e.detail.status !== 'idle') {
47+
keepListening = false;
48+
document.removeEventListener('TaskStatusUpdate', statusListener);
49+
document.removeEventListener('ScreenshotUpdate', screenshotListener);
50+
document.removeEventListener('TaskHistoryUpdate', historyListener);
51+
callback({type: 'status', data: 'stopped', history: historyData});
52+
}
53+
};
54+
55+
var screenshotListener = function (e) {
56+
callback({type: 'screenshot', data: e.detail.imgData});
57+
};
58+
59+
var historyListener = function (e) {
60+
historyData = e.detail.history;
61+
};
62+
63+
document.addEventListener('TaskStatusUpdate', statusListener);
64+
document.addEventListener('ScreenshotUpdate', screenshotListener);
65+
document.addEventListener('TaskHistoryUpdate', historyListener);
66+
67+
// To keep the async script alive
68+
var checkInterval = setInterval(function () {
69+
if (!keepListening) {
70+
clearInterval(checkInterval);
71+
}
72+
}, 1000);
73+
"""
4474

75+
try:
76+
while True:
77+
event_data = driver.execute_async_script(script)
78+
if event_data:
79+
if event_data['type'] == 'screenshot':
80+
write_screenshots(task_id, event_data['data'])
81+
elif event_data['type'] == 'status' and event_data['data'] == 'stopped':
82+
if event_data.get('history'):
83+
write_history(task_id, event_data['history'])
84+
break
85+
except Exception as e:
86+
print(f"Error while listening for updates: {e}")
87+
88+
def write_history(task_id, task_history):
4589
task_dir = os.path.join(results_dir, f"test{task_id}")
4690
os.makedirs(task_dir, exist_ok=True)
4791
file_path = os.path.join(task_dir, 'interact_messages.json')
4892

4993
with open(file_path, 'w') as file:
5094
json.dump(task_history, file, indent=4)
51-
return jsonify({"message": "Got history"}), 200
5295

53-
@app.route('/screenshot/', methods=['POST'])
54-
def save_screenshot():
55-
global task_id
56-
# Decode the image data
57-
image_data = request.json.get('value', '')
96+
def write_screenshots(task_id, image_data):
5897
image_bytes = base64.b64decode(image_data)
5998

60-
# Ensure the directory exists
6199
task_dir = os.path.join(results_dir, f"test{task_id}")
62100
os.makedirs(task_dir, exist_ok=True)
63-
64-
# Create a unique filename for each screenshot
65101
timestamp = int(time.time())
66102
file_path = os.path.join(task_dir, f'screenshot_{timestamp}.png')
67103

68-
# Save the screenshot to the file
69104
with open(file_path, 'wb') as file:
70105
file.write(image_bytes)
71-
72-
return jsonify({"message": f"Screenshot saved to {file_path}"}), 200
73-
74-
def run_server():
75-
app.run(port=5000, debug=False, use_reloader=False)
76106

77-
def setup_driver():
78-
chrome_options = Options()
79-
# Load the unpacked webwand chrome extension
80-
chrome_options.add_argument("--load-extension=./dist")
81-
service = Service(ChromeDriverManager().install())
82-
driver = webdriver.Chrome(service=service, options=chrome_options)
83-
return driver
84-
85-
def dispatch_event(driver, event_name, event):
86-
script = f"""
87-
var event = new CustomEvent('{event_name}', {{ detail: {json.dumps(event)} }});
88-
document.dispatchEvent(event);
89-
"""
90-
driver.execute_script(script)
91-
92-
def check_task_status(driver):
93-
while task_status == "idle" or task_status == "running":
94-
dispatch_event(driver, 'GetTaskStatus', {})
95-
time.sleep(3)
96-
print(task_status)
97-
98-
def run_webwand_task(driver, task_description):
107+
def run_webwand_task(driver, task_id, task_description):
99108
dispatch_event(driver, 'SetAPIKey', {"value": api_key})
100109
dispatch_event(driver, 'SetTask', {"value": task_description})
101110
dispatch_event(driver, 'RunTask', {})
102-
103-
check_task_status(driver)
104-
dispatch_event(driver, 'GetTaskHistory', {})
105-
time.sleep(1)
111+
add_task_listener(driver, task_id)
106112

107113
def main():
108114
driver = setup_driver()
109-
110-
first_time = True
115+
initial_load = True
111116

112117
with open('tasks_test.jsonl', 'r') as file:
113-
global task_id
114-
global task_status
115-
global task_history
116118
for line in file:
117119
task = json.loads(line)
118120
task_id = task["id"]
119-
task_status = "idle"
120-
task_history = []
121121
driver.get(task['web'])
122122

123-
if first_time:
123+
if initial_load:
124124
# Simulate click to open side panel
125125
pyautogui.click(extensions_pos)
126126
pyautogui.click(web_wand_pos)
127-
first_time = False
127+
initial_load = False
128128

129-
run_webwand_task(driver, task['ques'])
129+
run_webwand_task(driver, task_id, task['ques'])
130130
driver.quit()
131131

132132
if __name__ == "__main__":
133-
# Start the Flask server in a new thread
134-
server_thread = threading.Thread(target=run_server)
135-
server_thread.start()
136-
137-
main()
138-
server_thread.join()
133+
main()

‎results/testGoogle--0/interact_messages.json

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,36 @@
11
[
22
{
3-
"prompt": "The user requests the following task:\n\n Search for bananas\n\n\nCurrent time: 5/8/2024, 2:18:03 PM\nCurrent URL: https://www.google.com/\nCurrent page scrolling position: 100.0%\n\n\nUse the following data as a reference of the annotated elements (using `===` as a delimiter between each annotation):\n\nlabel = 1\nname = About\ntagName = A\n===\nlabel = 2\nname = Store\ntagName = A\n===\nlabel = 3\nname = Gmail\ntagName = A\n===\nlabel = 4\nname = Search for Images\ntagName = A\n===\nlabel = 5\nname = Google apps\ntagName = A\nrole = button\n===\nlabel = 6\nname = Sign in\ntagName = A\n===\nlabel = 7\nname = Search\ntagName = TEXTAREA\ncurrentValue = \nplaceholder = \nrole = combobox\n===\nlabel = 8\nname = Search by voice\ntagName = DIV\nrole = button\n===\nlabel = 9\nname = Search by image\ntagName = DIV\nrole = button\n===\nlabel = 10\nname = Google Search\ntagName = INPUT\ncurrentValue = Google Search\nplaceholder = \nrole = button\n===\nlabel = 11\nname = I'm Feeling Lucky\ntagName = INPUT\ncurrentValue = I'm Feeling Lucky\nplaceholder = \nrole = button\n===\nlabel = 12\nname = Advertising\ntagName = A\n===\nlabel = 13\nname = Business\ntagName = A\n===\nlabel = 14\nname = How Search works\ntagName = A\n===\nlabel = 15\nname = Our third decade of climate action: join us\ntagName = A\n===\nlabel = 16\nname = Privacy\ntagName = A\n===\nlabel = 17\nname = Terms\ntagName = A\n===\nlabel = 18\nname = Settings\ntagName = DIV\nrole = button\n\nThis textarea currently has focus:\nlabel = 7\nname = Search\ntagName = TEXTAREA\ncurrentValue = \nplaceholder = \nrole = combobox\n",
4-
"response": "{\n \"thought\": \"I am entering 'bananas' into the search field and pressing ENTER to perform the search.\",\n \"action\": {\n \"name\": \"setValueAndEnter\",\n \"args\": {\n \"label\": \"7\",\n \"value\": \"bananas\"\n }\n }\n}",
53
"action": {
6-
"thought": "I am entering 'bananas' into the search field and pressing ENTER to perform the search.",
74
"operation": {
8-
"name": "setValueAndEnter",
95
"args": {
106
"label": "7",
117
"value": "bananas"
12-
}
13-
}
8+
},
9+
"name": "setValueAndEnter"
10+
},
11+
"thought": "I am setting the value 'bananas' in the search field and pressing ENTER to perform the search."
1412
},
13+
"prompt": "The user requests the following task:\n\n Search for bananas\n\n\nCurrent time: 5/10/2024, 4:25:57 PM\nCurrent URL: https://www.google.com/\nCurrent page scrolling position: 100.0%\n\n\nUse the following data as a reference of the annotated elements (using `===` as a delimiter between each annotation):\n\nlabel = 1\nname = About\ntagName = A\n===\nlabel = 2\nname = Store\ntagName = A\n===\nlabel = 3\nname = Gmail\ntagName = A\n===\nlabel = 4\nname = Search for Images\ntagName = A\n===\nlabel = 5\nname = Google apps\ntagName = A\nrole = button\n===\nlabel = 6\nname = Sign in\ntagName = A\n===\nlabel = 7\nname = Search\ntagName = TEXTAREA\ncurrentValue = \nplaceholder = \nrole = combobox\n===\nlabel = 8\nname = Search by voice\ntagName = DIV\nrole = button\n===\nlabel = 9\nname = Search by image\ntagName = DIV\nrole = button\n===\nlabel = 10\nname = Google Search\ntagName = INPUT\ncurrentValue = Google Search\nplaceholder = \nrole = button\n===\nlabel = 11\nname = I'm Feeling Lucky\ntagName = INPUT\ncurrentValue = I'm Feeling Lucky\nplaceholder = \nrole = button\n===\nlabel = 12\nname = Advertising\ntagName = A\n===\nlabel = 13\nname = Business\ntagName = A\n===\nlabel = 14\nname = How Search works\ntagName = A\n===\nlabel = 15\nname = Our third decade of climate action: join us\ntagName = A\n===\nlabel = 16\nname = Privacy\ntagName = A\n===\nlabel = 17\nname = Terms\ntagName = A\n===\nlabel = 18\nname = Settings\ntagName = DIV\nrole = button\n\nThis textarea currently has focus:\nlabel = 7\nname = Search\ntagName = TEXTAREA\ncurrentValue = \nplaceholder = \nrole = combobox\n",
14+
"response": "{\n \"thought\": \"I am setting the value 'bananas' in the search field and pressing ENTER to perform the search.\",\n \"action\": {\n \"name\": \"setValueAndEnter\",\n \"args\": {\n \"label\": \"7\",\n \"value\": \"bananas\"\n }\n }\n}",
1515
"usage": {
16+
"completion_tokens": 66,
1617
"prompt_tokens": 2501,
17-
"completion_tokens": 64,
18-
"total_tokens": 2565
18+
"total_tokens": 2567
1919
}
2020
},
2121
{
22-
"prompt": "The user requests the following task:\n\n Search for bananas\nYou have already taken the following actions: \nThought: I am entering 'bananas' into the search field and pressing ENTER to perform the search.\nAction:{\"name\":\"setValueAndEnter\",\"args\":{\"label\":\"7\",\"value\":\"bananas\"}}\n\n\n\nCurrent time: 5/8/2024, 2:18:26 PM\nCurrent URL: https://www.google.com/search\nCurrent page scrolling position: 27.9%\n\n\nUse the following data as a reference of the annotated elements (using `===` as a delimiter between each annotation):\n\nlabel = 1\nname = Google\ntagName = A\n===\nlabel = 2\nname = Search\ntagName = TEXTAREA\ncurrentValue = bananas\nplaceholder = \nrole = combobox\n===\nlabel = 3\nname = Clear\ntagName = DIV\nrole = button\n===\nlabel = 4\nname = Search by voice\ntagName = DIV\nrole = button\n===\nlabel = 5\nname = Search by image\ntagName = DIV\nrole = button\n===\nlabel = 6\nname = Search\ntagName = BUTTON\n===\nlabel = 7\nname = Images\ntagName = A\nrole = link\n===\nlabel = 8\nname = Shopping\ntagName = A\nrole = link\n===\nlabel = 9\nname = Videos\ntagName = A\nrole = link\n===\nlabel = 10\nname = Forums\ntagName = A\nrole = link\n===\nlabel = 11\nname = More\ntagName = DIV\nrole = button\n===\nlabel = 12\nname = Tools\ntagName = DIV\nrole = button\n===\nlabel = 13\nname = Add Nutrition facts\ntagName = A\nrole = link\n===\nlabel = 14\nname = Vibe\ntagName = DIV\nrole = button\n===\nlabel = 15\nname = Add Benefits\ntagName = A\nrole = link\n===\nlabel = 16\nname = Add Recipes\ntagName = A\nrole = link\n===\nlabel = 17\nname = Add Types\ntagName = A\nrole = link\n===\nlabel = 18\nname = Add Open now\ntagName = A\nrole = link\n===\nlabel = 19\nname = Add Takeout\ntagName = A\nrole = link\n===\nlabel = 20\nname = Add Tree\ntagName = A\nrole = link\n===\nlabel = 21\nname = Add Delivery\ntagName = A\nrole = link\n===\nlabel = 22\nname = Banana Wikipedia https://en.wikipedia.org \u203a wiki \u203a Banana\ntagName = A\n===\nlabel = 23\nname = About this result\ntagName = DIV\nrole = button\n===\nlabel = 24\nname = List of banana cultivars\ntagName = A\n===\nlabel = 25\nname = Cavendish banana\ntagName = A\n===\nlabel = 26\nname = Cooking banana\ntagName = A\n===\nlabel = 27\nname = Gros Michel\ntagName = A\n===\nlabel = 28\nname = About this result\ntagName = DIV\nrole = button\n===\nlabel = 29\nname = What are the benefits of eating bananas?\ntagName = DIV\nrole = button\n===\nlabel = 30\nname = Are bananas good to eat everyday?\ntagName = DIV\nrole = button\n===\nlabel = 31\nname = What does \ud83c\udf4c mean in chat from girl?\ntagName = DIV\nrole = button\n===\nlabel = 32\nname = What are bananas high in?\ntagName = DIV\nrole = button\n===\nlabel = 33\nname = Feedback\ntagName = DIV\nrole = button\n===\nlabel = 34\nname = About this result\ntagName = DIV\nrole = button\n",
23-
"response": "{\n \"thought\": \"The search for 'bananas' has been successfully performed. The results include a Wikipedia link about bananas, various questions related to bananas such as their benefits and nutritional content, and links to different types of bananas like Cavendish and Cooking bananas.\",\n \"action\": {\n \"name\": \"finish\"\n }\n}",
2422
"action": {
25-
"thought": "The search for 'bananas' has been successfully performed. The results include a Wikipedia link about bananas, various questions related to bananas such as their benefits and nutritional content, and links to different types of bananas like Cavendish and Cooking bananas.",
2623
"operation": {
2724
"name": "finish"
28-
}
25+
},
26+
"thought": "The search for 'bananas' has been successfully performed. The results include a Wikipedia link about bananas, various questions related to bananas such as their benefits and nutritional content, and links to different types of bananas like Cavendish and Gros Michel."
2927
},
28+
"prompt": "The user requests the following task:\n\n Search for bananas\nYou have already taken the following actions: \nThought: I am setting the value 'bananas' in the search field and pressing ENTER to perform the search.\nAction:{\"name\":\"setValueAndEnter\",\"args\":{\"label\":\"7\",\"value\":\"bananas\"}}\n\n\n\nCurrent time: 5/10/2024, 4:26:09 PM\nCurrent URL: https://www.google.com/search\nCurrent page scrolling position: 27.9%\n\n\nUse the following data as a reference of the annotated elements (using `===` as a delimiter between each annotation):\n\nlabel = 1\nname = Google\ntagName = A\n===\nlabel = 2\nname = Search\ntagName = TEXTAREA\ncurrentValue = bananas\nplaceholder = \nrole = combobox\n===\nlabel = 3\nname = Clear\ntagName = DIV\nrole = button\n===\nlabel = 4\nname = Search by voice\ntagName = DIV\nrole = button\n===\nlabel = 5\nname = Search by image\ntagName = DIV\nrole = button\n===\nlabel = 6\nname = Search\ntagName = BUTTON\n===\nlabel = 7\nname = Images\ntagName = A\nrole = link\n===\nlabel = 8\nname = Shopping\ntagName = A\nrole = link\n===\nlabel = 9\nname = Videos\ntagName = A\nrole = link\n===\nlabel = 10\nname = News\ntagName = A\nrole = link\n===\nlabel = 11\nname = More\ntagName = DIV\nrole = button\n===\nlabel = 12\nname = Tools\ntagName = DIV\nrole = button\n===\nlabel = 13\nname = Add Nutrition facts\ntagName = A\nrole = link\n===\nlabel = 14\nname = Add Benefits\ntagName = A\nrole = link\n===\nlabel = 15\nname = Vibe\ntagName = DIV\nrole = button\n===\nlabel = 16\nname = Add Recipes\ntagName = A\nrole = link\n===\nlabel = 17\nname = Add Takeout\ntagName = A\nrole = link\n===\nlabel = 18\nname = Add Types\ntagName = A\nrole = link\n===\nlabel = 19\nname = Add Delivery\ntagName = A\nrole = link\n===\nlabel = 20\nname = Add Fruit\ntagName = A\nrole = link\n===\nlabel = 21\nname = Add Open now\ntagName = A\nrole = link\n===\nlabel = 22\nname = Banana Wikipedia https://en.wikipedia.org \u203a wiki \u203a Banana\ntagName = A\n===\nlabel = 23\nname = About this result\ntagName = DIV\nrole = button\n===\nlabel = 24\nname = List of banana cultivars\ntagName = A\n===\nlabel = 25\nname = Cavendish banana\ntagName = A\n===\nlabel = 26\nname = Cooking banana\ntagName = A\n===\nlabel = 27\nname = Gros Michel\ntagName = A\n===\nlabel = 28\nname = About this result\ntagName = DIV\nrole = button\n===\nlabel = 29\nname = What are the benefits of eating bananas?\ntagName = DIV\nrole = button\n===\nlabel = 30\nname = Is it good to eat a banana every day?\ntagName = DIV\nrole = button\n===\nlabel = 31\nname = What are bananas high in?\ntagName = DIV\nrole = button\n===\nlabel = 32\nname = What does \ud83c\udf4c mean in chat from girl?\ntagName = DIV\nrole = button\n===\nlabel = 33\nname = Feedback\ntagName = DIV\nrole = button\n===\nlabel = 34\nname = About this result\ntagName = DIV\nrole = button\n",
29+
"response": "{\n \"thought\": \"The search for 'bananas' has been successfully performed. The results include a Wikipedia link about bananas, various questions related to bananas such as their benefits and nutritional content, and links to different types of bananas like Cavendish and Gros Michel.\",\n \"action\": {\n \"name\": \"finish\"\n }\n}",
3030
"usage": {
31-
"prompt_tokens": 2849,
3231
"completion_tokens": 70,
33-
"total_tokens": 2919
32+
"prompt_tokens": 2854,
33+
"total_tokens": 2924
3434
}
3535
}
3636
]

‎results/testGoogle--0/screenshot_1715192283.png

-267 KB
Binary file not shown.

‎results/testGoogle--0/screenshot_1715192306.png

-499 KB
Binary file not shown.

‎results/testGoogle--0/screenshot_1715372757.png

112 KB
Loading

‎results/testGoogle--0/screenshot_1715372769.png

203 KB
Loading

0 commit comments

Comments
 (0)