build_data.py

"""Extract showcase data from trajectories and workspace files. Copies images to showcase/images/."""
import json
import os
import glob
import shutil
import re

DATA_ROOT = os.path.join(os.path.dirname(__file__), "..", "datasets", "copaw-agent-trajectories-30datasets-fixed")
TRAJECTORIES = os.path.join(DATA_ROOT, "trajectories.jsonl")
DATASETS_DIR = os.path.join(DATA_ROOT, "datasets")
SHOWCASE_DIR = os.path.dirname(__file__)
IMAGES_DIR = os.path.join(SHOWCASE_DIR, "images")
OUTPUT = os.path.join(SHOWCASE_DIR, "data.js")

CATEGORY_MAP = {
    "amineipad_network-anoamly-dataset": "Cybersecurity",
    "anmolsingla30_natural-language-identification-dataset-for-ml": "NLP",
    "anushkasalkar_ecg-hrv-feature-dataset-for-machine-learning": "Healthcare",
    "arashnic_learn-time-series-forecasting-from-gold-price": "Finance",
    "arnavvvvv_spotify-music": "Entertainment",
    "azminetoushikwasi_ucl-202122-uefa-champions-league": "Sports",
    "brendaso_2019-coronavirus-dataset-01212020-01262020": "Healthcare",
    "chadwambles_sample-jira-data-for-agile-analysis": "Project Management",
    "itssuru_hr-employee-attrition": "HR Analytics",
    "jaderz_hospital-beds-management": "Healthcare",
    "madisonwilson123_retail-loyalty-and-churn-behavior": "Retail",
    "mexwell_drug-consumption-classification": "Healthcare",
    "mrsimple07_energy-consumption-prediction": "Energy",
    "mysarahmadbhat_toyota-used-car-listing": "Automotive",
    "nancyalaswad90_breast-cancer-dataset": "Healthcare",
    "new-york-city_nyc-dog-names": "Urban Data",
    "notkrishna_cricket-statistics-for-all-formats": "Sports",
    "nudratabbas_sql-practice-dataset-1-easy-queries": "Education",
    "rabieelkharoua_cancer-prediction-dataset": "Healthcare",
    "rabieelkharoua_predict-restaurant-customer-satisfaction-dataset": "Food & Beverage",
    "rishikeshkonapure_hr-analytics-prediction": "HR Analytics",
    "syedjaferk_top-200-youtubers-cleaned": "Social Media",
    "thedevastator_nike-usa-products-prices-descriptions-and-custom": "Retail",
    "uciml_breast-cancer-wisconsin-data": "Healthcare",
    "viramatv_coffee-shop-data": "Food & Beverage",
    "vivovinco_nba-player-stats": "Sports",
    "winston56_fortune-500-data-2021": "Business",
    "zkskhurram_lung-cancer-clinical-dataset-20152025": "Healthcare",
    "zzettrkalpakbal_full-filled-brain-stroke-dataset": "Healthcare",
}


def slug_to_title(slug):
    parts = slug.split("_", 1)
    name = parts[1] if len(parts) > 1 else parts[0]
    return name.replace("-", " ").replace("_", " ").title()


def extract_assistant_summary(messages):
    for msg in reversed(messages):
        if msg["role"] == "assistant":
            text = re.sub(r'<tool_call>.*?</tool_call>', '', msg["content"], flags=re.DOTALL).strip()
            if len(text) > 200:
                return text
    return ""


def main():
    os.makedirs(IMAGES_DIR, exist_ok=True)
    cases = []

    with open(TRAJECTORIES) as f:
        trajectories = [json.loads(line) for line in f]

    for traj in trajectories:
        meta = traj["metadata"]
        slug = meta["dataset_slug"]
        workspace_dir = os.path.join(DATASETS_DIR, slug, "workspace")

        # Copy images to showcase/images/<slug>/
        images = []
        if os.path.isdir(workspace_dir):
            slug_img_dir = os.path.join(IMAGES_DIR, slug)
            os.makedirs(slug_img_dir, exist_ok=True)
            for img_path in sorted(glob.glob(os.path.join(workspace_dir, "*.png"))):
                fname = os.path.basename(img_path)
                dest = os.path.join(slug_img_dir, fname)
                shutil.copy2(img_path, dest)
                images.append(f"images/{slug}/{fname}")

        # Read markdown report if exists
        report = ""
        if os.path.isdir(workspace_dir):
            for md_path in glob.glob(os.path.join(workspace_dir, "*.md")):
                with open(md_path) as mf:
                    report = mf.read()
                break

        # Read python scripts (first 3, truncated)
        scripts = []
        if os.path.isdir(workspace_dir):
            for py_path in sorted(glob.glob(os.path.join(workspace_dir, "*.py")))[:3]:
                fname = os.path.basename(py_path)
                with open(py_path) as pf:
                    code = pf.read()
                scripts.append({"filename": fname, "code": code[:5000]})

        if not report:
            report = extract_assistant_summary(traj["messages"])

        completed = meta["tool_iterations"] < 50
        total_tokens = meta["total_input_tokens"] + meta["total_output_tokens"]

        case = {
            "slug": slug,
            "title": slug_to_title(slug),
            "category": CATEGORY_MAP.get(slug, "Other"),
            "iterations": meta["tool_iterations"],
            "total_tokens": total_tokens,
            "input_tokens": meta["total_input_tokens"],
            "output_tokens": meta["total_output_tokens"],
            "messages_count": len(traj["messages"]),
            "completed_naturally": completed,
            "images": images,
            "report": report[:10000],
            "scripts": scripts,
        }
        cases.append(case)

    cases.sort(key=lambda c: (not c["completed_naturally"], c["iterations"]))

    stats = {
        "total_cases": len(cases),
        "total_tokens": sum(c["total_tokens"] for c in cases),
        "avg_iterations": round(sum(c["iterations"] for c in cases) / len(cases), 1),
        "natural_completion_rate": round(sum(1 for c in cases if c["completed_naturally"]) / len(cases) * 100, 1),
        "categories": sorted(set(c["category"] for c in cases)),
    }

    output = f"const SHOWCASE_DATA = {json.dumps(cases, ensure_ascii=False, indent=None)};\n\nconst STATS = {json.dumps(stats, ensure_ascii=False, indent=2)};\n"

    with open(OUTPUT, "w") as f:
        f.write(output)

    print(f"Generated {OUTPUT}")
    print(f"  Cases: {len(cases)}")
    print(f"  Total images copied: {sum(len(c['images']) for c in cases)}")
    print(f"  Cases with reports: {sum(1 for c in cases if c['report'])}")
    print(f"  Output size: {os.path.getsize(OUTPUT) / 1024:.0f} KB")


if __name__ == "__main__":
    main()