main.py

"""Script to run end-to-end evaluation on the benchmark"""
import argparse
import glob
import json
import logging
import os
import random
import copy
import subprocess
import tempfile
import time
import pickle
from pathlib import Path
import time
import random
import openai
from gen_guidelines import gen_guideline
from collections import defaultdict
from pdb import set_trace as st

from agent import (
    Agent,
    PromptAgent,
    TeacherForcingAgent,
    construct_agent,
)
from agent.prompts import *
from browser_env import (
    Action,
    ActionTypes,
    ScriptBrowserEnv,
    StateInfo,
    Trajectory,
    create_stop_action,
)
from browser_env.actions import is_equivalent, create_id_based_action
from browser_env.auto_login import get_site_comb_from_filepath
from browser_env.helper_functions import (
    RenderHelper,
    get_action_description,
)
from evaluation_harness import evaluator_router
from sum_test import get_status

LOG_FOLDER = "log_files"
Path(LOG_FOLDER).mkdir(parents=True, exist_ok=True)
LOG_FILE_NAME = f"{LOG_FOLDER}/log_{time.strftime('%Y%m%d%H%M%S', time.localtime())}_{random.randint(0, 10000)}.log"

logger = logging.getLogger("logger")
logger.setLevel(logging.INFO)

console_handler = logging.StreamHandler()
console_handler.setLevel(logging.DEBUG)
logger.addHandler(console_handler)

file_handler = logging.FileHandler(LOG_FILE_NAME)
file_handler.setLevel(logging.DEBUG)
logger.addHandler(file_handler)

# Set the log format
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)


def config() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run end-to-end evaluation on the benchmark"
    )
    parser.add_argument(
        "--render", action="store_true", help="Render the browser"
    )
    parser.add_argument(
        "--slow_mo",
        type=int,
        default=0,
        help="Slow down the browser by the specified amount",
    )
    parser.add_argument(
        "--action_set_tag", default="id_accessibility_tree", help="Action type"
    )
    parser.add_argument(
        "--observation_type",
        choices=["accessibility_tree", "html", "image"],
        default="accessibility_tree",
        help="Observation type",
    )
    parser.add_argument(
        "--current_viewport_only",
        action="store_true",
        help="Only use the current viewport for the observation",
    )
    parser.add_argument("--viewport_width", type=int, default=1280)
    parser.add_argument("--viewport_height", type=int, default=720)
    parser.add_argument("--save_trace_enabled", action="store_true")
    parser.add_argument("--sleep_after_execution", type=float, default=0.0)

    parser.add_argument("--max_steps", type=int, default=20)

    # agent config
    parser.add_argument("--agent_type", type=str, default="prompt")
    parser.add_argument(
        "--instruction_path",
        type=str,
        default="agents/prompts/state_action_agent.json",
    )
    parser.add_argument(
        "--parsing_failure_th",
        help="When concesecutive parsing failure exceeds this threshold, the agent will stop",
        type=int,
        default=3,
    )
    parser.add_argument(
        "--repeating_action_failure_th",
        help="When concesecutive repeating action exceeds this threshold, the agent will stop",
        type=int,
        default=5
    )

    # lm config
    parser.add_argument("--provider", type=str, default="openai")
    parser.add_argument("--model", type=str, default="gpt-3.5-turbo-0613")
    parser.add_argument("--mode", type=str, default="chat")
    parser.add_argument("--temperature", type=float, default=0)
    parser.add_argument("--top_p", type=float, default=0.9)
    parser.add_argument("--context_length", type=int, default=0)
    parser.add_argument("--max_tokens", type=int, default=384)
    parser.add_argument("--stop_token", type=str, default=None)
    parser.add_argument(
        "--max_retry",
        type=int,
        help="max retry times to perform generations when parsing fails",
        default=1,
    )
    parser.add_argument(
        "--max_obs_length",
        type=int,
        help="when not zero, will truncate the observation to this length before feeding to the model",
        default=1920,
    )
    parser.add_argument(
        "--model_endpoint",
        help="huggingface model endpoint",
        type=str,
        default="",
    )

    # example config
    parser.add_argument("--test_start_idx", type=int, default=0)
    parser.add_argument("--test_end_idx", type=int, default=1000)

    # logging related
    parser.add_argument("--result_dir", type=str, default="")

    parser.add_argument("--apply_guidelines", action="store_true")
    parser.add_argument("--extract_guidelines", action="store_true")
    args = parser.parse_args()

    # check the whether the action space is compatible with the observation space
    if (
        args.action_set_tag == "id_accessibility_tree"
        and args.observation_type != "accessibility_tree"
    ):
        raise ValueError(
            f"Action type {args.action_set_tag} is incompatible with the observation type {args.observation_type}"
        )

    return args


