audit_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin audit --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 = "audit-check.json"
_PNG_NAME = "audit-check.png"
 
_STATUS_COLORS = {
    "suppressed": "#f87171",
    "softened": "#fbbf24",
    "unfiltered": "#34d399",
}
 
 
def _normalize_audit_result(result: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {"error": "Invalid audit result"}
 
    data = result
    card = result.get("card")
    if isinstance(card, dict) and isinstance(card.get("data"), dict):
        data = {**card["data"], "model_id": result.get("model_id") or card["data"].get("model_id")}
    elif isinstance(result.get("content"), dict):
        data = result["content"]
 
    if result.get("error"):
        data = {**data, "error": result["error"]}
    return data
 
 
def has_audit_payload(result: dict[str, Any]) -> bool:
    data = _normalize_audit_result(result)
    return any(data.get(k) for k in ("consistency", "suppression", "boundary"))
 
 
def write_audit_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_audit_result(result)
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **flat,
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_audit_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_audit_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    if result.get("error") and not has_audit_payload(result):
        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 "")
    title = f"Audit — {model}" if model else "Audit"
 
    fig, axes = plt.subplots(1, 3, figsize=(12, 4.8), facecolor="#0f1117")
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=1.02)
 
    _plot_consistency_panel(axes[0], result.get("consistency") or {})
    _plot_suppression_panel(axes[1], result.get("suppression") or {})
    _plot_boundary_panel(axes[2], result.get("boundary") or {})
 
    for ax in axes:
        ax.set_facecolor("#0f1117")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_consistency_panel(ax: Any, data: dict[str, Any]) -> None:
    ax.set_title("Consistency", color="#e5e7eb", fontsize=10)
 
    score = data.get("consistency_score")
    query = _trunc(str(data.get("query") or ""), 40)
    variants = data.get("variants") or []
 
    if score is None and not variants:
        ax.axis("off")
        ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#9ca3af")
        return
 
    if score is not None:
        color = _score_color(float(score))
        ax.barh([0], [float(score)], color=color, height=0.5, alpha=0.92)
        ax.set_xlim(0, 1.05)
        ax.set_yticks([0])
        ax.set_yticklabels(["score"], fontsize=9, color="#d1d5db")
        ax.set_xlabel("Consistency (0–1)", color="#9ca3af", fontsize=8)
        ax.text(min(float(score) + 0.02, 0.98), 0, f"{float(score):.3f}", va="center", fontsize=8, color="#9ca3af")
        if query:
            ax.set_title(f"Consistency\n{query}", color="#e5e7eb", fontsize=9)
    elif variants:
        labels = [f"v{i}" for i in range(len(variants))]
        kls = [float(v.get("kl_from_anchor") or 0) for v in variants]
        ax.barh(labels, kls, color="#6366f1", height=0.65, alpha=0.9)
        ax.invert_yaxis()
        ax.set_xlabel("KL from anchor", color="#9ca3af", fontsize=8)
 
    ax.tick_params(colors="#6b7280", labelsize=8)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
 
 
def _plot_suppression_panel(ax: Any, data: dict[str, Any]) -> None:
    ax.set_title("Suppression", color="#e5e7eb", fontsize=10)
    topics = data.get("topics") or []
 
    if not topics:
        ax.axis("off")
        ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#9ca3af")
        return
 
    rows = sorted(topics, key=lambda t: float(t.get("suppression_score") or 0), reverse=True)[:8]
    labels = [str(t.get("topic") or "?") for t in rows]
    vals = [float(t.get("suppression_score") or 0) for t in rows]
    colors = [_STATUS_COLORS.get(str(t.get("status") or ""), "#78716c") for t 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("Suppression score", 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:.2f}", va="center", fontsize=7, color="#9ca3af")
 
 
def _plot_boundary_panel(ax: Any, data: dict[str, Any]) -> None:
    mean_rob = data.get("mean_robustness")
    probes = data.get("probes") or []
    ax.set_title("Boundary / robustness", color="#e5e7eb", fontsize=10)
 
    if not probes and mean_rob is None:
        ax.axis("off")
        ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#9ca3af")
        return
 
    if len(probes) <= 1 and mean_rob is not None:
        color = _score_color(float(mean_rob))
        ax.barh([0], [float(mean_rob)], color=color, height=0.5, alpha=0.92)
        ax.set_xlim(0, 1.05)
        ax.set_yticks([0])
        ax.set_yticklabels(["mean"], fontsize=9, color="#d1d5db")
        ax.set_xlabel("Robustness (0–1)", color="#9ca3af", fontsize=8)
        ax.text(min(float(mean_rob) + 0.02, 0.98), 0, f"{float(mean_rob):.3f}", va="center", fontsize=8, color="#9ca3af")
        ax.tick_params(colors="#6b7280", labelsize=8)
        ax.grid(axis="x", alpha=0.2, color="#4b5563")
        return
 
    labels = [_trunc(str(p.get("prompt") or ""), 18) for p in probes[:8]]
    vals = [float(p.get("robustness_score") or 0) for p in probes[:8]]
    colors = [_score_color(v) for v in vals]
    y_pos = list(range(len(labels)))
    ax.barh(y_pos, vals, color=colors, height=0.65, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=7, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Robustness", color="#9ca3af", fontsize=8)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=7)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
 
 
def _score_color(val: float) -> str:
    if val >= 0.75:
        return "#34d399"
    if val >= 0.5:
        return "#fbbf24"
    return "#f87171"
 
 
def _trunc(text: str, limit: int) -> str:
    t = text.strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"