card_mapper.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""
Maps tool result dicts to PanelCardData-shaped dicts for the web UI.
Returns None for tools that produce no card (UI-only events, state mutations, etc.).
"""
from __future__ import annotations
 
# Tools that produce no card — UI events and state-only mutations
_NO_CARD_TOOLS: frozenset[str] = frozenset({
    # session memory
    "write_session_memory",
    "read_session_memory",
})
 
# Direct type mappings — tool_name → card type string
_TYPE_MAP: dict[str, str] = {
    "run_full_inspection": "inspectionFull",
    "run_benchmarks_on_top_feature": "interpScore",
    "run_consistency_eval": "evals",
    "run_suppression_eval": "evals",
    "run_boundary_eval": "evals",
    "run_audit": "evals",
    "get_feature_logits": "featureLogits",
    "get_feature_neighbors": "featureNeighbors",
    "run_steer_and_show": "steer",
    "extract_steer_vector": "steerVector",
    "run_multi_steer": "steer",
    "check_weights": "trojan",
    "ensure_umap_loaded": "umap",
    "run_layer_analysis": "layerAnalysis",
    "run_perturbation_sensitivity": "perturbation",
    "run_attention_routing": "attentionRouting",
    "run_simulation": "simulationFull",
    "list_simulation_runs": "simulationList",
    "load_simulation_run": "simulationFull",
    "compare_simulations": "simulationComparison",
    "run_custom_eval": "customEval",
    "run_red_team": "redTeam",
    "run_sae_diff": "saeDiff",
    "run_find_feature": "findFeature",
    "run_sae_stats": "saeStats",
    "run_weight_diff": "weightDiff",
    "run_merge_analysis": "mergeAnalysis",
    "run_trajectory_analysis": "trajectoryAnalysis",
    "run_residual_drift": "residualDrift",
    "run_confidence_analysis": "confidenceAnalysis",
}
 
 
def _simulation_method_from_meta(meta: dict) -> str:
    algo = meta.get("algoConfig") if isinstance(meta.get("algoConfig"), dict) else {}
    title = str(algo.get("title") or "")
    for name in ("QLoRA", "Full FT", "SFT", "RLHF", "DPO", "Distil", "CPT", "LoRA"):
        if name in title:
            return name
    return "LoRA"
 
 
def _simulation_config_from_meta(meta: dict) -> dict:
    config: dict = {"lr": 2e-4, "epochs": 3, "optimizer": "adamw"}
    algo = meta.get("algoConfig") if isinstance(meta.get("algoConfig"), dict) else {}
    rows = algo.get("config")
    if not isinstance(rows, list):
        return config
    for item in rows:
        if not isinstance(item, dict):
            continue
        key = str(item.get("k", "")).lower()
        val = item.get("v")
        if key == "lr":
            try:
                config["lr"] = float(str(val).replace("e", "E"))
            except (TypeError, ValueError):
                pass
        elif key == "rank":
            try:
                config["rank"] = int(val)
            except (TypeError, ValueError):
                pass
        elif key == "alpha":
            try:
                config["alpha"] = int(val)
            except (TypeError, ValueError):
                pass
        elif key == "epochs" or "epoch" in key:
            try:
                config["epochs"] = int(str(val).split()[0])
            except (TypeError, ValueError):
                pass
    return config
 
 
def simulation_list_card_data(result: dict) -> dict:
    """Map list_simulation_runs → SimulationListResult for the web panel."""
    raw = result.get("runs") or []
    runs: list[dict] = []
    for entry in raw:
        if isinstance(entry, str):
            runs.append({"run_id": entry})
        elif isinstance(entry, dict):
            runs.append({
                "run_id": str(entry.get("run_id") or ""),
                "saved_at": entry.get("saved_at"),
                "model_id": entry.get("model_id"),
                "topic": entry.get("topic"),
                "dataset_path": entry.get("dataset_path"),
                "kind": entry.get("kind"),
                "n_samples": entry.get("n_samples"),
                "harmful_count": entry.get("harmful_count"),
                "method": entry.get("method"),
            })
    return {
        "runs": runs,
        "count": int(result.get("count") if result.get("count") is not None else len(runs)),
    }
 
 
def steer_card_data(result: dict) -> dict | None:
    """Map run_steer_and_show / run_multi_steer → steer panel card fields."""
    baseline = result.get("baseline") or result.get("baseline_response") or ""
    steered = result.get("steered") or result.get("steered_response") or ""
    if not baseline and not steered:
        return None
 
    feature_idx = int(result.get("feature_idx", 0))
    feature_label = str(result.get("feature_label") or result.get("feature_ref") or "")
    features = result.get("features")
    if not feature_label and isinstance(features, list) and features:
        parts: list[str] = []
        for f in features:
            if not isinstance(f, dict):
                continue
            parts.append(str(f.get("label") or f.get("feature_ref") or f"F{f.get('feature_idx', '?')}"))
            if feature_idx == 0 and f.get("feature_idx") is not None:
                feature_idx = int(f["feature_idx"])
        feature_label = ", ".join(parts) if parts else f"F{feature_idx}"
 
    prompt = str(result.get("prompt") or "")
    words_changed = result.get("words_changed")
    if words_changed is None and baseline and steered:
        b_words = baseline.split()
        s_words = steered.split()
        words_changed = sum(1 for a, b in zip(b_words, s_words) if a != b) + abs(len(b_words) - len(s_words))
 
    data: dict = {
        "featureIdx": feature_idx,
        "featureLabel": feature_label,
        "prompt": prompt,
        "baseline": str(baseline),
        "steered": str(steered),
    }
    if words_changed is not None:
        data["wordsChanged"] = int(words_changed)
    if result.get("vector_path"):
        data["vectorPath"] = str(result["vector_path"])
    if result.get("vector_source_model_id"):
        data["vectorSourceModelId"] = str(result["vector_source_model_id"])
    return data
 
 
def steer_vector_card_data(result: dict) -> dict | None:
    if result.get("error"):
        return None
    path = result.get("output_path")
    if not path:
        return None
    return {
        "modelId": result.get("model_id"),
        "layer": result.get("layer"),
        "featureIdx": result.get("feature_idx"),
        "featureLabel": result.get("feature_label"),
        "probeId": result.get("probe_id"),
        "outputPath": path,
        "dModel": result.get("d_model"),
        "vectorL2Norm": result.get("vector_l2_norm"),
        "norm": result.get("norm"),
        "status": result.get("status", "done"),
    }
 
 
def load_simulation_card_data(result: dict) -> dict:
    """Map load_simulation_run → SimulationFullResult."""
    if result.get("status") == "not_found" or result.get("error"):
        return {}
    run_id = str(result.get("run_id") or "")
    payload = {k: v for k, v in result.items() if k not in ("run_id", "error", "status")}
    wrapped = {
        "model_id": payload.get("model_id") or (payload.get("meta") or {}).get("modelId"),
        "run_id": run_id,
        "result": payload,
        "events": [],
    }
    data = simulation_full_card_data(wrapped)
    data["savedRunId"] = run_id
    return data
 
 
 
 
 
 
def simulation_comparison_card_data(result: dict) -> dict:
    """Map compare_simulations → SimulationComparisonResult."""
    comp = result.get("comparison")
    if isinstance(comp, dict) and comp.get("featureDiffs") is not None:
        return comp
    return {}
 
 
def simulation_full_card_data(result: dict) -> dict:
    """Map CLI simulate result → SimulationFullResult for the web panel."""
    assembled = result.get("result") if isinstance(result.get("result"), dict) else {}
    events = result.get("events") if isinstance(result.get("events"), list) else []
 
    meta = dict(assembled.get("meta") or {})
    loss_history: list[float] = list(assembled.get("lossHistory") or [])
    grad_heatmap: list[dict] = []
    signals: list[dict] = list(assembled.get("signals") or [])
    logs: list[str] = []
 
    for ev in events:
        if not isinstance(ev, dict):
            continue
        t = ev.get("type")
        if t == "meta":
            meta.update({k: v for k, v in ev.items() if k != "type"})
        elif t == "step" and ev.get("loss") is not None:
            loss_history.append(float(ev["loss"]))
        elif t == "gradHeatmap":
            grad_heatmap.append({k: v for k, v in ev.items() if k != "type"})
        elif t == "signal":
            signals.append({
                "signalType": ev.get("signalType", ""),
                "severity": ev.get("severity", ""),
                "message": ev.get("message", ""),
                "step": ev.get("step", 0),
            })
        elif t == "log" and ev.get("line"):
            logs.append(str(ev["line"]))
 
    model_id = str(result.get("model_id") or meta.get("modelId") or "")
    if not model_id and meta.get("modelClass"):
        model_id = str(meta["modelClass"])
 
    model_diff = assembled.get("modelDiff")
    if isinstance(model_diff, dict):
        model_diff = {**model_diff, "isSimulation": model_diff.get("isSimulation", True)}
 
    return {
        "modelId": model_id,
        "method": _simulation_method_from_meta(meta),
        "modelClass": meta.get("modelClass") or (model_id.split("/")[-1] if model_id else "model"),
        "trainableParams": meta.get("trainableParams"),
        "config": _simulation_config_from_meta(meta),
        "datasetQuality": assembled.get("datasetQuality"),
        "saePrediction": assembled.get("saePrediction"),
        "influenceScores": assembled.get("influenceScores"),
        "rlhfPrediction": assembled.get("rlhfPrediction"),
        "effectiveLR": assembled.get("effectiveLR"),
        "lossSharpness": assembled.get("lossSharpness"),
        "modelDiff": model_diff,
        "gradHeatmap": grad_heatmap or assembled.get("gradHeatmap"),
        "lossHistory": loss_history,
        "signals": signals,
        "logs": logs,
        "calibration": assembled.get("calibration"),
        "savedRunId": result.get("run_id"),
        "isSimulation": True,
    }
 
 
def to_card(tool_name: str, result: dict) -> dict | None:
    """
    Convert a tool result dict into a PanelCardData-shaped dict.
    Returns None for tools that don't produce a visual card.
    """
    if tool_name in _NO_CARD_TOOLS:
        return None
 
    # If the result is already a UI event (from UI-only pass-through), skip
    if result.get("type") == "ui_event":
        return None
 
    # If the result signals not-implemented, no card
    if result.get("status") == "not_implemented":
        return None
 
    card_type = _TYPE_MAP.get(tool_name)
    if card_type is None:
        # Unknown tool or one that returns its own card structure
        return None
 
    # Special case: tools that wrap data under a sub-key
    if tool_name == "run_full_inspection":
        data = result.get("content", result)
        return {"type": card_type, "data": data}
 
    if tool_name == "get_feature_logits":
        data = result.get("content", result)
        if isinstance(data, dict):
            boosts = data.get("boosts") or data.get("top") or []
            suppresses = data.get("suppresses") or data.get("bottom") or []
            data = {**data, "boosts": boosts, "suppresses": suppresses}
        return {"type": card_type, "data": data}
 
    if tool_name == "get_feature_neighbors":
        data = result.get("content", result)
        return {"type": card_type, "data": data}
 
    if tool_name == "check_weights":
        if result.get("error"):
            return None
        src = result.get("trojan", result)
        if not isinstance(src, dict):
            return None
        data = {
            "model_id": str(result.get("model_id", src.get("model_id", ""))),
            "generated_at": str(result.get("generated_at", "")),
            "layers_analysed": src.get("layers_analysed", 0),
            "composite_risk": src.get("composite_risk", 0),
            "verdict": src.get("verdict", "clean"),
            "pct_flagged": src.get("pct_flagged", 0),
            "high_risk_count": src.get("high_risk_count", 0),
            "suspicious_count": src.get("suspicious_count", 0),
            "clean_count": src.get("clean_count", 0),
            "all_flags": src.get("all_flags", []),
            "scored_tensors": src.get("scored_tensors", []),
            "signals": src.get("signals", []),
        }
        return {"type": card_type, "data": data}
 
    if tool_name == "run_audit":
        return {
            "type": card_type,
            "data": {
                "consistency": result.get("consistency"),
                "suppression": result.get("suppression"),
                "boundary": result.get("boundary"),
            },
        }
 
    if tool_name == "run_consistency_eval":
        return {"type": card_type, "data": {"consistency": result}}
    if tool_name == "run_suppression_eval":
        return {"type": card_type, "data": {"suppression": result}}
    if tool_name == "run_boundary_eval":
        return {"type": card_type, "data": {"boundary": result}}
 
    if tool_name in ("run_steer_and_show", "run_multi_steer"):
        if result.get("error"):
            return None
        data = steer_card_data(result)
        if data is None:
            return None
        return {"type": card_type, "data": data}
 
    if tool_name == "extract_steer_vector":
        if result.get("error"):
            return None
        data = steer_vector_card_data(result)
        if data is None:
            return None
        return {"type": card_type, "data": data}
 
    if tool_name == "run_benchmarks_on_top_feature":
        if result.get("error"):
            return None
        payload = result
        if result.get("type") == "benchmark" and isinstance(result.get("data"), dict):
            payload = result["data"]
        return {"type": card_type, "data": payload}
 
    if tool_name == "run_find_feature":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "modelId": result.get("model_id"),
                "layer": result.get("layer"),
                "scorer": result.get("scorer"),
                "direction": result.get("direction"),
                "conditioning": result.get("conditioning"),
                "behavior": result.get("behavior"),
                "nHonest": result.get("n_honest"),
                "nDeceptive": result.get("n_deceptive"),
                "promptsPath": result.get("prompts_path"),
                "checkpoint": result.get("checkpoint"),
                "chosenFeatureIdx": result.get("chosen_feature_idx"),
                "chosenDelta": result.get("chosen_delta"),
                "warning": result.get("warning"),
                "persistedKey": result.get("persisted_key"),
                "experimentPath": result.get("experiment_path"),
                "rankings": result.get("rankings") or [],
                "status": result.get("status", "done"),
            },
        }
 
    if tool_name == "run_layer_analysis":
        return {
            "type": card_type,
            "data": {
                "stability": result.get("stability"),
                "ood": result.get("ood"),
                "localize": result.get("localize"),
            },
        }
 
    if tool_name == "run_sae_stats":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "modelId": result.get("model_id"),
                "mode": result.get("mode"),
                "nProbes": result.get("n_probes"),
                "topK": result.get("top_k"),
                "layersRequested": result.get("layers_requested") or [],
                "layerProfile": result.get("layer_profile") or [],
                "layerStats": result.get("layer_stats") or [],
                "heatmap": result.get("heatmap") or {},
                "probes": result.get("probes") or [],
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "run_confidence_analysis":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "modelId": result.get("model_id"),
                "mode": result.get("mode"),
                "nProbes": result.get("n_probes"),
                "threshold": result.get("threshold"),
                "meanConfidence": result.get("mean_confidence"),
                "aggregateEceProxy": result.get("aggregate_ece_proxy"),
                "lowConfidenceCount": result.get("low_confidence_count"),
                "joinSae": result.get("join_sae"),
                "saeLayer": result.get("sae_layer"),
                "stressorSummary": result.get("stressor_summary") or [],
                "heatmap": result.get("heatmap") or {},
                "probes": result.get("probes") or [],
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "run_weight_diff":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "baseModelId": result.get("baseModelId"),
                "ftCheckpointName": result.get("ftCheckpointName"),
                "checkpointPath": result.get("checkpointPath"),
                "modelMode": result.get("modelMode"),
                "deltaMode": result.get("deltaMode"),
                "trainingStep": result.get("trainingStep"),
                "nMatrices": result.get("nMatrices"),
                "totalDeltaL2": result.get("totalDeltaL2"),
                "maxDeltaL2": result.get("maxDeltaL2"),
                "meanDeltaStableRank": result.get("meanDeltaStableRank"),
                "layerProfile": result.get("layerProfile") or [],
                "topChanged": result.get("topChanged") or [],
                "matrices": result.get("matrices") or [],
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "run_merge_analysis":
        if result.get("error"):
            return None
        behavioral = result.get("behavioralDiff")
        return {
            "type": card_type,
            "data": {
                "baseModelId": result.get("baseModelId"),
                "ftCheckpointName": result.get("ftCheckpointName"),
                "checkpointPath": result.get("checkpointPath"),
                "modelMode": result.get("modelMode"),
                "deltaMode": result.get("deltaMode"),
                "trainingStep": result.get("trainingStep"),
                "mergeVerdict": result.get("mergeVerdict"),
                "warnings": result.get("warnings") or [],
                "nMatrices": result.get("nMatrices"),
                "totalDeltaL2": result.get("totalDeltaL2"),
                "maxDeltaL2": result.get("maxDeltaL2"),
                "meanDeltaStableRank": result.get("meanDeltaStableRank"),
                "collapseSignals": result.get("collapseSignals") or [],
                "topChanged": result.get("topChanged") or [],
                "layerProfile": result.get("layerProfile") or [],
                "withBehavioral": result.get("withBehavioral"),
                "behavioralDiff": behavioral,
                "behavioralError": result.get("behavioralError"),
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "run_trajectory_analysis":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "baseModelId": result.get("baseModelId"),
                "nCheckpoints": result.get("nCheckpoints"),
                "nAnalyzed": result.get("nAnalyzed"),
                "peakStep": result.get("peakStep"),
                "peakTotalDeltaL2": result.get("peakTotalDeltaL2"),
                "steps": result.get("steps") or [],
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "run_residual_drift":
        if result.get("error"):
            return None
        return {
            "type": card_type,
            "data": {
                "baseModelId": result.get("baseModelId"),
                "ftCheckpointName": result.get("ftCheckpointName"),
                "checkpointPath": result.get("checkpointPath"),
                "modelMode": result.get("modelMode"),
                "activationMode": result.get("activationMode"),
                "trainingStep": result.get("trainingStep"),
                "nProbes": result.get("nProbes"),
                "nLayers": result.get("nLayers"),
                "meanDrift": result.get("meanDrift"),
                "maxDrift": result.get("maxDrift"),
                "peakLayer": result.get("peakLayer"),
                "layerProfile": result.get("layerProfile") or [],
                "topLayers": result.get("topLayers") or [],
                "perProbe": result.get("perProbe") or [],
                "savedTo": result.get("saved_to"),
            },
        }
 
    if tool_name == "ensure_umap_loaded":
        points = result.get("points") or []
        return {
            "type": card_type,
            "data": {
                "modelId": str(result.get("model_id", "")),
                "nFeatures": int(result.get("n_features") or len(points)),
                "nPoints": int(result.get("n_points") or len(points)),
                "points": points,
            },
        }
 
    if tool_name == "run_simulation":
        if result.get("error"):
            return None
        return {"type": card_type, "data": simulation_full_card_data(result)}
 
        if result.get("error"):
            return None
        summary = result.get("result")
        if not isinstance(summary, dict) or summary.get("meanApSim") is None:
            return None
        return {"type": card_type, "data": summary}
 
    if tool_name == "list_simulation_runs":
        return {"type": card_type, "data": simulation_list_card_data(result)}
 
    if tool_name == "load_simulation_run":
        if result.get("status") == "not_found" or result.get("error"):
            return None
        data = load_simulation_card_data(result)
        if not data:
            return None
        return {"type": card_type, "data": data}
 
    if tool_name == "compare_simulations":
        if result.get("error"):
            return None
        kind = result.get("kind")
        comp = result.get("comparison") if isinstance(result.get("comparison"), dict) else {}
        data = simulation_comparison_card_data(result)
        if not data:
            return None
        return {"type": card_type, "data": data}
 
    if tool_name == "dataset_generate":
        if result.get("error"):
            return None
        return {"type": "datasetGenerated", "data": result}
 
        anisotropy = result.get("anisotropy")
        intrinsic = result.get("intrinsicDim") or result.get("intrinsic_dim")
        model_id = result.get("model_id")
        if isinstance(anisotropy, dict) and model_id and not anisotropy.get("model_id"):
            anisotropy = {**anisotropy, "model_id": model_id}
        if isinstance(intrinsic, dict) and model_id and not intrinsic.get("model_id"):
            intrinsic = {**intrinsic, "model_id": model_id}
        if not anisotropy and not intrinsic:
            return None
        return {
            "type": card_type,
            "data": {
                "anisotropy": anisotropy,
                "intrinsicDim": intrinsic,
            },
        }
 
    # Default: wrap full result as data
    return {"type": card_type, "data": result}