mp-ai-module-controller/scripts/generate_training_data.py

#!/usr/bin/env python3
"""Generate module-controller training data from dictionary artifacts."""

from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

INSTRUCTION = "Map the following natural language request to the correct action and payload."
DATASET_PREFIX = "module_controller_intents"
DEFAULT_MIN_RECORDS_PER_ACTION = 20
CUSTOM_LOCAL_LLM_ROOT = Path("/Volumes/Zoe/custom-local-llm")
DEFAULT_OUTPUT_DIR = CUSTOM_LOCAL_LLM_ROOT / "training-data"
DEFAULT_MANIFEST = CUSTOM_LOCAL_LLM_ROOT / "manifest.llm.json"
MIMIR_DATA_DESIGNER_BIN = "/home/aaron-pressey/.venvs/data-designer/bin/data-designer"
MIMIR_DATA_DESIGNER_ENV = "/home/aaron-pressey/.config/home-grown-llm-data/data-designer.env"


def main() -> int:
    args = parse_args()
    try:
        summary = run_bridge(args)
    except Exception as exc:  # noqa: BLE001 - CLI should report any bridge failure clearly.
        print(f"error: {exc}", file=sys.stderr)
        return 1

    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dictionary-dir", default="dictionary/")
    parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR))
    parser.add_argument("--manifest", default=str(DEFAULT_MANIFEST))
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--sdg-host")
    parser.add_argument("--sdg-user")
    parser.add_argument("--sdg-key")
    parser.add_argument("--sdg-model-alias", default="nvidia-text")
    parser.add_argument("--num-records", type=int, default=1000)
    parser.add_argument(
        "--timestamp",
        type=int,
        help="Override generation timestamp. Intended for deterministic tests.",
    )
    return parser.parse_args()


def run_bridge(args: argparse.Namespace) -> dict[str, Any]:
    dictionary_dir = Path(args.dictionary_dir)
    output_root = Path(args.output_dir)
    manifest_path = Path(args.manifest)
    generated_at = args.timestamp or int(time.time())
    dataset_key = f"{DATASET_PREFIX}_{generated_at}"

    artifacts = load_dictionary_artifacts(dictionary_dir)
    records = build_seed_records(artifacts, generated_at)
    train_records, validation_records = split_records_by_action(records)

    summary = {
        "dataset_key": dataset_key,
        "generated_at": generated_at,
        "dry_run": bool(args.dry_run),
        "action_count": len(artifacts["actions"]),
        "train_records": len(train_records),
        "validation_records": len(validation_records),
        "output_dir": str(output_root / dataset_key),
        "manifest": str(manifest_path),
        "sdg_requested": has_sdg_args(args),
        "sdg_model_alias": args.sdg_model_alias,
    }

    if args.dry_run:
        return summary

    dataset_dir = output_root / dataset_key
    dataset_dir.mkdir(parents=True, exist_ok=False)
    write_jsonl(dataset_dir / "train.jsonl", train_records)
    write_jsonl(dataset_dir / "validation.jsonl", validation_records)
    write_dataset_readme(dataset_dir / "README.md", summary, artifacts)

    expanded_path = None
    sdg_error = None
    if has_sdg_args(args):
        try:
            expanded_path = run_mimir_expansion(
                args=args,
                dataset_dir=dataset_dir,
                dataset_key=dataset_key,
                generated_at=generated_at,
                seed_records=train_records + validation_records,
                action_ids=sorted(artifacts["actions"].keys()),
                model_alias=args.sdg_model_alias,
            )
        except Exception as exc:  # noqa: BLE001 - keep local seed output and report remote failure.
            sdg_error = str(exc)
            (dataset_dir / "SDG_ERROR.txt").write_text(
                f"Mimir/Data Designer expansion failed after local seed generation.\n\n{sdg_error}\n"
            )

    patch_manifest(
        manifest_path=manifest_path,
        dataset_key=dataset_key,
        dataset_dir=dataset_dir,
        generated_at=generated_at,
        artifacts=artifacts,
        expanded_path=expanded_path,
        sdg_requested=has_sdg_args(args),
        sdg_model_alias=args.sdg_model_alias,
        sdg_error=sdg_error,
    )
    summary["expanded_path"] = str(expanded_path) if expanded_path else None
    summary["sdg_error"] = sdg_error
    if sdg_error:
        raise RuntimeError(f"local seed dataset was written, but Mimir expansion failed: {sdg_error}")
    return summary


