src/forge_lite/bundle.py - master - zhmax/FORGE_LITE

src/forge_lite/bundle.py
121 строка4 KB
wearetyomsmnv
add kernal version
08 апр 2026, 17:51
08 апр 2026, 17:51f147cb6
from __future__ import annotations

import csv
import io
import json
import zipfile
from pathlib import Path
from typing import Optional

# prompts file can be .txt or .csv
_PROMPTS_NAMES = ("prompts.txt", "prompts.csv")
REQUIRED_FILES = {"attack.py", "meta.json"}


def pack(
    attack_py: str | Path,
    prompts_file: Optional[str | Path],
    meta_json: str | Path,
    output: str | Path,
) -> Path:
    """Pack files into a .attack bundle (zip).

    prompts_file can be .txt / .csv, or None when meta.json contains a
    "dataset" key that points to a centrally-managed dataset in data/.
    """
    attack_py = Path(attack_py)
    meta_json = Path(meta_json)
    output = Path(output)

    for f in (attack_py, meta_json):
        if not f.exists():
            raise FileNotFoundError(f"File not found: {f}")

    meta = json.loads(meta_json.read_text(encoding="utf-8"))   # validate JSON

    with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.write(attack_py, arcname="attack.py")
        zf.write(meta_json, arcname="meta.json")

        if prompts_file is not None:
            prompts_file = Path(prompts_file)
            if not prompts_file.exists():
                raise FileNotFoundError(f"File not found: {prompts_file}")
            if prompts_file.suffix not in (".txt", ".csv"):
                raise ValueError("prompts_file must be .txt or .csv")
            zf.write(prompts_file, arcname=f"prompts{prompts_file.suffix}")
        elif "dataset" not in meta:
            raise ValueError(
                "prompts_file is required unless meta.json contains a 'dataset' key"
            )

    return output


def unpack(bundle_path: str | Path, output_dir: str | Path) -> Path:
    """Unpack a .attack bundle into a directory."""
    bundle_path = Path(bundle_path)
    output_dir = Path(output_dir)

    if not bundle_path.exists():
        raise FileNotFoundError(f"Bundle not found: {bundle_path}")

    output_dir.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(bundle_path, "r") as zf:
        zf.extractall(output_dir)

    return output_dir


def read_bundle(bundle_path: str | Path) -> tuple[str, Optional[list[str]], dict]:
    """
    Read bundle contents without extracting to disk.
    Returns: (attack_code, prompts_list_or_None, meta_dict)

    prompts_list is None when the bundle has no embedded prompts — the caller
    should resolve the prompts from meta["dataset"].
    """
    bundle_path = Path(bundle_path)

    with zipfile.ZipFile(bundle_path, "r") as zf:
        names = set(zf.namelist())
        _check_required(bundle_path.name, names)

        attack_code = zf.read("attack.py").decode("utf-8")
        meta = json.loads(zf.read("meta.json").decode("utf-8"))

        if "prompts.csv" in names:
            raw_csv = zf.read("prompts.csv").decode("utf-8")
            prompts: Optional[list[str]] = _parse_csv_prompts(raw_csv)
        elif "prompts.txt" in names:
            raw_txt = zf.read("prompts.txt").decode("utf-8")
            prompts = [line.strip() for line in raw_txt.splitlines() if line.strip()]
        else:
            prompts = None   # caller resolves via meta["dataset"]

    return attack_code, prompts, meta


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _check_required(name: str, names: set[str]) -> None:
    missing = REQUIRED_FILES - names
    if missing:
        raise ValueError(f"Bundle '{name}' is missing required files: {missing}")


def _parse_csv_prompts(raw: str) -> list[str]:
    reader = csv.DictReader(io.StringIO(raw))
    rows = list(reader)
    if not rows:
        return []
    for col in ("goal", "prompt", "text", "instruction"):
        if col in rows[0]:
            return [r[col].strip() for r in rows if r[col].strip()]
    cols = [c for c in rows[0].keys() if c not in ("", "Unnamed: 0")]
    if cols:
        return [r[cols[0]].strip() for r in rows if r[cols[0]].strip()]
    return []