OpenJarvis/scripts/orchestrator/upload_to_braintrust.py

#!/usr/bin/env python
"""Upload orchestrator-SFT JSONL datasets to Braintrust so they're browsable/
filterable in the UI (by area / difficulty / dataset / routed-model / correct).

Datasets land in the SAME Braintrust project as the rollout TRACES (the
``research`` project) so a run's data + traces sit side by side. The target is
resolved from ``OJ_BRAINTRUST_PROJECT_ID`` (default: the research project id);
``--project`` overrides by NAME if you want a different project.

Each record -> a Braintrust dataset row:
  input    = the problem/question
  expected = the FINAL_ANSWER value
  tags     = [gen_model, domain, correct|incorrect, clean|dirty, kept|dropped]
  metadata = domain, area, difficulty, dataset, subsector, task_id, correct,
             clean, kept, reward, gen_model, orchestrator_model, num_tool_calls,
             n_turns, routed_models (real names)
  (the full conversation is kept under metadata.conversation for inspection)

The dataset itself carries run-level metadata (run_label, specific gen_model,
git sha, config knobs, counts, domain distribution) so future experiments are
distinguishable in the UI.

Usage (CLI):
  .venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
      data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707 \
      data/orchestrator/sft/qwen_holdout_0707.jsonl=qwen_holdout_0707

  # name is optional; defaults to {model_short}_{split}_{date}
  .venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
      data/orchestrator/sft/qwen_holdout_0707.jsonl

Programmatic (used by the pipeline auto-upload hook in make_splits.py):
  from upload_to_braintrust import autoupload
  autoupload(["data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707"],
             run_label="...")
"""
import argparse
import json
import logging
import os
import re
import subprocess
from collections import Counter
from datetime import datetime
from pathlib import Path

logger = logging.getLogger(__name__)

# The `research` project — datasets go here, same place the rollout traces land.
DEFAULT_PROJECT_ID = "c707124e-ce9d-4187-ad11-f49f19f777ad"

_FA = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")


def _ordered_turns(convs, anon_map, orchestrator):
    """Rebuild each turn as role/model/content, in that display order.

    Braintrust sorts object keys alphabetically on write, which would push the
    (very long) `content` to the top. We prefix the keys (`1_role`, `2_model`,
    `3_content`) so they sort into role -> model -> content and the content
    renders last. `model` is the model behind the turn:
      * assistant turns  -> the orchestrator model
      * tool turns       -> the expert model that answered (de-anonymized)
      * system/user      -> None
    """
    out = []
    for c in convs:
        role = c.get("role")
        if role == "tool":
            model = anon_map.get(c.get("name")) or c.get("name")
        elif role == "assistant":
            model = orchestrator
        else:
            model = None
        out.append({"1_role": role, "2_model": model, "3_content": c.get("content")})
    return out


def _dataset_desc(dataset):
    """Fallback dataset blurb for records generated before the field existed."""
    try:
        from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
            dataset_description,
        )
        return dataset_description(dataset)
    except Exception:
        return ""


def _git_sha():
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=Path(__file__).resolve().parent, stderr=subprocess.DEVNULL,
        ).decode().strip() or None
    except Exception:
        return None


def _model_short(rows, path):
    """Short family tag for the default dataset name (qwen / gemma / ...)."""
    stem = Path(path).stem.lower()
    for fam in ("qwen", "gemma"):
        if stem.startswith(fam) or fam in stem:
            return fam
    gm = next((r.get("gen_model") for r in rows if r.get("gen_model")), "") or ""
    gm = gm.lower()
    for fam in ("qwen", "gemma"):
        if fam in gm:
            return fam
    return "orch"


def _split_of(path):
    stem = Path(path).stem.lower()
    for s in ("holdout", "train", "clean", "overfit", "partial", "8k", "4k", "2k", "1k"):
        if s in stem:
            return s
    return "data"


def _default_name(rows, path):
    return f"{_model_short(rows, path)}_{_split_of(path)}_{datetime.now():%m%d}"


def _config_from_env():
    """Config knobs from the run env (set by build_orchestrator_sft.py). Omit unset."""
    cfg = {}
    for env_key, meta_key, cast in [
        ("OJ_CFG_TEMPERATURE", "temperature", float),
        ("OJ_CFG_MAX_TURNS", "max_turns", int),
        ("OJ_CFG_ANONYMIZE", "anonymize", lambda v: v.strip().lower() in ("1", "true", "yes")),
        ("OJ_CFG_REJECTION_ONLY", "rejection_only", lambda v: v.strip().lower() in ("1", "true", "yes")),
    ]:
        v = os.getenv(env_key)
        if v not in (None, ""):
            try:
                cfg[meta_key] = cast(v)
            except Exception:
                cfg[meta_key] = v
    return cfg