def load_dictionary_artifacts(dictionary_dir: Path) -> dict[str, Any]:
    actions_index = read_json(dictionary_dir / "actions.index.json")
    codebook = read_json(dictionary_dir / "model.codebook.json")
    examples = read_jsonl(dictionary_dir / "model.examples.jsonl")
    actions = actions_index.get("actions", {})
    if not actions:
        raise ValueError("actions.index.json contains no actions")
    return {
        "actions_index": actions_index,
        "codebook": codebook,
        "examples": examples,
        "actions": actions,
        "action_to_code": invert_unique(codebook["adaptive"]["actions"], "action codes"),
        "field_to_code": invert_unique(codebook["adaptive"]["fields"], "field codes"),
    }


def build_seed_records(artifacts: dict[str, Any], generated_at: int) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for action_id in sorted(artifacts["actions"]):
        action = artifacts["actions"][action_id]
        action_records: list[dict[str, Any]] = []
        sample_payload = sample_payload_for_schema(action.get("payload_schema", {}))

        for phrase in action.get("intent_examples", []):
            action_records.append(
                build_record(artifacts, action_id, phrase, sample_payload, "seed_intent_example", generated_at)
            )

        for alias in action.get("aliases", []):
            phrase = f"{alias} {payload_phrase(action_id, sample_payload)}".strip()
            action_records.append(
                build_record(artifacts, action_id, phrase, sample_payload, "alias_variant", generated_at)
            )

        for phrase, payload in field_variants(action_id, action.get("payload_schema", {})):
            action_records.append(
                build_record(artifacts, action_id, phrase, payload, "field_variant", generated_at)
            )

        for phrase, payload in template_variants(action_id, sample_payload, artifacts["examples"]):
            action_records.append(
                build_record(artifacts, action_id, phrase, payload, "template_expansion", generated_at)
            )

        while len(action_records) < DEFAULT_MIN_RECORDS_PER_ACTION:
            ordinal = len(action_records) + 1
            phrase = f"please run {action_id} example {ordinal} with {payload_phrase(action_id, sample_payload)}"
            action_records.append(
                build_record(artifacts, action_id, phrase, sample_payload, "template_expansion", generated_at)
            )

        records.extend(dedupe_records(action_records)[:DEFAULT_MIN_RECORDS_PER_ACTION])

    return records


def build_record(
    artifacts: dict[str, Any],
    action_id: str,
    input_text: str,
    payload: dict[str, Any],
    source: str,
    generated_at: int,
) -> dict[str, Any]:
    verbose = {"action_id": action_id, "payload": payload}
    compact = compact_target(artifacts, action_id, payload)
    return {
        "instruction": INSTRUCTION,
        "input": input_text,
        "output": json.dumps(verbose, sort_keys=True),
        "metadata": {
            "action_id": action_id,
            "compact": json.dumps(compact, separators=(",", ":"), sort_keys=True),
            "source": source,
            "split": "",
            "generated_at": generated_at,
        },
    }


def compact_target(artifacts: dict[str, Any], action_id: str, payload: dict[str, Any]) -> dict[str, Any]:
    codebook = artifacts["codebook"]
    action_to_code = artifacts["action_to_code"]
    field_to_code = artifacts["field_to_code"]
    if action_id not in action_to_code:
        raise ValueError(f"missing compact action code for {action_id}")

    compact_payload = {}
    for field, value in sorted(payload.items()):
        if field not in field_to_code:
            raise ValueError(f"missing compact field code for {field}")
        compact_payload[field_to_code[field]] = value

    return {
        "v": codebook["dictionary_version"],
        "c": codebook["codebook_checksum"],
        "a": action_to_code[action_id],
        "p": compact_payload,
    }


