scripts/build_preference_pairs.py

"""Generate R2Vul-style preference pairs (DPO/RLAIF) for OpsGuard.

For each (scenario, seed, step), get the correct action from
keyword_security_triager and pair it against a contrastive rejected sample
(plausible-but-wrong reasoning OR confident wrong action).

Usage:
    python scripts/build_preference_pairs.py --out data/preference_pairs.jsonl
"""
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from eval.policies import build_policy
from models import ActionType, OpsguardAction
from scripts.system_prompt import SYSTEM_PROMPT, format_observation
from server.opsguard_environment import OpsguardEnvironment


DEFAULT_SCENARIOS = [
    "E0_normal_intake",
    "E1_typosquat_wave",
    "E2_social_eng_buildup",
    "E3_compromised_maintainer",
    "E4_multi_vector",
]


_PLAUSIBLE_WRONG_REASONS = {
    "reject_pr": [
        "performance optimization, looks fine",
        "minor refactor — author has good track record",
        "stylistic cleanup, no functional impact",
        "diff is short and well-formatted, safe to land",
    ],
    "merge_pr": [
        "small change but author is unverified, hold off",
        "looks suspicious — defer until reviewed",
        "diff modifies sensitive paths, needs review",
    ],
    "close_spam": [
        "valid bug report, just needs more detail",
        "user is frustrated but issue may be real",
        "low-effort but possibly legitimate",
    ],
    "label": [
        "irrelevant chatter, ignore",
        "spam-adjacent, lean toward closing",
        "duplicate of an unrelated issue",
    ],
    "request_info": [
        "issue is clear, just merge as-is",
        "self-explanatory, label and move on",
        "no need for repro, the title says it all",
    ],
    "flag_security": [
        "ordinary feature request, no security angle",
        "user is just being polite, normal request",
        "looks like a standard support ticket",
    ],
    "comment": [
        "ignore — low-priority noise",
        "close as duplicate of nothing in particular",
    ],
    "assign": [
        "leave unassigned, will sort itself out",
        "queue for triage later, no owner needed",
    ],
    "link_duplicate": [
        "treat as new issue, ignore prior context",
        "different root cause, keep separate",
    ],
    "query_history": [
        "decide blindly, history won't help",
        "skip the lookup, act now",
    ],
    "wait": [
        "act immediately even without context",
        "no need to defer, make a call now",
    ],
}


_CONFIDENT_WRONG_REASON = {
    "merge_pr": "trusted contributor, small change",
    "close_spam": "appears to be spam",
    "label": "obvious bug, no further triage needed",
    "wait": "nothing actionable here",
}


def _action_to_dict(action: OpsguardAction) -> dict:
    return action.model_dump(mode="json", exclude_none=True)


def _serialize(action: OpsguardAction) -> str:
    return json.dumps(_action_to_dict(action), ensure_ascii=False)


def _make_contrastive_reasoning(chosen: OpsguardAction, rng: random.Random) -> OpsguardAction:
    at = chosen.action_type.value
    pool = _PLAUSIBLE_WRONG_REASONS.get(at, ["unclear rationale", "going with gut feeling"])
    new_reasoning = rng.choice(pool)
    if new_reasoning == (chosen.reasoning or ""):
        new_reasoning = new_reasoning + " (revised)"
    return chosen.model_copy(update={"reasoning": new_reasoning})


def _make_wrong_action_confident(
    chosen: OpsguardAction,
    is_attack: bool,
    rng: random.Random,
) -> OpsguardAction:
    target = chosen.target_issue_id
    if is_attack:
        return OpsguardAction(
            action_type=ActionType.MERGE_PR,
            target_issue_id=target,
            reasoning=_CONFIDENT_WRONG_REASON["merge_pr"],
        )
    if chosen.action_type == ActionType.CLOSE_SPAM:
        return OpsguardAction(
            action_type=ActionType.LABEL,
            target_issue_id=target,
            label="bug",
            reasoning="seems like a real bug worth tracking",
        )
    return OpsguardAction(
        action_type=ActionType.CLOSE_SPAM,
        target_issue_id=target,
        reasoning=_CONFIDENT_WRONG_REASON["close_spam"],
    )


def _build_prompt(obs) -> str:
    return SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + format_observation(obs)


