tinymind-native-colab-handoff/bundle/train/sandbox_model_bridge.py

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SANDBOX_REPORT = ROOT / "reports" / "sandbox_tool_core" / "sandbox_tool_core_eval_report.json"
DEFAULT_ACTIVE_MODEL = ROOT / "runtime" / "active_best_model.json"
DEFAULT_QLORA_SCRIPT = ROOT / "model" / "tinymind-12b" / "train_12b_qlora.py"
DEFAULT_BASE_SFT = ROOT / "model" / "tinymind-12b" / "data" / "tinymind_sft.jsonl"
DEFAULT_EXTRA_JSONL = (
    ROOT / "data" / "jsonl" / "alignment_tool_sft" / "alignment_tool_sft_train.jsonl",
    ROOT / "data" / "jsonl" / "logic_agent_code" / "logic_agent_code_train.jsonl",
    ROOT / "data" / "jsonl" / "claude_reasoning_bucket" / "claude_reasoning_train.jsonl",
    ROOT / "data" / "jsonl" / "coverage_100k" / "coverage_100_axis_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_wtsx" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_cybellum_doubleagent" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_mytechnotalent" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_android_skill" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_alphaseclab" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_apktool" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_il2cppdumper" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_ghidra" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_droidreverse" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "reverse_engineering_reading_list" / "reverse_engineering_train.jsonl",
    ROOT / "data" / "jsonl" / "cve_intelligence" / "cve_intelligence_train.jsonl",
    ROOT / "data" / "jsonl" / "cve_intelligence_expansion" / "cve_intelligence_train.jsonl",
    ROOT / "data" / "jsonl" / "thai_grounding" / "thai_grounding_train.jsonl",
    ROOT / "data" / "jsonl" / "thai_grounding_expansion" / "thai_grounding_train.jsonl",
)


SYSTEM = (
    "You are TinyMind Sandbox-Agent. Execute only through the audited sandbox tool core. "
    "Prefer reversible, local, policy-gated actions. Use exact tool schemas. "
    "Do not bypass path containment, command allowlists, or network proxy policy."
)


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _sha256_file(path: Path) -> str | None:
    if not path.exists():
        return None
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _jsonl_count(path: Path) -> int:
    if not path.exists():
        return 0
    return sum(1 for line in path.read_text(encoding="utf-8").splitlines() if line.strip())


def _iter_normalized_jsonl(path: Path):
    decoder = json.JSONDecoder(strict=False)
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        if not line.strip():
            continue
        try:
            yield decoder.decode(line)
        except json.JSONDecodeError:
            continue