def early_stop(
    trajectory: Trajectory, max_steps: int, thresholds: dict[str, int]
) -> tuple[bool, str]:
    """Check whether need to early stop"""

    # reach the max step
    num_steps = (len(trajectory) - 1) / 2
    if num_steps >= max_steps:
        return True, f"Reach max steps {max_steps}"
    
    last_k_actions: list[Action]
    action_seq: list[Action]

    # Case: parsing failure for k times
    k = thresholds["parsing_failure"]
    last_k_actions = trajectory[1::2][-k:]  # type: ignore[assignment]

    if last_k_actions and isinstance(last_k_actions[-1], str):
        print(last_k_actions[-1])
        return False, ''
    
    last_k_actions = [x for x in last_k_actions if not isinstance(x, str)]

    if len(last_k_actions) >= k:
        if all(
            [
                action["action_type"] == ActionTypes.NONE
                for action in last_k_actions
            ]
        ):
            return True, f"Failed to parse actions for {k} times"

    # Case: same action for k times
    k = thresholds["repeating_action"]
    last_k_actions = trajectory[1::2][-k:]  # type: ignore[assignment]
    action_seq = trajectory[1::2]  # type: ignore[assignment]
    last_k_actions = [x for x in last_k_actions if not isinstance(x, str)]

    if len(action_seq) == 0:
        return False, ""

    last_action: Action = action_seq[-1]

    if last_action["action_type"] != ActionTypes.TYPE:
        if len(last_k_actions) >= k:
            if all(
                [
                    is_equivalent(action, last_action)
                    for action in last_k_actions
                ]
            ):
                return True, f"Same action for {k} times"

    else:
        # check the action sequence
        if (
            sum([is_equivalent(action, last_action) for action in action_seq if not isinstance(action, str)])
            >= k
        ):
            return True, f"Same typing action for {k} times"

    return False, ""

def identical_action(act1, act2): #type content is ignored here
    if not (isinstance(act1, str) or isinstance(act2, str)):
        for key in ['action_type', 'element_role', 'element_name', 'page_number', 'element_id']:
            if act1[key] != act2[key]:
                print(key, act1[key], act2[key])
                return False
        return True 
    elif isinstance(act1, str) and isinstance(act2, str):
        return act1 == act2
    return False
    

def get_intent(config_file, colored_logger):
    with open(config_file) as f:
        _c = json.load(f)
        intent = _c["intent"]
        task_id = _c["task_id"]
        # automatically login
        if _c["storage_state"]:
            cookie_file_name = os.path.basename(_c["storage_state"])
            comb = get_site_comb_from_filepath(cookie_file_name)
            temp_dir = tempfile.mkdtemp()
            # subprocess to renew the cookie
            subprocess.run(
                [
                    "python",
                    "browser_env/auto_login.py",
                    "--auth_folder",
                    temp_dir,
                    "--site_list",
                    *comb,
                ]
            )
            _c["storage_state"] = f"{temp_dir}/{cookie_file_name}"
            assert os.path.exists(_c["storage_state"])
            # update the config file
            config_file = f"{temp_dir}/{os.path.basename(config_file)}"
            with open(config_file, "w") as f:
                json.dump(_c, f)
    colored_logger.log(f"[Config file]: {config_file}", option='print')
    colored_logger.log(f"[Intent]: {intent}", option='print')
    return config_file, intent, task_id