def _is_correct(action: OpsguardAction, last_obs_metadata: dict | None) -> bool:
    # Treat any non-trivial action with non-negative step reward as "correct".
    if last_obs_metadata is None:
        return True
    step = last_obs_metadata.get("step") or last_obs_metadata
    if isinstance(step, dict):
        total = step.get("total")
        if isinstance(total, (int, float)):
            return total >= 0.0
    return True


def generate_pairs(
    scenarios: list[str],
    seeds: list[int],
    pairs_per_scenario: int,
    rng: random.Random,
) -> list[dict]:
    policy = build_policy("keyword_security_triager")
    out: list[dict] = []

    for sid in scenarios:
        scenario_pairs: list[dict] = []
        if pairs_per_scenario <= 0:
            continue

        for seed in seeds:
            if len(scenario_pairs) >= pairs_per_scenario:
                break
            env = OpsguardEnvironment()
            try:
                obs = env.reset(scenario_id=sid, seed=seed)
            except Exception:
                break

            step_idx = 0
            while not obs.done:
                if obs.current_issue is None:
                    break
                ci = obs.current_issue
                action = policy(obs)

                if action.action_type == ActionType.WAIT:
                    try:
                        obs = env.step(action)
                    except Exception:
                        break
                    step_idx += 1
                    continue

                prompt = _build_prompt(obs)
                next_obs = None
                try:
                    next_obs = env.step(action)
                except Exception:
                    break

                last_meta = next_obs.metadata if next_obs is not None else None
                if not _is_correct(action, last_meta):
                    obs = next_obs
                    step_idx += 1
                    continue

                is_attack = False
                if last_meta:
                    recent = last_meta.get("step") if isinstance(last_meta, dict) else None
                    _ = recent  # kept for future structured reads
                if next_obs is not None and next_obs.recent_actions:
                    last = next_obs.recent_actions[-1]
                    is_attack = bool(last.get("is_attack"))

                kind = "contrastive_reasoning" if rng.random() < 0.5 else "wrong_action_confident"
                if kind == "contrastive_reasoning":
                    rejected = _make_contrastive_reasoning(action, rng)
                else:
                    rejected = _make_wrong_action_confident(action, is_attack, rng)

                chosen_str = _serialize(action)
                rejected_str = _serialize(rejected)
                if chosen_str == rejected_str:
                    rejected = _make_contrastive_reasoning(
                        action.model_copy(update={"reasoning": (action.reasoning or "") + " "}),
                        rng,
                    )
                    rejected_str = _serialize(rejected)
                    if chosen_str == rejected_str:
                        forced = action.model_copy(update={"reasoning": "alternate rationale placeholder"})
                        rejected_str = _serialize(forced)
                        kind = "contrastive_reasoning"

                record = {
                    "prompt": prompt,
                    "chosen": chosen_str,
                    "rejected": rejected_str,
                    "scenario": sid,
                    "step": step_idx,
                    "issue_id": ci.issue_id,
                    "is_attack": is_attack,
                    "kind": kind,
                }
                scenario_pairs.append(record)
                step_idx += 1
                obs = next_obs

                if len(scenario_pairs) >= pairs_per_scenario:
                    break

        out.extend(scenario_pairs[:pairs_per_scenario])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/preference_pairs.jsonl")
    ap.add_argument("--max-pairs", type=int, default=1000)
    ap.add_argument("--scenarios", nargs="*", default=DEFAULT_SCENARIOS)
    ap.add_argument("--seeds", nargs="*", type=int, default=[0, 1, 2, 3])
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    n_scen = max(1, len(args.scenarios))
    pairs_per_scenario = max(1, args.max_pairs // n_scen)

    pairs = generate_pairs(args.scenarios, args.seeds, pairs_per_scenario, rng)
    pairs = pairs[: args.max_pairs]

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "w", encoding="utf-8") as f:
        for rec in pairs:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")

    by_scen: dict[str, int] = {}
    by_kind: dict[str, int] = {}
    for r in pairs:
        by_scen[r["scenario"]] = by_scen.get(r["scenario"], 0) + 1
        by_kind[r["kind"]] = by_kind.get(r["kind"], 0) + 1
    print(f"DONE: wrote {len(pairs)} preference pairs to {out}")
    print(f"  by scenario: {by_scen}")
    print(f"  by kind:     {by_kind}")


if __name__ == "__main__":
    main()