def split_records_by_action(records: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    grouped: dict[str, list[dict[str, Any]]] = {}
    for record in records:
        grouped.setdefault(record["metadata"]["action_id"], []).append(record)

    train: list[dict[str, Any]] = []
    validation: list[dict[str, Any]] = []
    for action_id in sorted(grouped):
        action_records = grouped[action_id]
        validation_count = max(1, round(len(action_records) * 0.1)) if len(action_records) > 1 else 0
        split_at = len(action_records) - validation_count
        for record in action_records[:split_at]:
            train.append(with_split(record, "train"))
        for record in action_records[split_at:]:
            validation.append(with_split(record, "validation"))
    return train, validation


def with_split(record: dict[str, Any], split: str) -> dict[str, Any]:
    cloned = json.loads(json.dumps(record))
    cloned["metadata"]["split"] = split
    return cloned


def sample_payload_for_schema(schema: dict[str, Any]) -> dict[str, Any]:
    properties = schema.get("properties", {})
    return {field: sample_value(field, field_schema) for field, field_schema in sorted(properties.items())}


def sample_value(field: str, schema: dict[str, Any]) -> Any:
    if "enum" in schema:
        return schema["enum"][0]
    if field == "expression":
        return "2 + 2 * 3"
    if field == "value":
        return 100
    if field == "from":
        return "km"
    if field == "to":
        return "miles"
    if field == "min":
        return 1
    if field == "max":
        return 100
    if field == "max_chars":
        return 100
    if field == "filename":
        return "hello.txt"
    if field == "content":
        return "hello from codex"
    if field in {"text", "message"}:
        return "hello from codex"
    if schema.get("type") in {"number", "integer"}:
        return 1
    return f"{field} value"


def field_variants(action_id: str, schema: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]:
    properties = schema.get("properties", {})
    if not properties:
        return [(f"{action_id} with empty payload", {})]

    variants: list[tuple[str, dict[str, Any]]] = []
    base = sample_payload_for_schema(schema)
    for field, field_schema in sorted(properties.items()):
        if "enum" in field_schema:
            for enum_value in field_schema["enum"]:
                payload = dict(base)
                payload[field] = enum_value
                variants.append((f"{action_id} where {field} is {enum_value}", payload))
        elif field_schema.get("type") in {"number", "integer"}:
            for value in [0, sample_value(field, field_schema), 999]:
                payload = dict(base)
                payload[field] = value
                variants.append((f"{action_id} with {field} {value}", payload))
        else:
            payload = dict(base)
            payload[field] = sample_value(field, field_schema)
            variants.append((f"{action_id} using {field} {payload[field]}", payload))
    return variants


def template_variants(
    action_id: str,
    sample_payload: dict[str, Any],
    examples: list[dict[str, Any]],
) -> list[tuple[str, dict[str, Any]]]:
    template_count = max(1, sum(1 for example in examples if example.get("kind") == "synthetic_template"))
    variants = []
    for index in range(1, max(5, template_count) + 1):
        variants.append((f"template request {index} for {action_id}: {payload_phrase(action_id, sample_payload)}", sample_payload))
    return variants


def payload_phrase(action_id: str, payload: dict[str, Any]) -> str:
    if not payload:
        return "with empty payload"
    parts = [f"{key} {value}" for key, value in sorted(payload.items())]
    return f"{action_id} with " + ", ".join(parts)


def dedupe_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
    seen = set()
    unique = []
    for record in records:
        key = (record["input"], record["output"])
        if key in seen:
            continue
        seen.add(key)
        unique.append(record)
    return unique


def run_mimir_expansion(
    args: argparse.Namespace,
    dataset_dir: Path,
    dataset_key: str,
    generated_at: int,
    seed_records: list[dict[str, Any]],
    action_ids: list[str],
    model_alias: str,
) -> Path:
    require_sdg_args(args)
    remote_base = "/mnt/storage/data-designer/managed-assets"
    remote_artifact_path = f"{remote_base}/module_controller_artifacts_{generated_at}"
    remote_seed = f"{remote_base}/module_controller_seed_{generated_at}.jsonl"
    remote_config = f"{remote_base}/module_controller_{generated_at}.py"
    remote_expanded = f"{remote_artifact_path}/{dataset_key}.jsonl"
    expanded_path = dataset_dir / "expanded.jsonl"
    raw_expanded_path = dataset_dir / "expanded.raw.jsonl"

    with tempfile.TemporaryDirectory() as tmp:
        tmp_path = Path(tmp)
        local_seed = tmp_path / "seed.jsonl"
        local_config = tmp_path / "module_controller.py"
        write_jsonl(local_seed, seed_records)
        local_config.write_text(data_designer_config(dataset_key, generated_at, remote_seed, action_ids, model_alias))

        target = f"{args.sdg_user}@{args.sdg_host}"
        scp_base = ["scp", "-i", str(Path(args.sdg_key).expanduser())]
        ssh_base = ["ssh", "-i", str(Path(args.sdg_key).expanduser()), target]
        subprocess.run([*scp_base, str(local_seed), f"{target}:{remote_seed}"], check=True)
        subprocess.run([*scp_base, str(local_config), f"{target}:{remote_config}"], check=True)
        subprocess.run(
            [
                *ssh_base,
                remote_bash_command(
                    remote_data_designer_create_command(
                        remote_config=remote_config,
                        num_records=args.num_records,
                        dataset_key=dataset_key,
                        artifact_path=remote_artifact_path,
                    )
                ),
            ],
            check=True,
        )
        subprocess.run([*scp_base, f"{target}:{remote_expanded}", str(raw_expanded_path)], check=True)
        normalize_expanded_jsonl(raw_expanded_path, expanded_path)
    return expanded_path


def remote_data_designer_create_command(
    remote_config: str,
    num_records: int,
    dataset_key: str,
    artifact_path: str,
) -> str:
    command = [
        MIMIR_DATA_DESIGNER_BIN,
        "create",
        remote_config,
        "--num-records",
        str(num_records),
        "--dataset-name",
        dataset_key,
        "--artifact-path",
        artifact_path,
        "--output-format",
        "jsonl",
        "--no-tui",
    ]
    quoted_command = " ".join(shlex.quote(part) for part in command)
    quoted_env = shlex.quote(MIMIR_DATA_DESIGNER_ENV)
    return f"set -a; [ ! -f {quoted_env} ] || . {quoted_env}; set +a; {quoted_command}"


def remote_bash_command(command: str) -> str:
    return f"bash -lc {shlex.quote(command)}"


def data_designer_config(
    dataset_key: str,
    generated_at: int,
    remote_seed: str,
    action_ids: list[str],
    model_alias: str = "nvidia-text",
) -> str:
    action_ids_repr = repr(action_ids)
    action_ids_json = json.dumps(action_ids, separators=(",", ":"))
    return f'''# auto-generated by generate_training_data.py - do not edit by hand
# generated_at: {generated_at}

from data_designer.config import DataDesignerConfigBuilder, LocalFileSeedSource


ACTION_IDS = {action_ids_repr}


def load_config_builder():
    builder = DataDesignerConfigBuilder(
        model_configs="/mnt/storage/data-designer/model_configs.yaml",
    )
    builder.with_seed_dataset(LocalFileSeedSource(path="{remote_seed}"))
    builder.add_column(
        name="expanded_record",
        column_type="llm-structured",
        model_alias="{model_alias}",
        prompt="""You are generating additional instruction-tuning rows for a local action dispatcher.

Given this seed natural-language request and target JSON, create one new varied request
that maps to the same action_id. Keep the output schema strict.

Allowed action_id values: {action_ids_json}

Seed request:
{{{{ input }}}}

Seed output JSON:
{{{{ output }}}}

Return a complete instruction dataset row. The output field must be a JSON string with
exactly action_id and payload. The metadata.action_id must match that output action_id.""",
        output_format={{
            "type": "object",
            "required": ["instruction", "input", "output", "metadata"],
            "additionalProperties": False,
            "properties": {{
                "instruction": {{"type": "string", "minLength": 1}},
                "input": {{"type": "string", "minLength": 1}},
                "output": {{"type": "string", "minLength": 1}},
                "metadata": {{
                    "type": "object",
                    "required": ["action_id", "source", "split", "generated_at"],
                    "additionalProperties": True,
                    "properties": {{
                        "action_id": {{"type": "string", "enum": ACTION_IDS}},
                        "source": {{"type": "string"}},
                        "split": {{"type": "string"}},
                        "generated_at": {{"type": "integer"}},
                    }},
                }},
            }},
        }},
    )
    return builder
'''


def normalize_expanded_jsonl(raw_path: Path, output_path: Path) -> None:
    normalized: list[dict[str, Any]] = []
    for row in read_jsonl(raw_path):
        candidate = row.get("expanded_record") if isinstance(row, dict) else None
        if isinstance(candidate, dict):
            normalized.append(candidate)
        elif isinstance(row, dict) and {"instruction", "input", "output", "metadata"}.issubset(row):
            normalized.append(row)
    if not normalized:
        raise ValueError(f"no expanded instruction rows found in {raw_path}")
    write_jsonl(output_path, normalized)


def has_sdg_args(args: argparse.Namespace) -> bool:
    return bool(args.sdg_host or args.sdg_user or args.sdg_key)


def require_sdg_args(args: argparse.Namespace) -> None:
    missing = [name for name in ["sdg_host", "sdg_user", "sdg_key"] if not getattr(args, name)]
    if missing:
        raise ValueError(f"Mimir expansion requires all SDG flags; missing {', '.join('--' + name.replace('_', '-') for name in missing)}")


def patch_manifest(
    manifest_path: Path,
    dataset_key: str,
    dataset_dir: Path,
    generated_at: int,
    artifacts: dict[str, Any],
    expanded_path: Path | None,
    sdg_requested: bool,
    sdg_model_alias: str | None,
    sdg_error: str | None = None,
) -> None:
    manifest = read_json(manifest_path)
    datasets = manifest.setdefault("datasets", {})
    if dataset_key in datasets:
        raise ValueError(f"manifest already contains dataset {dataset_key}")

    codebook = artifacts["codebook"]
    actions_index = artifacts["actions_index"]
    datasets[dataset_key] = {
        "name": "Module Controller Intents",
        "category": "synthetic_sft",
        "schema_profile": "instruction",
        "source_url": None,
        "acquisition_method": "local_synthetic_generator",
        "license_spdx": "UNLICENSED",
        "redistribution_allowed": False,
        "estimated_tokens": None,
        "status": "generated_local_remote_failed" if sdg_error else "generated_local",
        "local_path": relative_data_path(dataset_dir),
        "local_format": "jsonl",
        "generator": "mp-ai-module-controller/scripts/generate_training_data.py",
        "sdg_engine": "data-designer" if sdg_requested else None,
        "sdg_host": "mimir (100.80.52.47)" if sdg_requested else None,
        "sdg_model": sdg_model_alias if sdg_requested else None,
        "sdg_model_alias": sdg_model_alias if sdg_requested else None,
        "expanded_path": str(expanded_path) if expanded_path else None,
        "sdg_error": sdg_error,
        "codebook_version": codebook["dictionary_version"],
        "codebook_checksum": codebook["codebook_checksum"],
        "registry_checksum": actions_index["registry_checksum"],
        "generated_at": generated_at,
        "doc": f"{relative_data_path(dataset_dir)}/README.md",
    }
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=False) + "\n")