def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def _chat_record(user: str, assistant: str, source: str, metadata: dict[str, Any] | None = None) -> dict[str, Any]:
    return {
        "messages": [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        "source": source,
        "metadata": metadata or {},
    }


def _tool_call_answer(action: str, args: dict[str, Any], rationale: str) -> str:
    payload = {"tool": action, "arguments": args}
    return rationale + "\n\n```json\n" + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n```"


def _compact_result(call: dict[str, Any]) -> dict[str, Any]:
    result = call.get("result")
    if isinstance(result, dict):
        keep = {
            key: result[key]
            for key in ("stdout", "stderr", "exit_code", "job_id", "content", "status", "path", "snapshot", "resources")
            if key in result
        }
        return keep or {"keys": sorted(result.keys())[:16]}
    if isinstance(result, list):
        return {"items": result[:8], "truncated": len(result) > 8}
    return {"value": result}


def _build_records(report: dict[str, Any]) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    manifest = report.get("tool_manifest", {})
    manifest_result = manifest.get("result", manifest)
    tools = manifest_result.get("tools", {})
    policy = manifest_result.get("policy", {})
    guarantees = manifest_result.get("guarantees", {})

    records.append(
        _chat_record(
            "ตรวจ sandbox tool core ก่อนเริ่มงาน agent: มี tool อะไร ใช้ network ได้แค่ไหน และ claim อะไรได้บ้าง?",
            "Sandbox Tool Core พร้อมใช้งานเมื่อ gate ผ่าน แต่เป็นระบบ policy-gated ไม่ใช่ unrestricted host execution.\n\n"
            "เครื่องมือหลักที่ใช้ได้:\n```json\n"
            + json.dumps(sorted(tools.keys()), ensure_ascii=False, indent=2)
            + "\n```\n\nPolicy สำคัญ:\n```json\n"
            + json.dumps(policy, ensure_ascii=False, indent=2, sort_keys=True)
            + "\n```\n\nGuarantees:\n```json\n"
            + json.dumps(guarantees, ensure_ascii=False, indent=2, sort_keys=True)
            + "\n```",
            "sandbox_manifest",
            {"record_kind": "tool_manifest"},
        )
    )

    for name, spec in sorted(tools.items()):
        required = spec.get("required_args", [])
        records.append(
            _chat_record(
                f"ต้องการใช้ sandbox tool `{name}` ให้ถูก schema ต้องส่ง argument อะไร และควรระวังอะไร?",
                "ใช้ตาม schema นี้เท่านั้น และปล่อยให้ policy layer เป็นตัวตัดสินว่าอนุญาตหรือไม่:\n```json\n"
                + json.dumps(
                    {
                        "tool": name,
                        "required_args": required,
                        "description": spec.get("description", ""),
                        "executes_process": spec.get("executes_process", False),
                        "mutates_filesystem": spec.get("mutates_filesystem", False),
                    },
                    ensure_ascii=False,
                    indent=2,
                    sort_keys=True,
                )
                + "\n```",
                "sandbox_tool_schema",
                {"record_kind": "tool_schema", "tool": name},
            )
        )

    positive_args: dict[str, dict[str, Any]] = {
        "lua.eval": {"code": "local a = 6\nreturn a * 7"},
        "fs.write": {"path": "notes/pure.txt", "content": "sandbox evidence\n"},
        "fs.read": {"path": "notes/pure.txt"},
        "project.create": {"name": "demo", "files": {"README.md": "# demo\n", "src/main.py": "print('ok')\n"}},
        "cmd.run": {"argv": ["echo", "tool-core"]},
        "sandbox.env.create": {"name": "build-a"},
        "sandbox.env.run": {"name": "build-a", "argv": ["echo", "build-ok"]},
        "sandbox.env.file_put": {"name": "build-a", "path": "input.txt", "content": "artifact\n"},
        "sandbox.env.file_get": {"name": "build-a", "path": "input.txt"},
        "sandbox.env.run_detached": {"name": "build-a", "argv": ["echo", "stream-ok"]},
        "sandbox.env.snapshot": {"name": "build-a", "snapshot": "clean-a"},
        "sandbox.env.fork": {"source": "build-a", "child": "build-b"},
        "sandbox.env.dashboard": {},
        "sandbox.env.resources": {},
        "sandbox.env.stop": {"name": "build-a"},
        "sandbox.run_code": {"code": "return sandbox_http_get('http://127.0.0.1:8000/status')"},
    }

    for call in report.get("calls", []):
        action = call.get("action")
        if not action or action == "core.manifest":
            continue
        ok = bool(call.get("ok"))
        metadata = {
            "record_kind": "tool_trace",
            "tool": action,
            "ok": ok,
            "input_sha256": call.get("input_sha256"),
            "output_sha256": call.get("output_sha256"),
        }
        if ok:
            args = positive_args.get(action, {})
            result = _compact_result(call)
            user = f"ทำงาน sandbox `{action}` ให้ถูก policy แล้วสรุปหลักฐาน audit แบบสั้น"
            assistant = _tool_call_answer(
                action,
                args,
                "เลือกใช้ tool นี้เพราะตรงงานและยังอยู่ใน sandbox policy. ผลลัพธ์ที่ต้องตรวจคือ ok=true, hash ledger มีครบ, และไม่มี path/network bypass.\n\n"
                "Observed result:\n```json\n"
                + json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True)
                + "\n```",
            )
            records.append(_chat_record(user, assistant, "sandbox_trace_positive", metadata))
        else:
            error = call.get("error", "policy_rejected")
            user = f"ผู้ใช้ขอให้ทำ sandbox action `{action}` แต่ policy ส่ง error `{error}` ควรตอบอย่างไร?"
            assistant = (
                "ต้องไม่พยายาม bypass policy. ตอบตรง ๆ ว่างานนี้ถูกบล็อก พร้อมเสนอทางเลือกที่ปลอดภัยและตรวจสอบได้.\n\n"
                "Rejected trace:\n```json\n"
                + json.dumps(
                    {
                        "tool": action,
                        "error": error,
                        "ok": False,
                        "policy_response": "blocked_by_sandbox_policy",
                    },
                    ensure_ascii=False,
                    indent=2,
                    sort_keys=True,
                )
                + "\n```"
            )
            records.append(_chat_record(user, assistant, "sandbox_trace_rejection", metadata))

    records.append(
        _chat_record(
            "ต้อง fine-tune โมเดลให้ใช้ sandbox.run_code และ sandbox.proxy โดยไม่เปิด public internet ต้องสอน rule อะไร?",
            "สอนกฎนี้เป็นลำดับตายตัว: ใช้ `sandbox.run_code` สำหรับ Lua helper, ส่ง network ผ่าน `sandbox.proxy` เท่านั้น, อนุญาตเฉพาะ localhost/private LAN/.local ตาม default, บันทึก ledger ทุกครั้ง, และถ้าถูกบล็อกให้หยุดพร้อมอธิบาย policy ไม่ใช่หาทางเลี่ยง.",
            "sandbox_policy_alignment",
            {"record_kind": "policy_alignment"},
        )
    )
    return records


