trace_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin trace --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 = "trace-check.json"
_PNG_NAME = "trace-check.png"
 
 
def _flatten_trace_result(result: dict[str, Any]) -> dict[str, Any]:
    if result.get("error"):
        return dict(result)
 
    content = result.get("content") or {}
    card_data = (result.get("card") or {}).get("data") or {}
 
    def pick(snake: str, camel: str, default: Any = None) -> Any:
        val = content.get(snake)
        if val is None or val == "" or val == []:
            val = card_data.get(camel, default)
        return default if val is None else val
 
    flat = {
        "prompt": pick("prompt", "prompt", ""),
        "response": pick("response", "response", ""),
        "model_id": pick("model_id", "modelId", ""),
        "sae_layer": pick("sae_layer", "saeLayer"),
        "top_features": pick("top_features", "topFeatures", []),
        "prompt_tokens": card_data.get("promptTokens", content.get("prompt_tokens", [])),
        "response_tokens": card_data.get("responseTokens", content.get("response_tokens", [])),
        "attribution": card_data.get("attribution", content.get("attribution", [])),
        "logit_lens": card_data.get("logitLensResults", content.get("logit_lens", [])),
        "trace_results": card_data.get("traceResults", content.get("trace_results", [])),
        "trace_target": card_data.get("traceTarget", content.get("trace_target", "")),
        "color_map": card_data.get("colorMap", content.get("color_map")),
    }
    if result.get("error"):
        flat["error"] = result["error"]
    return flat
 
 
def write_trace_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 = _flatten_trace_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_trace_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_trace_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 "")
    prompt = str(result.get("prompt") or "")[:64]
    response = str(result.get("response") or "")[:80]
    sae_layer = result.get("sae_layer", "?")
 
    fig = plt.figure(figsize=(12, 7), facecolor="#0f1117")
    gs = fig.add_gridspec(2, 2, hspace=0.38, wspace=0.28)
    ax_feat = fig.add_subplot(gs[0, 0])
    ax_lens = fig.add_subplot(gs[0, 1])
    ax_trace = fig.add_subplot(gs[1, 0])
    ax_text = fig.add_subplot(gs[1, 1])
 
    fig.suptitle(f"Full trace — {model}  ·  SAE layer {sae_layer}", color="#e5e7eb", fontsize=11, y=0.98)
 
    _plot_top_features(ax_feat, result.get("top_features") or [])
    _plot_logit_lens(ax_lens, result.get("logit_lens") or [], result.get("trace_target") or "")
    _plot_trace(ax_trace, result.get("trace_results") or [])
    _plot_summary(ax_text, prompt, response, result.get("attribution") or [])
 
    for ax in (ax_feat, ax_lens, ax_trace):
        ax.set_facecolor("#0f1117")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    ax_text.set_facecolor("#0f1117")
    for spine in ax_text.spines.values():
        spine.set_color("#374151")
 
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_top_features(ax: Any, features: list[dict[str, Any]]) -> None:
    rows = sorted(features, key=lambda f: float(f.get("activation") or 0), reverse=True)[:12]
    if not rows:
        ax.axis("off")
        ax.text(0.5, 0.5, "No top features", ha="center", va="center", color="#9ca3af")
        ax.set_title("Top SAE features", color="#e5e7eb", fontsize=10)
        return
 
    labels = []
    vals = []
    for f in rows:
        ref = f.get("feature_ref") or f.get("feature_idx", "?")
        label = f.get("causal_label") or f.get("label") or ""
        tok = f.get("token") or ""
        suffix = f" · {tok!r}" if tok else (f" · {label[:18]}" if label else "")
        labels.append(f"{ref}{suffix}"[:28])
        vals.append(float(f.get("activation") or 0))
 
    y_pos = list(range(len(rows)))
    ax.barh(y_pos, vals, color="#6366f1", height=0.72, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=7, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlabel("Activation", color="#9ca3af", fontsize=8)
    ax.set_title("Top SAE features", color="#e5e7eb", fontsize=10)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=7)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
 
 
def _plot_logit_lens(ax: Any, logit_lens: list[dict[str, Any]], trace_target: str) -> None:
    if not logit_lens:
        ax.axis("off")
        ax.text(0.5, 0.5, "No logit lens data", ha="center", va="center", color="#9ca3af")
        ax.set_title("Logit lens", color="#e5e7eb", fontsize=10)
        return
 
    layers = []
    probs = []
    top_token = trace_target or ""
    for row in logit_lens:
        layer = int(row.get("layer", len(layers)))
        tops = row.get("top_tokens") or []
        if not tops:
            continue
        best = tops[0]
        layers.append(layer)
        probs.append(float(best.get("prob") or 0))
        if not top_token:
            top_token = str(best.get("token") or "")
 
    if not layers:
        ax.axis("off")
        ax.text(0.5, 0.5, "No logit lens data", ha="center", va="center", color="#9ca3af")
        ax.set_title("Logit lens", color="#e5e7eb", fontsize=10)
        return
 
    ax.plot(layers, probs, color="#34d399", marker="o", markersize=3, linewidth=1.5)
    ax.set_xlabel("Layer", color="#9ca3af", fontsize=8)
    ax.set_ylabel("P(top token)", color="#9ca3af", fontsize=8)
    title = f'Logit lens — top token {top_token!r}' if top_token else "Logit lens"
    ax.set_title(title[:48], color="#e5e7eb", fontsize=9)
    ax.tick_params(colors="#6b7280", labelsize=7)
    ax.grid(alpha=0.2, color="#4b5563")
 
 
def _plot_trace(ax: Any, trace_results: list[dict[str, Any]]) -> None:
    if not trace_results:
        ax.axis("off")
        ax.text(0.5, 0.5, "No trace data", ha="center", va="center", color="#9ca3af")
        ax.set_title("Layer probability delta", color="#e5e7eb", fontsize=10)
        return
 
    layers = [int(r.get("layer", i)) for i, r in enumerate(trace_results)]
    drops = [float(r.get("drop") or 0) for r in trace_results]
    ax.bar(layers, drops, color="#fbbf24", width=0.75, alpha=0.85)
    ax.set_xlabel("Layer", color="#9ca3af", fontsize=8)
    ax.set_ylabel("Δ prob", color="#9ca3af", fontsize=8)
    ax.set_title("Layer probability delta", color="#e5e7eb", fontsize=10)
    ax.tick_params(colors="#6b7280", labelsize=7)
    ax.grid(axis="y", alpha=0.2, color="#4b5563")
 
 
def _plot_summary(ax: Any, prompt: str, response: str, attribution: list[dict[str, Any]]) -> None:
    ax.axis("off")
    ax.set_title("Summary", color="#e5e7eb", fontsize=10, loc="left")
 
    lines = [
        f'Prompt: "{prompt}"',
        "",
        f'Response: "{response}"',
        "",
        f"Attribution: {len(attribution)} response tokens with driving features",
    ]
    if attribution:
        lines.append("")
        for a in attribution[:4]:
            tok = a.get("response_token", "")
            driven = a.get("driven_by_features") or []
            refs = ", ".join(
                str(d.get("feature_ref") or d.get("feature_idx", "?"))
                for d in driven[:3]
            )
            lines.append(f"  {tok!r} ← {refs}")
 
    ax.text(
        0.02,
        0.96,
        "\n".join(lines),
        transform=ax.transAxes,
        va="top",
        ha="left",
        fontsize=8,
        color="#d1d5db",
        family="monospace",
        wrap=True,
    )