diff --git a/README.md b/README.md index c422340..f638542 100644 --- a/README.md +++ b/README.md @@ -53,20 +53,26 @@ python scripts/create_personas.py ### Step 2: Generate task intents (via exploration) ```bash -# 2a. Run the exploration agent. Trajectories are saved as agentlab pickles +# 2a. Build the exploration.tasks.json config that a3-explore consumes. +# One task per (persona, site) pair, written to +# agent_as_annotators/configs//exploration.tasks.json so +# that a3-explore picks it up via importlib.resources. +python scripts/generate_exploration_tasks.py -m + +# 2b. Run the exploration agent. Trajectories are saved as agentlab pickles # under $AGENTLAB_EXP_ROOT///step_*.pkl.gz. a3-explore -# 2b. Extract chat messages from each step pickle into a parallel JSON tree +# 2c. Extract chat messages from each step pickle into a parallel JSON tree # at outputs/chat_messages///step_*.json. python scripts/extract_chat_messages.py --find-latest -# 2c. For each trajectory, randomly sample N steps (default 3, skipping step 0), +# 2d. For each trajectory, randomly sample N steps (default 3, skipping step 0), # append the TASK_INTENT_PROMPT_TEMPLATE as a final user turn, and write each # prompt as outputs/task_intents/prompts//task_.step_.json. python scripts/prepare_tasks_intents_prompts.py --find-latest -# 2d. Send each prepared prompt to the Task Designer LLM. Completions land in +# 2e. Send each prepared prompt to the Task Designer LLM. Completions land in # outputs/task_intents/completions///. python scripts/generate_task_intents.py \ --exploration-model \ diff --git a/scripts/clean_trajectories.py b/scripts/clean_trajectories.py deleted file mode 100644 index 26ff570..0000000 --- a/scripts/clean_trajectories.py +++ /dev/null @@ -1,207 +0,0 @@ -import argparse -import datetime -import itertools -import os -from pathlib import Path - -try: - import orjson -except ImportError: - import json as orjson - -import openai -from tqdm.auto import tqdm - -from agent_reward_bench.judge import create_chat_messages_from_trajectory -from agent_reward_bench.judge.existing import aer, nnetnav -from agent_reward_bench.judge.args import default_judge_args, judge_args - -from agent_reward_bench.utils import ( - get_api_key_from_env_var, - get_base_url_from_env_var, - CostEstimator, -) - -def run_judge_on_single_trajectory( - traj_path, - traj_dir, - save_dir, - estimator, - pbar: tqdm, - skip_errors: bool = False, -): - """ - Function to run the judge on a single trajectory. - """ - save_dir = Path(save_dir) - traj_path = Path(traj_path) - - save_path = save_dir.joinpath(f"{traj_path.stem}.json") - error_path = save_dir.joinpath(f"{traj_path.stem}_error.json") - - # if the file already exists, skip - if save_path.exists(): - tqdm.write(f"Skipping {traj_path.name} because the file already exists") - return "exists" - - if skip_errors and error_path.exists(): - tqdm.write(f"Skipping {traj_path.name} because the error file already exists") - return "exists" - - # with open(path, "r") as f: - # trajectory = json.load(f) - with open(traj_path, "rb") as f: - trajectory = orjson.loads(f.read()) - - # if it's not valid, skip - if not trajectory["valid"]: - tqdm.write(f"Skipping {traj_path} because it's not valid") - return None - - # show the trajectory id in the progress bar - pbar.set_postfix( - task_id=traj_path.stem.split(".")[-1], estimator=estimator.total_cost - ) - - if judge == "functional": - results = { - "benchmark": trajectory["benchmark"], - "goal": trajectory["goal"], - "agent": trajectory["agent"], - "judge": judge, - "judge_model_name": "judge", - "provider": provider, - "cost": 0, - "trajectory_info": { - "valid": trajectory["valid"], - "model": trajectory["model"], - "trajectory_dir": trajectory.get("trajectory_dir"), - "seed": trajectory["seed"], - "model_args": trajectory["model_args"], - "flags": trajectory["flags"], - "summary_info": trajectory["summary_info"], - }, - } - return results - - results = { - "benchmark": trajectory["benchmark"], - "goal": trajectory["goal"], - "agent": trajectory["agent"], - "trajectory_info": { - "valid": trajectory["valid"], - "model": trajectory["model"], - "trajectory_dir": trajectory.get("trajectory_dir"), - "seed": trajectory["seed"], - "model_args": trajectory["model_args"], - "flags": trajectory["flags"], - "summary_info": trajectory["summary_info"], - }, - "reward": trajectory["reward"], - } - - breakpoint() - - return results - -benchmarks = ["a3_synth"] -agents = [ - "GenericAgent-Qwen_Qwen2.5-VL-72B-Instruct", -] - - -parser = argparse.ArgumentParser( - description="Run the judge on a set of trajectories", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, -) -parser.add_argument( - "--base_dir", - type=str, - default="trajectories/cleaned", - help="The base directory where the trajectories are stored", -) -parser.add_argument( - "--base_save_dir", - type=str, - default="trajectories/judgments", - help="The base directory where the judgments will be saved", -) -parser.add_argument( - "--agents", - nargs="+", - default=agents, - help="Select the agent to use", -) -args = parser.parse_args() - -def main(): - agents = args.agents - base_dir = Path(args.base_dir) - base_save_dir = Path(args.base_save_dir) - - cost_across_runs = 0 - for agent, benchmark in itertools.product(agents, benchmarks): - # get all json files that in the path tree of the benchmark and model - traj_dir = Path(base_dir, benchmark, agent) - save_dir = Path(base_save_dir, benchmark, agent) - - trajectories_paths = list(traj_dir.glob("**/*.json")) - - print("\nAgent:", agent) - print("Benchmark:", benchmark) - - pbar = tqdm(trajectories_paths, desc="Running judge") - - for traj_path in pbar: - # if the lock file exists, skip - save_path_lock = save_dir.joinpath(f"{traj_path.stem}.lock") - if save_path_lock.exists(): - tqdm.write(f"Skipping {traj_path.name} because the lock file exists") - continue - - # create a lock file to prevent concurrent writes - with open(save_path_lock, "w") as f: - dt_now = datetime.datetime.now().isoformat() - f.write(dt_now) - try: - out = run_judge_on_single_trajectory( - traj_path=traj_path, - traj_dir=traj_dir, - save_dir=save_dir, - skip_errors=False, # set to False to handle errors in the function - ) - # except if the user interrupts the process, then we delete the lockfile - except KeyboardInterrupt: - tqdm.write(f"Interrupted while processing {traj_path.name}, deleting lock file") - save_path_lock.unlink(missing_ok=True) - raise - - if out == "exists": - save_path_lock.unlink(missing_ok=True) - continue - - if out is None: - # unlock the lock file - save_path_lock.unlink(missing_ok=True) - tqdm.write(f"Skipping {traj_path.name} because it was not valid or an error occurred") - continue - - results = out - try: - save_path = save_dir.joinpath(f"{traj_path.stem}.json") - save_dir.mkdir(parents=True, exist_ok=True) - with open(save_path, "wb") as f: - f.write(orjson.dumps(results)) - - except KeyboardInterrupt: - tqdm.write(f"Interrupted while processing {traj_path.name}, deleting lock file") - save_path_lock.unlink(missing_ok=True) - raise - - # unlock the lock file - save_path_lock.unlink(missing_ok=True) - - print(f"Total cost across runs: ${round(cost_across_runs, 6)}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/scripts/generate_exploration_tasks.py b/scripts/generate_exploration_tasks.py new file mode 100644 index 0000000..530ff49 --- /dev/null +++ b/scripts/generate_exploration_tasks.py @@ -0,0 +1,138 @@ +""" +Generate the `exploration.tasks.json` config consumed by `a3-explore`. + +Reads personas produced by `scripts/create_personas.py` +(`outputs/personas//personas.json`) and writes one +exploration task per (persona, site) into +`agent_as_annotators/configs//exploration.tasks.json`, +which is where `a3-explore` discovers it via `importlib.resources`. +""" + +import argparse +import json +from pathlib import Path + +from agent_as_annotators.prompts import ( + TASK_EXPLORATION_PROMPT_TEMPLATE, + TASK_EXPLORATION_PROMPT_TEMPLATE_WITH_MIN_STEPS, +) +from agent_as_annotators.utils import load_all_model_configs + + +SITE_NAME_TO_START_URL = { + "shopping_admin": "__SHOPPING_ADMIN__", + "shopping": "__SHOPPING__", + "map": "__MAP__", + "gitlab": "__GITLAB__", + "reddit": "__REDDIT__", + "wikipedia": "__WIKIPEDIA__", +} + + +def format_task_record_from_persona(persona_dict, site_name, intent_id, prompt=None, min_steps=10): + """Format a single exploration task record from a persona/site pair.""" + matched_string = "I am done exploring the websites." + + persona = f"{persona_dict['name']}:\n{persona_dict['description']}" + if min_steps is not None: + intent = TASK_EXPLORATION_PROMPT_TEMPLATE_WITH_MIN_STEPS.format(persona=persona, min_steps=min_steps) + else: + intent = TASK_EXPLORATION_PROMPT_TEMPLATE.format(persona=persona) + + return { + "sites": [site_name], + "task_id": intent_id, + "intent_template_id": intent_id, + "intent": intent, + "require_login": True, + "storage_state": f"./.auth/{site_name}_state.json", + "start_url": SITE_NAME_TO_START_URL[site_name], + "geolocation": None, + "intent_template": intent, + "instantiation_dict": {}, + "require_reset": False, + "eval": { + "eval_types": ["string_match"], + "reference_answers": {"exact_match": matched_string}, + "reference_url": "", + "program_html": [], + "string_note": "", + "reference_answer_raw_annotation": matched_string, + }, + "llm_prompt": prompt, + "persona": persona_dict, + } + + +def load_personas(model_name, base_dir="outputs/personas", filename="personas.json"): + personas_path = Path(base_dir, model_name.replace("/", "_"), filename) + with open(personas_path, "r") as f: + return json.load(f) + + +def main(model_name, output_dir, min_steps=10): + model_dir_name = model_name.replace("/", "_") + personas = load_personas(model_name=model_name) + site_names = list(SITE_NAME_TO_START_URL.keys()) + + tasks = [] + for persona in personas: + for site_name in site_names: + tasks.append( + format_task_record_from_persona( + persona, site_name, len(tasks), min_steps=min_steps + ) + ) + + output_path = Path(output_dir, model_dir_name, "exploration.tasks.json") + output_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Saving {len(tasks)} tasks to {output_path}") + with open(output_path, "w") as f: + json.dump(tasks, f, indent=4) + + sample_output_path = Path(output_dir, model_dir_name, "exploration.sample.json") + with open(sample_output_path, "w") as f: + json.dump(tasks[:4], f, indent=4) + print(f"Saved a 4-task sample to {sample_output_path}") + + +if __name__ == "__main__": + this_file = Path(__file__).resolve() + model_configs_path = this_file.parent.parent / "configs" / "model_configs.json" + shorthands_to_configs = load_all_model_configs(model_configs_path) + + # Default output dir is the package's configs/ directory so that a3-explore + # picks up the generated file automatically via importlib.resources. + default_output_dir = this_file.parent.parent / "agent_as_annotators" / "configs" + + parser = argparse.ArgumentParser( + description="Generate the exploration.tasks.json config for a3-explore", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "-m", + "--model", + type=str, + default="qwen3-vl-32b-thinking", + choices=list(shorthands_to_configs.keys()), + help="Exploration model shorthand (used to namespace the output directory)", + ) + parser.add_argument( + "-o", + "--output-dir", + type=str, + default=str(default_output_dir), + help="Base directory under which /exploration.tasks.json is written", + ) + parser.add_argument( + "-n", + "--min-steps", + type=int, + default=10, + help="Minimum number of exploration steps requested in the persona prompt", + ) + args = parser.parse_args() + + model_config = shorthands_to_configs[args.model] + model = model_config["kwargs"]["model_name"] + main(model_name=model, output_dir=args.output_dir, min_steps=args.min_steps)