-
Notifications
You must be signed in to change notification settings - Fork 4
/
search_main.py
50 lines (46 loc) · 2.61 KB
/
search_main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
from litewebagent.utils.playwright_manager import PlaywrightManager
from dotenv import load_dotenv
import argparse
import json
import playwright
_ = load_dotenv()
import time
import logging
from litewebagent.utils.utils import setup_logger
from litewebagent.agents.webagent import setup_search_agent
import os
def main(args):
log_folder = args.log_folder
logger = setup_logger(log_folder)
playwright_manager = PlaywrightManager(storage_state='state.json', video_dir=os.path.join(args.log_folder, 'videos'))
browser = playwright_manager.get_browser()
context = playwright_manager.get_context()
page = playwright_manager.get_page()
playwright_manager.playwright.selectors.set_test_id_attribute('data-unique-test-id')
# Use the features from command-line arguments
features = args.features.split(',') if args.features else None
branching_factor = args.branching_factor if args.branching_factor else None
agent = setup_search_agent(args.starting_url, args.goal, model_name=args.model, agent_type=args.agent_type, features=features, branching_factor=branching_factor, playwright_manager=playwright_manager, log_folder=args.log_folder)
trajectories = agent.send_prompt(args.plan, args.search_algorithm)
import pdb; pdb.set_trace()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run web automation tasks with different agent types.")
parser.add_argument('--agent_type', type=str, default="PromptSearchAgent",
choices=["PromptSearchAgent"],
help="Type of agent to use (default: PromptSearchAgent)")
parser.add_argument('--search_algorithm', type=str, default="bfs",
choices=["bfs", "dfs"])
parser.add_argument('--model', type=str, default="gpt-4o-mini",
help="Model to use for the agent (default: gpt-4o-mini)")
parser.add_argument('--starting_url', type=str, required=True,
help="Starting URL for the web automation task")
parser.add_argument('--plan', type=str, required=True,
help="Plan for the web automation task")
parser.add_argument('--goal', type=str, required=True,
help="Goal for the web automation task")
parser.add_argument('--features', type=str, default="axtree",
help="Comma-separated list of features to use (default: None, which uses all features)")
parser.add_argument('--branching_factor', type=int, default=None)
parser.add_argument('--log_folder', type=str, default='log', help='Path to the log folder')
args = parser.parse_args()
main(args)