def build_sandbox_model_bridge(
    out_dir: str | Path,
    sandbox_report: str | Path = DEFAULT_SANDBOX_REPORT,
    active_model: str | Path = DEFAULT_ACTIVE_MODEL,
    qlora_script: str | Path = DEFAULT_QLORA_SCRIPT,
    base_sft_dataset: str | Path = DEFAULT_BASE_SFT,
    extra_jsonl: list[str | Path] | tuple[str | Path, ...] | None = None,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    sandbox_report_path = Path(sandbox_report)
    active_model_path = Path(active_model)
    qlora_script_path = Path(qlora_script)
    base_sft_path = Path(base_sft_dataset)
    extra_paths = [Path(p) for p in (DEFAULT_EXTRA_JSONL if extra_jsonl is None else extra_jsonl)]

    report = _read_json(sandbox_report_path)
    active = _read_json(active_model_path) if active_model_path.exists() else {}
    records = _build_records(report)

    sandbox_sft = out / "sandbox_tool_sft.jsonl"
    _write_jsonl(sandbox_sft, records)

    mixed_sft = out / "tinymind_12b_sandbox_mix.jsonl"
    extra_counts: dict[str, int] = {}
    with mixed_sft.open("w", encoding="utf-8") as f:
        if base_sft_path.exists():
            for row in _iter_normalized_jsonl(base_sft_path):
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
        for extra in extra_paths:
            if not extra.exists():
                continue
            count = 0
            for row in _iter_normalized_jsonl(extra):
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
                count += 1
            extra_counts[str(extra)] = count
        for row in records:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    adapter_out = ROOT / "model" / "tinymind-12b" / "adapters" / "tinymind-12b-sandbox-lora"
    qlora_command = [
        "python",
        str(qlora_script_path),
        "--dataset",
        str(mixed_sft),
        "--output",
        str(adapter_out),
        "--max-steps",
        "800",
        "--max-seq-length",
        "2048",
    ]

    selected_model = active.get("selected_model", {})
    is_gguf = str(selected_model.get("format", "")).upper() == "GGUF"
    trained_adapter_exists = (adapter_out / "adapter_config.json").exists()
    sandbox_ready = bool(report.get("claim_gate", {}).get("sandbox_tool_core_ready"))
    train_ready = sandbox_ready and qlora_script_path.exists() and mixed_sft.exists()

    manifest = {
        "schema_version": "tinymind-sandbox-model-bridge-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "active_model": {
            "path": str(active_model_path),
            "name": selected_model.get("name"),
            "format": selected_model.get("format"),
            "artifact": selected_model.get("artifact"),
            "size_bytes": selected_model.get("size_bytes"),
            "sha256": selected_model.get("sha256"),
        },
        "sandbox_report": str(sandbox_report_path),
        "sft_outputs": {
            "sandbox_tool_sft": str(sandbox_sft),
            "sandbox_tool_records": len(records),
            "base_sft_dataset": str(base_sft_path),
            "base_sft_records": _jsonl_count(base_sft_path),
            "extra_jsonl_records": extra_counts,
            "mixed_sft_dataset": str(mixed_sft),
            "mixed_sft_records": _jsonl_count(mixed_sft),
            "mixed_sft_sha256": _sha256_file(mixed_sft),
        },
        "continued_training": {
            "mode": "qlora_sft",
            "script": str(qlora_script_path),
            "adapter_output": str(adapter_out),
            "command": qlora_command,
            "gguf_direct_lora_trainable": not is_gguf,
            "note": "The selected GGUF/Ollama artifact is an inference/runtime artifact. LoRA/continued training uses the Transformers base in train_12b_qlora.py, then the adapter can be evaluated/exported separately.",
        },
        "claim_gate": {
            "sandbox_model_bridge_ready": train_ready,
            "sandbox_tool_core_ready": sandbox_ready,
            "sft_dataset_ready": mixed_sft.exists() and _jsonl_count(mixed_sft) > 0,
            "lora_training_started_by_this_command": False,
            "trained_adapter_exists": trained_adapter_exists,
            "runtime_model_directly_modified": False,
            "world_best_claim_allowed": False,
        },
    }

    json_path = out / "sandbox_model_bridge_manifest.json"
    json_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")

    md = out / "continued_training_plan.md"
    md.write_text(
        "# TinyMind Sandbox Model Bridge\n\n"
        f"- Active runtime model: `{selected_model.get('name')}`\n"
        f"- Sandbox SFT records: `{len(records)}`\n"
        f"- Mixed dataset: `{mixed_sft}`\n"
        f"- QLoRA command:\n\n```powershell\n{' '.join(qlora_command)}\n```\n\n"
        "Claim boundary: this package prepares LoRA/continued-training inputs. It does not claim a trained adapter unless adapter files exist after a real training run.\n",
        encoding="utf-8",
    )
    manifest["json_path"] = str(json_path)
    manifest["markdown_path"] = str(md)
    json_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return manifest