\n"
+
+ if render_screenshot:
+ # image observation
+ img_obs = observation["image"]
+ image = Image.fromarray(img_obs)
+ byte_io = io.BytesIO()
+ image.save(byte_io, format="PNG")
+ byte_io.seek(0)
+ image_bytes = base64.b64encode(byte_io.read())
+ image_str = image_bytes.decode("utf-8")
+ new_content += f"
\n"
+
+ # meta data
+ new_content += f"
{meta_data['action_history'][-1]}
\n"
+
+ # action
+ action_str = get_render_action(
+ action,
+ info["observation_metadata"],
+ action_set_tag=self.action_set_tag,
+ )
+ # with yellow background
+ action_str = f"
{action_str}
"
+ new_content += f"{action_str}\n"
+
+ # add new content
+ self.render_file.seek(0)
+ html = self.render_file.read()
+ html_body = re.findall(r"(.*?)", html, re.DOTALL)[0]
+ html_body += new_content
+
+ html = HTML_TEMPLATE.format(body=html_body)
+ self.render_file.seek(0)
+ self.render_file.truncate()
+ self.render_file.write(html)
+ self.render_file.flush()
+
+ def close(self) -> None:
+ self.render_file.close()
diff --git a/browser_env/trajectory.py b/browser_env/trajectory.py
new file mode 100644
index 0000000..1c4c410
--- /dev/null
+++ b/browser_env/trajectory.py
@@ -0,0 +1,6 @@
+from typing import Union
+
+from .actions import Action
+from .utils import StateInfo
+
+Trajectory = list[Union[StateInfo, Action]]
diff --git a/environment_docker/webarena-homepage/app.py b/environment_docker/webarena-homepage/app.py
index 97a9dd7..0b092c7 100644
--- a/environment_docker/webarena-homepage/app.py
+++ b/environment_docker/webarena-homepage/app.py
@@ -4,22 +4,22 @@ app = Flask(__name__)
@app.route("/")
-def index():
+def index() -> str:
return render_template("index.html")
@app.route("/scratchpad.html")
-def scratchpad():
+def scratchpad() -> str:
return render_template("scratchpad.html")
@app.route("/calculator.html")
-def calculator():
+def calculator() -> str:
return render_template("calculator.html")
@app.route("/password.html")
-def password():
+def password() -> str:
return render_template("password.html")
diff --git a/minimal_example.py b/minimal_example.py
index 173a800..9cae4d2 100644
--- a/minimal_example.py
+++ b/minimal_example.py
@@ -51,8 +51,6 @@ print("Done generating config files with the correct URLs")
subprocess.run(["bash", "prepare.sh"])
print("Done saving account cookies")
-from agent.utils import Trajectory
-
# Init an environment
from browser_env import (
Action,
@@ -60,6 +58,7 @@ from browser_env import (
ObservationMetadata,
ScriptBrowserEnv,
StateInfo,
+ Trajectory,
action2str,
create_id_based_action,
create_stop_action,
diff --git a/run.py b/run.py
index b7e41ec..7c8a7b8 100644
--- a/run.py
+++ b/run.py
@@ -1,39 +1,37 @@
"""Script to run end-to-end evaluation on the benchmark"""
import argparse
-import base64
import glob
-import io
import json
import logging
import os
import random
-import re
-import subprocess
import time
-from itertools import chain
from pathlib import Path
-from typing import Any
import openai
-import tiktoken
from beartype import beartype
-from PIL import Image
-from prompt_toolkit import prompt
-from agent import Agent, PromptAgent, TeacherForcingAgent
+from agent import (
+ Agent,
+ PromptAgent,
+ TeacherForcingAgent,
+ construct_agent,
+)
from agent.prompts import *
from browser_env import (
Action,
ActionTypes,
- ObservationMetadata,
ScriptBrowserEnv,
StateInfo,
- action2str,
+ Trajectory,
create_stop_action,
)
from browser_env.actions import is_equivalent
+from browser_env.helper_functions import (
+ RenderHelper,
+ get_action_description,
+)
from evaluation_harness import evaluator_router
-from llms import lm_config
LOG_FOLDER = "log_files"
Path(LOG_FOLDER).mkdir(parents=True, exist_ok=True)
@@ -55,24 +53,6 @@ formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)
-Trajectory = list[Action | StateInfo]
-HTML_TEMPLATE = """
-
-
-
-
-
-
- {body}
-
-
-"""
-
def config() -> argparse.Namespace:
parser = argparse.ArgumentParser(
@@ -164,165 +144,6 @@ def config() -> argparse.Namespace:
return args
-@beartype
-def get_render_action(
- action: Action,
- observation_metadata: dict[str, ObservationMetadata],
- action_set_tag: str,
-) -> str:
- """Parse the predicted actions for rendering purpose. More comprehensive information"""
- match action_set_tag:
- case "id_accessibility_tree":
- text_meta_data = observation_metadata["text"]
- if action["element_id"] in text_meta_data["obs_nodes_info"]:
- node_content = text_meta_data["obs_nodes_info"][
- action["element_id"]
- ]["text"]
- else:
- node_content = "No match found"
-
- action_str = f"
{action['raw_prediction']}"
- action_str += f"
"
- action_str += f"
{action2str(action, action_set_tag, node_content)}"
-
- case "playwright":
- action_str = action["pw_code"]
- case _:
- raise ValueError(f"Unknown action type {action['action_type']}")
- return action_str
-
-
-@beartype
-def get_action_description(
- action: Action,
- observation_metadata: dict[str, ObservationMetadata],
- action_set_tag: str,
- prompt_constructor: PromptConstructor | None,
-) -> str:
- """Generate the text version of the predicted actions to store in action history for prompt use.
- May contain hint information to recover from the failures"""
-
- match action_set_tag:
- case "id_accessibility_tree":
- text_meta_data = observation_metadata["text"]
- if action["action_type"] in [
- ActionTypes.CLICK,
- ActionTypes.HOVER,
- ActionTypes.TYPE,
- ]:
- action_name = str(action["action_type"]).split(".")[1].lower()
- if action["element_id"] in text_meta_data["obs_nodes_info"]:
- node_content = text_meta_data["obs_nodes_info"][
- action["element_id"]
- ]["text"]
- node_content = " ".join(node_content.split()[1:])
- action_str = action2str(
- action, action_set_tag, node_content
- )
- else:
- action_str = f"Attempt to perfom \"{action_name}\" on element \"[{action['element_id']}]\" but no matching element found. Please check the observation more carefully."
- else:
- if (
- action["action_type"] == ActionTypes.NONE
- and prompt_constructor is not None
- ):
- action_splitter = prompt_constructor.instruction[
- "meta_data"
- ]["action_splitter"]
- action_str = f'The previous prediction you issued was "{action["raw_prediction"]}". However, the format was incorrect. Ensure that the action is wrapped inside a pair of {action_splitter} and enclose arguments within [] as follows: {action_splitter}action [arg] ...{action_splitter}.'
- else:
- action_str = action2str(action, action_set_tag, "")
-
- case "playwright":
- action_str = action["pw_code"]
-
- case _:
- raise ValueError(f"Unknown action type {action['action_type']}")
-
- return action_str
-
-
-class RenderHelper(object):
- """Helper class to render text and image observations and meta data in the trajectory"""
-
- def __init__(
- self, config_file: str, result_dir: str, action_set_tag: str
- ) -> None:
- with open(config_file, "r") as f:
- _config = json.load(f)
- _config_str = ""
- for k, v in _config.items():
- _config_str += f"{k}: {v}\n"
- _config_str = f"
{_config_str}\n"
- task_id = _config["task_id"]
-
- self.action_set_tag = action_set_tag
-
- self.render_file = open(
- Path(result_dir) / f"render_{task_id}.html", "a+"
- )
- self.render_file.truncate(0)
- # write init template
- self.render_file.write(HTML_TEMPLATE.format(body=f"{_config_str}"))
- self.render_file.read()
- self.render_file.flush()
-
- def render(
- self,
- action: Action,
- state_info: StateInfo,
- meta_data: dict[str, Any],
- render_screenshot: bool = False,
- ) -> None:
- """Render the trajectory"""
- # text observation
- observation = state_info["observation"]
- text_obs = observation["text"]
- info = state_info["info"]
- new_content = f"
New Page
\n"
- new_content += f"
\n"
- new_content += f"
{text_obs}\n"
-
- if render_screenshot:
- # image observation
- img_obs = observation["image"]
- image = Image.fromarray(img_obs)
- byte_io = io.BytesIO()
- image.save(byte_io, format="PNG")
- byte_io.seek(0)
- image_bytes = base64.b64encode(byte_io.read())
- image_str = image_bytes.decode("utf-8")
- new_content += f"
\n"
-
- # meta data
- new_content += f"
{meta_data['action_history'][-1]}
\n"
-
- # action
- action_str = get_render_action(
- action,
- info["observation_metadata"],
- action_set_tag=self.action_set_tag,
- )
- # with yellow background
- action_str = f"
{action_str}
"
- new_content += f"{action_str}\n"
-
- # add new content
- self.render_file.seek(0)
- html = self.render_file.read()
- html_body = re.findall(r"(.*?)", html, re.DOTALL)[0]
- html_body += new_content
-
- html = HTML_TEMPLATE.format(body=html_body)
- self.render_file.seek(0)
- self.render_file.truncate()
- self.render_file.write(html)
- self.render_file.flush()
-
- def close(self) -> None:
- self.render_file.close()
-
-
@beartype
def early_stop(
trajectory: Trajectory, max_steps: int, thresholds: dict[str, int]
@@ -383,7 +204,7 @@ def early_stop(
@beartype
def test(
args: argparse.Namespace,
- agent: Agent | PromptAgent,
+ agent: Agent | PromptAgent | TeacherForcingAgent,
config_file_list: list[str],
) -> None:
scores = []
@@ -504,55 +325,12 @@ def test(
f.write(f"[Unhandled Error] {repr(e)}\n")
f.write(traceback.format_exc()) # write stack trace to file
- # logger.info(f"[Render] {render_helper.render_file.name}")
- # subprocess.run(["open", render_helper.render_file.name])
render_helper.close()
env.close()
logger.info(f"Average score: {sum(scores) / len(scores)}")
-def construct_llm_config(args: argparse.Namespace) -> lm_config.LMConfig:
- llm_config = lm_config.LMConfig(
- provider=args.provider, model=args.model, mode=args.mode
- )
- if args.provider == "openai":
- llm_config.gen_config["temperature"] = args.temperature
- llm_config.gen_config["top_p"] = args.top_p
- llm_config.gen_config["context_length"] = args.context_length
- llm_config.gen_config["max_tokens"] = args.max_tokens
- llm_config.gen_config["stop_token"] = args.stop_token
- llm_config.gen_config["max_obs_length"] = args.max_obs_length
- else:
- raise NotImplementedError(f"provider {args.provider} not implemented")
- return llm_config
-
-
-def construct_agent(args: argparse.Namespace) -> Agent:
- llm_config = construct_llm_config(args)
-
- agent: Agent
- if args.agent_type == "teacher_forcing":
- agent = TeacherForcingAgent()
- elif args.agent_type == "prompt":
- with open(args.instruction_path) as f:
- constructor_type = json.load(f)["meta_data"]["prompt_constructor"]
- tokenizer = tiktoken.encoding_for_model(llm_config.model)
- prompt_constructor = eval(constructor_type)(
- args.instruction_path, lm_config=llm_config, tokenizer=tokenizer
- )
- agent = PromptAgent(
- action_set_tag=args.action_set_tag,
- lm_config=llm_config,
- prompt_constructor=prompt_constructor,
- )
- else:
- raise NotImplementedError(
- f"agent type {args.agent_type} not implemented"
- )
- return agent
-
-
def prepare(args: argparse.Namespace) -> None:
# convert prompt python files to json
from agent.prompts import to_json
+
+
+