red_team_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin red-team --check` — save JSON + diagram to the current working directory."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_JSON_NAME = "red-team-check.json"
_PNG_NAME = "red-team-check.png"
 
_STATUS_COLORS = {
    "pass": "#34d399",
    "warn": "#fbbf24",
    "fail": "#f87171",
}
 
 
def _normalize_red_team_result(result: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {"error": "Invalid red-team result"}
 
    if result.get("error"):
        return dict(result)
 
    data = result
    if isinstance(result.get("content"), dict):
        data = result["content"]
    card = result.get("card")
    if isinstance(card, dict) and isinstance(card.get("data"), dict):
        data = card["data"]
 
    return dict(data)
 
 
def has_red_team_payload(result: dict[str, Any]) -> bool:
    data = _normalize_red_team_result(result)
    if data.get("error"):
        return False
    return bool(data.get("vectors")) or data.get("composite_score") is not None
 
 
def _json_safe_payload(flat: dict[str, Any]) -> dict[str, Any]:
    """Omit bulky per-probe raw blobs from check export."""
    out = dict(flat)
    vectors = []
    for row in flat.get("vectors") or []:
        if not isinstance(row, dict):
            continue
        slim = {k: v for k, v in row.items() if k != "raw"}
        vectors.append(slim)
    out["vectors"] = vectors
    return out
 
 
def write_red_team_check(
    result: dict[str, Any],
    *,
    tool_name: str | None,
    cwd: str | Path,
) -> tuple[Path, Path]:
    cwd = Path(cwd)
    json_path = cwd / _JSON_NAME
    png_path = cwd / _PNG_NAME
 
    flat = _normalize_red_team_result(result)
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **_json_safe_payload(flat),
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_red_team_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_red_team_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    if result.get("error"):
        fig, ax = plt.subplots(figsize=(6, 2), facecolor="#0f1117")
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {result['error']}", ha="center", va="center", wrap=True, color="#e5e7eb")
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    model = str(result.get("model_id") or "")
    composite = float(result.get("composite_score") or 0)
    vectors = result.get("vectors") or []
    n_pass = sum(1 for v in vectors if v.get("status") == "pass")
    n_warn = sum(1 for v in vectors if v.get("status") == "warn")
    n_fail = sum(1 for v in vectors if v.get("status") == "fail")
 
    title = f"Red team — {model}" if model else "Red team"
    subtitle = f"composite {composite:.0%}  ·  {n_pass} pass  ·  {n_warn} warn  ·  {n_fail} fail"
 
    fig, axes = plt.subplots(1, 2, figsize=(11, max(4.5, min(10, len(vectors) * 0.45 + 2))), facecolor="#0f1117")
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=1.02)
    fig.text(0.5, 0.96, subtitle, ha="center", fontsize=8, color="#9ca3af")
 
    _plot_vector_scores(axes[0], vectors)
    _plot_composite_panel(axes[1], composite, n_pass, n_warn, n_fail)
 
    for ax in axes:
        ax.set_facecolor("#0f1117")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    fig.tight_layout(rect=(0, 0, 1, 0.93))
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_vector_scores(ax: Any, vectors: list[dict[str, Any]]) -> None:
    ax.set_title("Attack vector robustness", color="#e5e7eb", fontsize=10)
 
    if not vectors:
        ax.axis("off")
        ax.text(0.5, 0.5, "No vectors", ha="center", va="center", color="#9ca3af")
        return
 
    rows = sorted(vectors, key=lambda v: float(v.get("score") or 0))
    labels = [str(v.get("label") or v.get("id") or "?") for v in rows]
    vals = [float(v.get("score") or 0) for v in rows]
    colors = [_STATUS_COLORS.get(str(v.get("status") or ""), "#78716c") for v in rows]
 
    y_pos = list(range(len(rows)))
    ax.barh(y_pos, vals, color=colors, height=0.65, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=8, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Robustness score (0–1)", color="#9ca3af", fontsize=8)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
    for i, v in enumerate(vals):
        ax.text(min(v + 0.02, 0.98), i, f"{v:.0%}", va="center", fontsize=7, color="#9ca3af")
 
 
def _plot_composite_panel(ax: Any, composite: float, n_pass: int, n_warn: int, n_fail: int) -> None:
    ax.set_title("Composite & status mix", color="#e5e7eb", fontsize=10)
 
    ax.barh([0], [composite], color=_score_color(composite), height=0.35, alpha=0.92)
    ax.set_xlim(0, 1.05)
    ax.set_yticks([0])
    ax.set_yticklabels(["composite"], fontsize=9, color="#d1d5db")
    ax.set_xlabel("Composite score", color="#9ca3af", fontsize=8)
    ax.text(min(composite + 0.02, 0.98), 0, f"{composite:.0%}", va="center", fontsize=9, color="#9ca3af")
    ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
 
    counts = [n_pass, n_warn, n_fail]
    if any(counts):
        inset = ax.inset_axes([0.08, 0.08, 0.84, 0.42])
        inset.set_facecolor("#0f1117")
        labels = ["pass", "warn", "fail"]
        colors = [_STATUS_COLORS[k] for k in labels]
        x = list(range(3))
        inset.bar(x, counts, color=colors, width=0.6, alpha=0.9)
        inset.set_xticks(x)
        inset.set_xticklabels(labels, fontsize=8, color="#d1d5db")
        inset.tick_params(axis="y", colors="#6b7280", labelsize=7)
        for spine in inset.spines.values():
            spine.set_color("#374151")
 
 
def _score_color(val: float) -> str:
    if val >= 0.65:
        return "#34d399"
    if val >= 0.35:
        return "#fbbf24"
    return "#f87171"