def _row(rec):
    convs = rec.get("conversations", [])
    user = next((c["content"] for c in convs if c["role"] == "user"), "")
    question = re.sub(r"^\s*Problem:\s*", "", user, count=1).strip()
    asst = [c["content"] for c in convs if c["role"] == "assistant"]
    fa = asst[-1] if asst else ""
    m = list(_FA.finditer(fa))
    model_answer = fa[m[-1].end():].strip() if m else fa.strip()
    # Gold reference the verifier graded against (stamped at generation time as
    # `gold_answer`). `expected` in Braintrust is the GOLD so the UI shows
    # gold-vs-model; the model's own answer goes to metadata.model_answer.
    gold = rec.get("gold_answer") or None
    has_gold = gold is not None
    am = (rec.get("metrics", {}) or {}).get("anon_map", {}) or {}
    routed = []
    for c in convs:
        if c["role"] == "assistant":
            for t in re.findall(r"<tool_call>(\{.*?\})</tool_call>", c["content"], re.DOTALL):
                try:
                    real = am.get(json.loads(t).get("name"))
                    if real:
                        routed.append(real)
                except Exception:
                    pass
    metrics = rec.get("metrics", {}) or {}
    gen_model = rec.get("gen_model")
    domain = rec.get("domain")
    correct = rec.get("correct")
    clean = rec.get("clean")
    kept = rec.get("kept")
    tags = [t for t in [
        gen_model,
        f"domain:{domain}" if domain else None,
        "correct" if correct else "incorrect",
        "clean" if clean else "dirty",
        "kept" if kept else "dropped",
        "gold" if has_gold else "no_gold",
    ] if t]
    return {
        "input": question,
        "expected": gold,
        "tags": tags,
        "metadata": {
            "gold_answer": gold,
            "model_answer": model_answer,
            "has_gold": has_gold,
            "domain": domain,
            "area": rec.get("area"),
            "difficulty": rec.get("difficulty") or None,
            "dataset": rec.get("dataset"),
            "dataset_description": rec.get("dataset_description") or _dataset_desc(rec.get("dataset")),
            "subsector": rec.get("subsector"),
            "task_id": rec.get("task_id"),
            "correct": correct,
            "clean": clean,
            "kept": kept,
            "reward": rec.get("reward"),
            "gen_model": gen_model,
            "orchestrator_model": rec.get("orchestrator_model"),
            "num_tool_calls": metrics.get("num_tool_calls"),
            "n_turns": metrics.get("num_turns"),
            "routed_models": routed,
            "conversation": _ordered_turns(convs, am, rec.get("orchestrator_model")),
        },
    }


def _dist(rows, key):
    """{value: count} for a record field, most-common first, None/'' dropped."""
    c = Counter(r.get(key) for r in rows if (r.get(key) or "") != "")
    return {str(k): v for k, v in sorted(c.items(), key=lambda x: -x[1])}


def _routed_dist(rows):
    """Distribution of real (de-anonymized) expert models routed to."""
    c = Counter()
    for r in rows:
        am = (r.get("metrics", {}) or {}).get("anon_map", {}) or {}
        for cc in r.get("conversations", []):
            if cc.get("role") != "assistant":
                continue
            for t in re.findall(r"<tool_call>(\{.*?\})</tool_call>", cc.get("content", ""), re.DOTALL):
                try:
                    real = am.get(json.loads(t).get("name"))
                    if real:
                        c[real] += 1
                except Exception:
                    pass
    return {k: v for k, v in sorted(c.items(), key=lambda x: -x[1])}


def _avg(rows, *path):
    vals = []
    for r in rows:
        v = r.get("metrics", {}) or {}
        for p in path:
            v = (v or {}).get(p) if isinstance(v, dict) else None
        if isinstance(v, (int, float)):
            vals.append(v)
    return round(sum(vals) / len(vals), 2) if vals else None