def relative_data_path(dataset_dir: Path) -> str:
    parts = dataset_dir.parts
    if "training-data" in parts:
        index = len(parts) - 1 - list(reversed(parts)).index("training-data")
        return "/".join(parts[index:])
    if "data" in parts:
        index = len(parts) - 1 - list(reversed(parts)).index("data")
        return "/".join(parts[index:])
    return str(dataset_dir)


def write_dataset_readme(path: Path, summary: dict[str, Any], artifacts: dict[str, Any]) -> None:
    codebook = artifacts["codebook"]
    path.write_text(
        "\n".join(
            [
                f"# {summary['dataset_key']}",
                "",
                "Instruction-format seed dataset for module-controller action dispatch.",
                "",
                f"- Generated at: `{summary['generated_at']}`",
                f"- Actions: `{summary['action_count']}`",
                f"- Train records: `{summary['train_records']}`",
                f"- Validation records: `{summary['validation_records']}`",
                f"- Codebook version: `{codebook['dictionary_version']}`",
                f"- Codebook checksum: `{codebook['codebook_checksum']}`",
                "",
            ]
        )
    )


def read_json(path: Path) -> Any:
    return json.loads(path.read_text())


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    for line in path.read_text().splitlines():
        if line.strip():
            rows.append(json.loads(line))
    return rows


def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
    with path.open("w") as file:
        for record in records:
            file.write(json.dumps(record, sort_keys=True) + "\n")


def invert_unique(mapping: dict[str, str], label: str) -> dict[str, str]:
    inverted = {}
    for code, value in mapping.items():
        if value in inverted:
            raise ValueError(f"duplicate {label} value {value}")
        inverted[value] = code
    return inverted


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