def set_env(args):
    env = ScriptBrowserEnv(
        headless=not args.render,
        slow_mo=args.slow_mo,
        observation_type=args.observation_type,
        current_viewport_only=args.current_viewport_only,
        viewport_size={
            "width": args.viewport_width,
            "height": args.viewport_height,
        },
        save_trace_enabled=args.save_trace_enabled,
        sleep_after_execution=args.sleep_after_execution,
    )
    return env

def url_to_global(url):
    url = url.replace(f'<your_webarena_host>:9999/f', "<your_webarena_host>:9999/r")
    url = url.replace('<your_webarena_host>:9999/', 'www.reddit.com/')
    url = url.replace('127.0.0.1:9999/', 'www.reddit.com/')
    return url 

def run_one_task(max_steps,
                early_stop_thresholds,
                render_helper,
                intent,
                guidelines,
                env, 
                config_file,
                colored_logger
                ):
    agent.reset(config_file)
    trajectory = []
    obs, info = env.reset(options={"config_file": config_file})
    state_info: StateInfo = {"observation": obs, "info": info}
    trajectory.append(state_info)
    meta_data = {"action_history": ["None"]}
    meta_data['guidelines'] = guidelines
    meta_data['summarization'] = []
    step_i = 0
    while True:
        early_stop_flag, stop_info = early_stop(
            trajectory, max_steps, early_stop_thresholds
        ) # 20, {'parsing_failure': 3, 'repeating_action': 3}
        colored_logger.log(f"\033[34mObs {step_i}\033[0m: \n{obs['text']}")
        colored_logger.log(f"\033[34mURL {step_i}\033[0m: \n{agent.prompt_constructor.map_url_to_real(info['page'].url)}")
        if guidelines:
            curr_sum = get_status(obs['text'], agent.prompt_constructor.map_url_to_real(info['page'].url), colored_logger, method='self')
            colored_logger.log(f"\033[35mSummarization {step_i}\033[0m: \n{curr_sum}")
            meta_data['summarization'].append(curr_sum)

        if early_stop_flag:
            action = create_stop_action(f"Early stop: {stop_info}")
        else:
            try:
                action = agent.next_action(
                    trajectory, intent, meta_data=meta_data
                )
            except ValueError as e:
                # get the error message
                action = create_stop_action(f"ERROR: {str(e)}")
                colored_logger.log(f"\033[32mCreating an error stop action for \033[0m \n{str(e)}")
        colored_logger.log(f"\033[35mAction {step_i}\033[0m: \n{action}")
        colored_logger.log(f"\033[36mIntent\033[0m: {intent}")
        trajectory.append(action)
        step_i += 1

        if not isinstance(action, str):
            action_str = get_action_description(
                action,
                state_info["info"]["observation_metadata"],
                action_set_tag=args.action_set_tag,
                prompt_constructor=agent.prompt_constructor
                if isinstance(agent, PromptAgent)
                else None,
            )
            render_helper.render(
                action, state_info, meta_data, args.render_screenshot
            )
        elif action.startswith('record'):
            action_str = f'Information record from previous action: {action}'
        
        meta_data["action_history"].append(action_str.replace('127.0.0.1:9999/', 'www.reddit.com/'))

        if (not isinstance(action, str)) and action["action_type"] == ActionTypes.STOP:
            break
        # obs, _, terminated, _, info = env.step(redirect_action(action))
        obs, _, terminated, _, info = env.step(action)
        if 'You cannot post more. Wait a while before trying again.' in obs['text']:
            colored_logger.log(f"\033[35mNeed to Reset environment due to environment issue{step_i}\033[0m")
            return 'error'
        state_info = {"observation": obs, "info": info}
        trajectory.append(state_info)

        if terminated:
            # add a action place holder
            trajectory.append(create_stop_action(""))
            break
    
    evaluator = evaluator_router(config_file)
    new_trajectory = []
    for i in range(len(trajectory)//2):
        obs1, act1 = trajectory[2*i], trajectory[2*i+1]
        if not isinstance(act1, str):
            new_trajectory.append(obs1)
            new_trajectory.append(act1)
    score = evaluator(
        trajectory=new_trajectory,
        config_file=config_file,
        page=env.page,
        client=env.get_page_client(env.page),
    )
    
    if score == 1:
        colored_logger.log(f"[Result] (PASS) {config_file}")
        print('please either reset or delete the new post then continue.')
        st()
    else:
        colored_logger.log(f"[Result] (FAIL) {config_file}")

    return score

def action_same_state(demo, 
                      start_step,
                      end_step, 
                      curr_sum,
                      colored_logger,
                      ):
    later_actions = '' #actions on the same page. 
    if end_step is None:
        for step in range(start_step+1, len(demo['trajectory'])//2):
            next_sum = get_status(demo['trajectory'][2*step]['observation']['text'], demo['trajectory'][2*step]['info']['page'].url, colored_logger, method='self')
            if next_sum == curr_sum:
                url = url_to_global(demo['trajectory'][2*step]['info']['page'].url)
                obs = demo['trajectory'][2*step]['observation']['text']
                act = demo['metadata']['action_history'][step+1]
                later_actions += f'Observation for Future step {step-start_step}:\n{obs}\nURL:{url}\nDemonstration action for Future step {step-start_step}: {act}\n'
            else:
                break
    else:
        all_prevs = 0
        for step in list(range(end_step))[::-1]:
            next_sum = get_status(demo['trajectory'][2*step]['observation']['text'], demo['trajectory'][2*step]['info']['page'].url, colored_logger, method='self')
            if next_sum == curr_sum:
                all_prevs += 1

        for step in list(range(end_step))[::-1]:
            next_sum = get_status(demo['trajectory'][2*step]['observation']['text'], demo['trajectory'][2*step]['info']['page'].url, colored_logger, method='self')
            if next_sum == curr_sum:
                url = url_to_global(demo['trajectory'][2*step]['info']['page'].url)
                obs = demo['trajectory'][2*step]['observation']['text']
                act = demo['metadata']['action_history'][step+1]
                later_actions = f'Observation for Previous step {all_prevs-(end_step-step)+1}:\n{obs}\nURL:{url}\nDemonstration action for Previous step {all_prevs-(end_step-step)+1}: {act}\n' + later_actions
            else:
                break
    if later_actions == '':
        later_actions = 'None'
    # colored_logger.log(f"\033[32mLater Action: \033[0m \n{later_actions}")
    return later_actions

def extract_guidelines(config_file_list,
                       agent,
                       record_guidelines,
                       colored_logger,
                       args,
                       ):
    file_path = './seen_states.json'
    with open(file_path, 'w') as file:
        json.dump({'On the main page of a reddit-like forum.': [0,0],
                   'On the page of a Reddit post.': [0,0],
                   'On the biography edit page of a Reddit user.': [0, 0],
                   'On the main page of a Reddit forum.': [0, 0],
                   "On the All forums page of a Reddit-like Platform, with element RootWebArea 'All forums' at the beginning. ": [0, 0]}, file)

    guidelines = defaultdict(lambda:defaultdict(lambda:[0, 0]))
    for config_file in config_file_list: 
        render_helper = RenderHelper(
            config_file, args.result_dir, args.action_set_tag
        )
        config_file, intent, task_id = get_intent(config_file, colored_logger)
        # demo_file = open(f"./extracted_data/reddit/{task_id}/human.pkl", 'rb')
        demo_file = open(f"./extracted_data/reddit/{task_id}/human_reason.pkl", 'rb')
        demo = pickle.load(demo_file)
        # keys: 'trajectory', 'metadata', 'intent', 'instruction_path', 'raw_actions', 'result'
        agent.reset(config_file)
        meta_data = {"action_history": ["None"]}
        meta_data['guidelines'] = guidelines
        meta_data['summarization'] = []
        for demo_step in range(len(demo['trajectory'])//2): #follow the steps in human demonstration
            state_info = demo['trajectory'][2*demo_step]
            obs = state_info['observation'] #in the demonstration 
            info = state_info['info']
            info['page'].url = url_to_global(info['page'].url)
            
            curr_sum = get_status(obs['text'], info['page'].url, colored_logger, 'self')
            meta_data['summarization'].append(curr_sum)

            colored_logger.log(f"\033[34mObs {demo_step}\033[0m: \n{obs['text']}")
            colored_logger.log(f"\033[34mURL {demo_step}\033[0m: \n{info['page'].url}")
            colored_logger.log(f"\033[35mSummarization {demo_step}\033[0m: \n{curr_sum}")
            action = agent.next_action(
                demo['trajectory'][:2*demo_step+1], intent, meta_data=meta_data
            )
            prompt = agent.prompt_constructor.construct(
                demo['trajectory'][:2*demo_step+1], intent, meta_data
            )
            demo_raw_action = demo['raw_actions'][demo_step]
            action_splitter = "```"
            pattern = rf"{action_splitter}((.|\n)*?){action_splitter}"
            match = re.findall(pattern, demo_raw_action, re.DOTALL)
            formatted_action = match[-1][0]
            demo_action = formatted_action if formatted_action.startswith('record') else create_id_based_action(formatted_action)

            if not identical_action(demo_action, action) or random.random() < 0.1: #or do we give guidelines regardless?
                colored_logger.log(f"\033[35mOriginal Action {demo_step}\033[0m: \n{action}")
                action = create_stop_action("")
                later_actions = action_same_state(demo, demo_step, None, curr_sum, colored_logger)
                prev_actions = action_same_state(demo, 0, demo_step, curr_sum, colored_logger)
                guidelines = gen_guideline(prompt[-1]['content'],
                                            curr_sum,
                                            action if isinstance(action, str) else action['raw_prediction'],
                                            demo_raw_action,
                                            guidelines,
                                            later_actions=later_actions,
                                            prev_actions=prev_actions,
                                            )
                # retry with the new guideline?
                meta_data['guidelines'] = guidelines
                action = agent.next_action(
                    demo['trajectory'][:2*demo_step+1], intent, meta_data=meta_data
                )
                corrected = identical_action(demo_action, action)
                colored_logger.log(f"\033[35mNew Action {demo_step}\033[0m: \n{action}, {demo_action}")
                colored_logger.log(f"\033[32mDesired Action\033[0m: {demo_raw_action}")
                colored_logger.log(f"\033[33mAfter adding guidelines\033[0m: {corrected}")
                # st()
            else:
                colored_logger.log(f"\033[31mGenerated the correct action {action['raw_prediction']}, moving on to the next step. \033[0m")
            action = demo_action
            meta_data["action_history"] = demo['metadata']['action_history'][:demo_step+2]

        colored_logger.log(f"\033[35mAction{demo_step}\033[0m: \n{action}")
        render_helper.close()
        colored_logger.log(record_guidelines)

        for key in guidelines:
            if key not in record_guidelines:
                record_guidelines[key] = copy.deepcopy(guidelines[key])
            else:
                for subkey in guidelines[key]:
                    record_guidelines[key][subkey] = guidelines[key][subkey]

    json_guidelines = {}
    for key in record_guidelines:
        json_guidelines[key] = {}
        for subkey in record_guidelines[key]:
            json_guidelines[key][subkey] = record_guidelines[key][subkey]
    
    with open(f"./saved_guidelines/guidelines.json", "w") as outfile:
        json.dump(json_guidelines, outfile)

from agent.colored_logger import ColoredLogger

def run(
    args: argparse.Namespace,
    agent: Agent | PromptAgent | TeacherForcingAgent,
    config_file_list: list[str],
) -> None:
    max_steps = args.max_steps

    early_stop_thresholds = {
        "parsing_failure": args.parsing_failure_th,
        "repeating_action": args.repeating_action_failure_th,
    }

    env=set_env(args)

    record_guidelines = defaultdict(lambda: defaultdict(lambda:[0, 0])) #weight, occurance
    colored_logger = ColoredLogger(args)
    agent.prompt_constructor.colored_logger = colored_logger
    if args.extract_guidelines:
        extract_guidelines(config_file_list,
                           agent,
                           record_guidelines,
                           colored_logger,
                           args,)

    if args.apply_guidelines:
        f = open(f"./saved_guidelines/guidelines.json") 
        guidelines = json.load(f)
    else:
        guidelines = None
    # task_range = [28, 29, 30, 31, 66, 67, 68, 399, 400, 402, 403, 405, 406, 407, 408, 409, 410, 580, 581, 583, 584, 595, 596, 597, 598, 600, 602, 603, 604, 606, 607, 608, 609, 610, 611, 613, 614, 615, 616, 617, 618, 620, 621, 622, 623, 624, 625, 627, 628, 629, 630, 632, 633, 634, 635, 636, 637, 638, 639, 640, 642, 643, 644, 646, 647, 648, 649, 650, 651, 715, 716, 717, 718, 719, 721, 722, 723, 724, 725, 726, 727, 728, 730, 731, 732, 734, 735]
    task_range = [399, 400]
    results, task_i = {}, 0
    while task_i < len(task_range):
        i = task_range[task_i]
        new_config_file = f"config_files/{i}.json"
        render_helper = RenderHelper(
            new_config_file, args.result_dir, args.action_set_tag
        )
        new_config_file, intent, task_id = get_intent(new_config_file, colored_logger)
        score = run_one_task(max_steps,
                            early_stop_thresholds,
                            render_helper,
                            intent,
                            guidelines,
                            env, 
                            new_config_file,
                            colored_logger,
                            )
        if score != 'error':
            task_i += 1
            results[i] = score
        else:
            env.close()
            env=set_env(args)
        render_helper.close()
        for key in results:
            print(key, results[key])
        print('results:', np.mean(list(results.values())))
    env.close()


def prepare(args: argparse.Namespace) -> None:
    # convert prompt python files to json
    from agent.prompts import to_json

    to_json.run()

    # prepare result dir
    result_dir = args.result_dir
    if not result_dir:
        result_dir = (
            f"cache/results_{time.strftime('%Y%m%d%H%M%S', time.localtime())}"
        )
    if not Path(result_dir).exists():
        Path(result_dir).mkdir(parents=True, exist_ok=True)
        args.result_dir = result_dir
        logger.info(f"Create result dir: {result_dir}")

    if not (Path(result_dir) / "traces").exists():
        (Path(result_dir) / "traces").mkdir(parents=True)

    # log the log file
    with open(os.path.join(result_dir, "log_files.txt"), "a+") as f:
        f.write(f"{LOG_FILE_NAME}\n")


def dump_config(args: argparse.Namespace) -> None:
    config_file = Path(args.result_dir) / "config.json"
    if not config_file.exists():
        with open(config_file, "w") as f:
            json.dump(vars(args), f, indent=4)
            logger.info(f"Dump config to {config_file}")


if __name__ == "__main__":
    args = config()
    args.sleep_after_execution = 3.0
    prepare(args)
    random.seed(0)

    if args.apply_guidelines:
        args.instruction_path = 'agent/prompts/jsons/reddit_guideline.json'
    else:
        args.instruction_path = 'agent/prompts/jsons/reddit.json'

    args.log_name = f"::{args.apply_guidelines}::from{args.extract_guidelines}::{args.model}::{args.instruction_path.split('/')[-1].split('.json')[0]}"

    test_file_list = []
    if args.extract_guidelines:
        task_range = [401, 626, 729, 582, 641, 720, 605, 733, 714, 652, 631, 599, 601, 69, 645, 404, 626, 27, 612, 619]
        task_range = [401]
    else:
        task_range = []

    for i in task_range:
        test_file_list.append(f"config_files/{i}.json")

    args.render = False
    args.render_screenshot = True
    args.save_trace_enabled = True

    args.current_viewport_only = True
    dump_config(args)

    agent = construct_agent(args)
    run(args, agent, test_file_list)

# python main.py --model gpt-4-1106-preview --result_dir ./exp_logs --apply_guidelines --extract_guidelines

# conda env create -f environment.yaml
# pip install -r requirements.txt
# playwright install
# pip install -e .