activation_scanner_cli.py

#!/usr/bin/env python3
"""Product-preview CLI for the cached activation scanner."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

try:
    from .activation_scanner_core import (
        CORE_VERSION,
        DECISION_POLICY_NAME,
        load_artifact,
        read_batch_inputs,
        scan_texts,
    )
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    from research.activation_scanner_core import (  # type: ignore
        CORE_VERSION,
        DECISION_POLICY_NAME,
        load_artifact,
        read_batch_inputs,
        scan_texts,
    )


RESEARCH_DIR = Path(__file__).resolve().parent
DEFAULT_ARTIFACT = (
    RESEARCH_DIR
    / "_results"
    / "activation_scanner_artifacts"
    / "qwen-pooled-curated-core-l13-15-v2"
)
DECISION_RANK = {
    "allow": 0,
    "warn": 1,
    "block": 2,
    "quarantine": 3,
}


def artifact_complete(path: Path) -> bool:
    artifact = path.parent if path.is_file() else path
    return (artifact / "metadata.json").exists() and (artifact / "probe_weights.npz").exists()


def add_runtime_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--artifact", type=Path, default=DEFAULT_ARTIFACT, help="Cached scanner artifact directory.")
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--device", default=None)
    parser.add_argument("--dtype", choices=("auto", "float32", "bfloat16"), default=None)
    parser.add_argument("--local-files-only", action="store_true")
    parser.add_argument("--warn-threshold", type=float, default=None)
    parser.add_argument("--block-threshold", type=float, default=None)
    parser.add_argument("--top-sae-features", type=int, default=8)
    parser.add_argument(
        "--fail-on",
        choices=("never", "warn", "block", "quarantine"),
        default="never",
        help="Exit with status 2 when any result reaches this decision or higher.",
    )
    parser.add_argument(
        "--format",
        choices=("json", "summary"),
        default="json",
        help="Output JSON for hooks or a short human-readable summary.",
    )
    parser.add_argument("--pretty", action="store_true", help="Pretty-print JSON output.")


def read_scan_text(args: argparse.Namespace) -> str:
    selected = [bool(args.text), bool(args.file), bool(args.stdin)]
    if sum(selected) > 1:
        raise SystemExit("Use only one of --text, --file, or --stdin.")
    if args.text:
        return args.text
    if args.file:
        return args.file.read_text()
    if args.stdin or not sys.stdin.isatty():
        return sys.stdin.read()
    raise SystemExit("Provide --text, --file, or pipe text on stdin.")


def max_decision(results: list[dict[str, Any]]) -> str:
    if not results:
        return "allow"
    return max((str(row.get("decision", "allow")) for row in results), key=lambda decision: DECISION_RANK.get(decision, -1))


def exit_code_for(results: list[dict[str, Any]], fail_on: str) -> int:
    if fail_on == "never":
        return 0
    threshold = DECISION_RANK[fail_on]
    return 2 if any(DECISION_RANK.get(str(row.get("decision")), -1) >= threshold for row in results) else 0


def scan_with_args(texts: list[str], args: argparse.Namespace) -> list[dict[str, Any]]:
    if not artifact_complete(args.artifact):
        raise SystemExit(
            "Missing default scanner artifact. Build it with:\n"
            "research/.venv-audit/bin/python -m research.train_probe_artifact "
            "--model qwen2.5-0.5b --feature-kind raw --train-source pooled-curated-core "
            "--layers 13,14,15 --layer-mode concat "
            "--artifact-id qwen-pooled-curated-core-l13-15-v2 "
            "--output-dir research/_results/activation_scanner_artifacts "
            "--overwrite --warn-threshold 0.30 --block-threshold 0.85 --pretty"
        )
    return scan_texts(
        texts,
        args.artifact,
        batch_size=args.batch_size,
        max_length=args.max_length,
        device=args.device,
        dtype=args.dtype,
        local_files_only=args.local_files_only,
        warn_threshold=args.warn_threshold,
        block_threshold=args.block_threshold,
        top_sae_features=args.top_sae_features,
    )


def print_json(payload: Any, pretty: bool) -> None:
    print(json.dumps(payload, indent=2 if pretty else None, ensure_ascii=False))


def print_summary(results: list[dict[str, Any]]) -> None:
    for idx, row in enumerate(results, start=1):
        label = row.get("input_id") or f"input-{idx}"
        print(
            f"{label}: decision={row.get('decision')} "
            f"risk={float(row.get('risk_score', 0)):.3f} "
            f"activation={float(row.get('activation_score', 0)):.3f} "
            f"static={float(row.get('static_score', 0)):.3f}"
        )
        for reason in row.get("risk_reasons", [])[:3]:
            print(f"  - {reason}")


def command_scan(args: argparse.Namespace) -> int:
    result = scan_with_args([read_scan_text(args)], args)[0]
    if args.format == "summary":
        print_summary([result])
    else:
        print_json(result, args.pretty)
    return exit_code_for([result], args.fail_on)


def command_batch(args: argparse.Namespace) -> int:
    rows = read_batch_inputs(args.batch_file)
    results = scan_with_args([text for _, text in rows], args)
    for idx, (input_id, _) in enumerate(rows):
        if input_id is not None:
            results[idx]["input_id"] = input_id
    payload = {
        "mode": "activation_scanner_cli_batch",
        "scanner_version": CORE_VERSION,
        "decision_policy": DECISION_POLICY_NAME,
        "artifact": str(args.artifact),
        "count": len(results),
        "max_decision": max_decision(results),
        "results": results,
    }
    if args.format == "summary":
        print_summary(results)
    else:
        print_json(payload, args.pretty)
    return exit_code_for(results, args.fail_on)


def command_doctor(args: argparse.Namespace) -> int:
    complete = artifact_complete(args.artifact)
    payload: dict[str, Any] = {
        "scanner_version": CORE_VERSION,
        "decision_policy": DECISION_POLICY_NAME,
        "artifact": str(args.artifact),
        "artifact_complete": complete,
    }
    if complete:
        artifact_dir, metadata, _weights = load_artifact(args.artifact)
        payload.update(
            {
                "artifact_dir": str(artifact_dir),
                "artifact_id": metadata.get("artifact_id", artifact_dir.name),
                "model_id": metadata.get("model_id"),
                "sensor_model": metadata.get("sensor_model"),
                "feature_kind": metadata.get("feature_kind"),
                "feature_dim": metadata.get("feature_dim"),
                "thresholds": metadata.get("thresholds"),
                "layer_policy": metadata.get("layer_policy"),
            }
        )
    print_json(payload, args.pretty)
    return 0 if complete else 1


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--version", action="store_true", help="Print scanner version and exit.")
    subparsers = parser.add_subparsers(dest="command")

    scan = subparsers.add_parser("scan", help="Scan one tool, MCP, skill, hook, or prompt text.")
    scan.add_argument("--text", help="Text to scan.")
    scan.add_argument("--file", type=Path, help="File containing text to scan.")
    scan.add_argument("--stdin", action="store_true", help="Read text from stdin.")
    add_runtime_args(scan)
    scan.set_defaults(func=command_scan)

    batch = subparsers.add_parser("batch", help="Scan a JSON array of strings or {id,text} objects.")
    batch.add_argument("--batch-file", type=Path, required=True)
    add_runtime_args(batch)
    batch.set_defaults(func=command_batch)

    doctor = subparsers.add_parser("doctor", help="Check the cached scanner artifact.")
    doctor.add_argument("--artifact", type=Path, default=DEFAULT_ARTIFACT)
    doctor.add_argument("--pretty", action="store_true")
    doctor.set_defaults(func=command_doctor)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.version:
        print(CORE_VERSION)
        return 0
    if not hasattr(args, "func"):
        parser.print_help()
        return 2
    return int(args.func(args))


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