training/dataset.py

"""Generate fixed-seed training dataset for PM-Ops triage task.

Each seed is validated against the env's own RNG so the task type, difficulty,
org_config, and scenario the env generates at reset(seed=S) EXACTLY matches the
brief embedded in the prompt. Previously, the dataset generated triage briefs but
the env silently ran a different task type (release_notes, dep_update, etc.) for
the same seed — causing env_score=0 for every episode.
"""
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from server.world.org_generator import generate_org_config
from server.world.scenario_gen import generate_scenario

SEED_PREFIX = "SEED:"

# Must match pm_ops_environment.py constants exactly
_ENV_DIFFICULTY_POOL = ["easy", "medium", "medium", "hard"]
_ENV_TASK_TYPES = ["triage", "incident_routing", "release_notes", "dep_update"]


def _env_params(seed: int) -> tuple[str, str]:
    """Predict difficulty + task_type the env will choose for this seed.

    Replicates the first two RNG calls in PMOpsEnvironment.reset() so we can
    filter seeds to those that produce the task type we want.
    """
    rng = random.Random(seed)
    difficulty = rng.choice(_ENV_DIFFICULTY_POOL)
    task_type  = rng.choice(_ENV_TASK_TYPES)
    return difficulty, task_type


def _valid_triage(scenario: dict) -> bool:
    exp = scenario.get("expected", {})
    return bool(exp.get("channel")) and bool(exp.get("team"))


def generate_triage_dataset(n_episodes: int = 150, base_seed: int = 42) -> list[dict]:
    """Return dataset rows where the env WILL run a triage episode for the embedded seed.

    We pre-simulate the env's RNG to only include seeds where:
      1. env.reset(seed) picks task_type="triage"
      2. The resulting org + scenario are valid (have expected channel + team)
      3. The difficulty, org_config, and brief exactly match what the env will use

    This eliminates the mismatch where the dataset had triage briefs but the env
    graded as release_notes → guaranteed env_score=0.
    """
    rng = random.Random(base_seed)
    rows = []

    while len(rows) < n_episodes:
        seed = rng.randint(0, 2**31)

        # Only use seeds the env will run as triage
        difficulty, task_type = _env_params(seed)
        if task_type != "triage":
            continue

        # Generate org + scenario using the SAME difficulty + seed the env will use
        org, scenario = None, None
        for attempt in range(10):
            org = generate_org_config(seed + attempt, difficulty)
            scenario = generate_scenario("triage", org, seed + attempt)
            if _valid_triage(scenario):
                break

        if not _valid_triage(scenario):
            continue

        rows.append({
            "prompt": f"{SEED_PREFIX}{seed} | {scenario['brief']}",
            "seed": seed,
            "difficulty": difficulty,
        })

    return rows


def parse_seed_from_prompt(prompt: str) -> int | None:
    """Extract seed embedded by generate_triage_dataset."""
    if not prompt.startswith(SEED_PREFIX):
        return None
    try:
        return int(prompt[len(SEED_PREFIX):].split(" | ")[0])
    except ValueError:
        return None


def save_dataset(rows: list[dict], path: str) -> None:
    with open(path, "w") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")
    print(f"Saved {len(rows)} episodes → {path}")


def load_dataset(path: str) -> list[dict]:
    with open(path) as f:
        return [json.loads(line) for line in f]


if __name__ == "__main__":
    out = os.path.join(os.path.dirname(__file__), "triage_dataset.jsonl")
    rows = generate_triage_dataset(n_episodes=150)
    save_dataset(rows, out)