from __future__ import annotations import json import logging import traceback from argparse import ArgumentParser from pathlib import Path import yaml from tqdm.auto import tqdm try: from .server import load_content except ImportError: from server import load_content logger = logging.getLogger(__name__) logging.getLogger("simple_parsing").setLevel(logging.INFO) TEMPLATE = """ Trajectory Viewer
{file_path_tree}

Conversation History

{file_content}
""" try: with open(Path(__file__).parent / "style.css") as infile: STYLE_SHEET = infile.read() except Exception as e: style_file = Path(__file__).parent / "style.css" logger.error(f"Failed to load style sheet from {style_file}: {traceback.format_exc()}") raise e def _load_file(file_name, gold_patches, test_patches): try: role_map = { "user": "Computer", "assistant": "SWE-Agent", "subroutine": "SWE-Agent subroutine", "default": "Default", "system": "System", "demo": "Demonstration", } content = load_content(file_name, gold_patches, test_patches) if "history" in content and isinstance(content["history"], list): history_content = "" for index, item in enumerate(content["history"]): item_content = item.get("content", "").replace("<", "<").replace(">", ">") if item.get("agent") and item["agent"] != "primary": role_class = "subroutine" else: role_class = item.get("role", "default").lower().replace(" ", "-") element_id = f"historyItem{index}" role_name = role_map.get(item.get("role", ""), item.get("role", "")) history_content += ( f"""
""" f"""
{role_name}
""" f"""
""" f"""
{item_content}
""" f"""
""" f"""
""" f"""
""" ) return history_content else: return "No history content found." except Exception: return f"Error loading content. {traceback.format_exc()}" def _make_file_path_tree(file_path): path_parts = file_path.split("/") relevant_parts = path_parts[-3:] html_string = '
\n' for part in relevant_parts: html_string += f'
{part}
\n' html_string += "
" return html_string def save_static_viewer(file_path): if not isinstance(file_path, Path): file_path = Path(file_path) data = [] if "args.yaml" in list(map(lambda x: x.name, file_path.parent.iterdir())): args = yaml.safe_load(Path(file_path.parent / "args.yaml").read_text()) if "environment" in args and "data_path" in args["environment"]: data_path = Path(__file__).parent.parent / args["environment"]["data_path"] if data_path.exists(): with open(data_path) as f: data = json.load(f) if not isinstance(data, list) or not data or "patch" not in data[0] or "test_patch" not in data[0]: data = [] gold_patches = {x["instance_id"]: x["patch"] for x in data} test_patches = {x["instance_id"]: x["test_patch"] for x in data} content = _load_file(file_path, gold_patches, test_patches) file_path_tree = _make_file_path_tree(file_path.absolute().as_posix()) icons_path = Path(__file__).parent / "icons" relative_icons_path = find_relative_path(file_path, icons_path) style_sheet = STYLE_SHEET.replace("url('icons/", f"url('{relative_icons_path.as_posix()}/").replace( 'url("icons/', f'url("{relative_icons_path.as_posix()}/', ) data = TEMPLATE.format(file_content=content, style_sheet=style_sheet, file_path_tree=file_path_tree) output_file = file_path.with_suffix(".html") with open(output_file, "w+") as outfile: print(data, file=outfile) logger.info(f"Saved static viewer to {output_file}") def find_relative_path(from_path, to_path): # Convert paths to absolute for uniformity from_path = from_path.resolve() to_path = to_path.resolve() if from_path.is_file(): from_path = from_path.parent if to_path.is_file(): to_path = to_path.parent if not from_path.is_dir() or not to_path.is_dir(): msg = f"Both from_path and to_path must be directories, but got {from_path} and {to_path}" raise ValueError(msg) # Identify the common ancestor and the parts of each path beyond it common_parts = 0 for from_part, to_part in zip(from_path.parts, to_path.parts): if from_part != to_part: break common_parts += 1 # Calculate the '../' needed to get back from from_path to the common ancestor back_to_ancestor = [".."] * (len(from_path.parts) - common_parts) # Direct path from common ancestor to to_path to_target = to_path.parts[common_parts:] # Combine to get the relative path return Path(*back_to_ancestor, *to_target) def save_all_trajectories(directory): if not isinstance(directory, Path): directory = Path(directory) all_files = list(directory.glob("**/*.traj")) logger.info(f"Found {len(all_files)} trajectory files in {directory}") for file_path in tqdm(all_files, desc="Saving static viewers"): save_static_viewer(file_path) logger.info(f"Saved static viewers for all trajectories in {args.directory}") if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("directory", type=str, help="Directory containing trajectory files") args = parser.parse_args() save_all_trajectories(args.directory)