def _dataset_metadata(rows, path, run_label, gen_model):
    n = len(rows)
    meta = {
        "model": _model_short(rows, path),
        "split": _split_of(path),
        "task": "orchestrator-SFT unified-tool routing traces",
        "date": "2026-07-07",
        "run_label": run_label,
        "gen_model": gen_model,
        "git_sha": _git_sha(),
        "source_file": str(path),
        "source_datasets": ["GeneralThought-430K-filtered", "OpenThoughts3-1.2M"],
        "task_mix": "balanced (GeneralThought + OpenThoughts code/math/science)",
        "filter": "correct + clean (rejects file-write echoes, garble, "
                  "reasoning-degeneration, truncated tails, essays/markdown dumps)",
        "leak_free": "train/holdout disjoint by task_id",
        "n_total": n,
        "n_correct": sum(1 for r in rows if r.get("correct")),
        "n_clean": sum(1 for r in rows if r.get("clean")),
        "n_kept": sum(1 for r in rows if r.get("kept")),
        "n_with_gold": sum(1 for r in rows if (r.get("gold_answer") or "").strip()),
        "avg_tool_calls": _avg(rows, "num_tool_calls"),
        "avg_turns": _avg(rows, "num_turns"),
        "area_distribution": _dist(rows, "area"),
        "difficulty_distribution": _dist(rows, "difficulty"),
        "dataset_distribution": _dist(rows, "dataset"),
        "routed_model_distribution": _routed_dist(rows),
    }
    cfg = _config_from_env()
    if cfg:
        meta["config"] = cfg
    return {k: v for k, v in meta.items() if v is not None}


def upload_dataset(path, name=None, *, project_id=None, project=None,
                   run_label=None, gen_model=None, description=""):
    """Upload one JSONL file as a Braintrust dataset. Returns (name, n_rows, url).

    Target project: ``project`` (name) if given, else ``project_id`` /
    ``OJ_BRAINTRUST_PROJECT_ID`` / the research default.
    """
    import braintrust

    rows = [json.loads(l) for l in open(path) if l.strip()]
    name = name or _default_name(rows, path)
    gen_model = gen_model or os.getenv("OJ_GEN_MODEL") or \
        next((r.get("gen_model") for r in rows if r.get("gen_model")), None)
    run_label = run_label or os.getenv("OJ_RUN_LABEL") or name

    init_kwargs = {"name": name, "description": description or None,
                   "metadata": _dataset_metadata(rows, path, run_label, gen_model)}
    if project:
        init_kwargs["project"] = project
    else:
        init_kwargs["project_id"] = project_id or os.getenv(
            "OJ_BRAINTRUST_PROJECT_ID", DEFAULT_PROJECT_ID)

    ds = braintrust.init_dataset(**init_kwargs)
    for r in rows:
        ds.insert(**_row(r))
    ds.flush()
    summ = ds.summarize()
    url = getattr(summ, "dataset_url", None) or (project or init_kwargs.get("project_id"))
    return name, len(rows), url


def autoupload(specs, *, run_label=None, gen_model=None, description=""):
    """No-op-safe wrapper for pipeline hooks. Honors OJ_BRAINTRUST_AUTOUPLOAD
    (default ON) and never raises: any failure (missing key/pkg, network) is
    logged and swallowed so the data pipeline is never broken by telemetry."""
    if os.getenv("OJ_BRAINTRUST_AUTOUPLOAD", "1").strip().lower() in ("0", "false", "no", "off", ""):
        logger.info("[braintrust] autoupload disabled (OJ_BRAINTRUST_AUTOUPLOAD)")
        return
    if not os.getenv("BRAINTRUST_API_KEY"):
        logger.info("[braintrust] autoupload skipped — BRAINTRUST_API_KEY unset")
        return
    try:
        import braintrust  # noqa: F401
    except Exception:
        logger.info("[braintrust] autoupload skipped — braintrust not installed")
        return
    for spec in specs:
        path, _, name = spec.partition("=")
        try:
            nm, n, url = upload_dataset(path, name or None, run_label=run_label,
                                        gen_model=gen_model, description=description)
            logger.info("[braintrust] uploaded %s: %d rows -> %s", nm, n, url)
            print(f"[braintrust] uploaded {nm}: {n} rows -> {url}")
        except Exception as exc:
            logger.warning("[braintrust] autoupload FAILED for %s (%s) — continuing", path, exc)
            print(f"[braintrust] autoupload FAILED for {path} ({exc}) — pipeline unaffected")


def main():
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    p = argparse.ArgumentParser()
    p.add_argument("--project", default=None,
                   help="Target project by NAME (overrides OJ_BRAINTRUST_PROJECT_ID).")
    p.add_argument("--project-id", default=None,
                   help="Target project by id (default: OJ_BRAINTRUST_PROJECT_ID or research).")
    p.add_argument("--run-label", default=None)
    p.add_argument("--gen-model", default=None, help="Specific gen model id override.")
    p.add_argument("--description", default="")
    p.add_argument("specs", nargs="+", help="path.jsonl[=dataset_name]")
    args = p.parse_args()
    for spec in args.specs:
        path, _, name = spec.partition("=")
        nm, n, url = upload_dataset(
            path, name or None, project_id=args.project_id, project=args.project,
            run_label=args.run_label, gen_model=args.gen_model,
            description=args.description)
        print(f"[braintrust] {nm}: {n} rows -> {url}")


if __name__ == "__main__":
    raise SystemExit(main())