inspect.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.
 
from __future__ import annotations
 
from typing import Any
 
from .registry import register
 
 
def _resolve_inspection_sae_layer(
    model_id: str,
    layer_raw: Any,
    *,
    memory: dict[str, Any] | None = None,
    command: str = "trace",
) -> int:
    """Resolve SAE layer from explicit arg or session memory; never silently default."""
    from aquin.compute.model_loader import require_sae_layer
 
    layer: int | None = int(layer_raw) if layer_raw is not None else None
    if layer is None and memory:
        mem_layer = memory.get("lastSaeLayer")
        if mem_layer is not None:
            layer = int(mem_layer)
    return require_sae_layer(model_id, layer, command=command)
 
 
def _build_color_map(attribution: list[dict]) -> dict | None:
    feature_color: dict[int, int] = {}
    counter = [0]
 
    def _color(fidx: int) -> int:
        if fidx not in feature_color:
            feature_color[fidx] = counter[0] % 10
            counter[0] += 1
        return feature_color[fidx]
 
    prompt_cm: dict[int, int] = {}
    response_cm: dict[int, int] = {}
    for entry in attribution:
        ri = entry["response_ti"]
        for feat in entry.get("driven_by_features", []):
            cidx = _color(feat["feature_idx"])
            response_cm[ri] = cidx
            for pi in feat.get("also_in_prompt_positions", []):
                prompt_cm[pi] = cidx
 
    if not prompt_cm and not response_cm:
        return None
    return {"prompt": prompt_cm, "response": response_cm}
 
 
def _build_trace_results(logit_lens: list[dict]) -> list[dict]:
    if not logit_lens:
        return []
    results = []
    final_prob = logit_lens[-1]["top_tokens"][0]["prob"] if logit_lens[-1].get("top_tokens") else 0.0
    for i, row in enumerate(logit_lens):
        prob = row["top_tokens"][0]["prob"] if row.get("top_tokens") else 0.0
        next_row = logit_lens[i + 1] if i + 1 < len(logit_lens) else None
        next_prob = next_row["top_tokens"][0]["prob"] if next_row and next_row.get("top_tokens") else prob
        delta = max(0.0, next_prob - prob)
        results.append({
            "layer": row["layer"],
            "drop": round(delta, 4),
            "attn_drop": round(delta * 0.6, 4),
            "mlp_drop": round(delta * 0.4, 4),
            "baseline_prob": round(final_prob, 4),
        })
    return results
 
 
@register("run_full_inspection", {
    "type": "function",
    "function": {
        "name": "run_full_inspection",
        "description": (
            "Run full inspection: generate a response, extract top SAE features at the given layer, "
            "compute attribution, and run logit lens. layer must match a downloaded SAE "
            "(aquin load sae <model>-l<n>). Omit layer to reuse lastSaeLayer from session memory. "
            "Errors with available layers if the checkpoint is missing."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "The prompt to inspect"},
                "layer": {
                    "type": "integer",
                    "description": (
                        "SAE layer index (aquin load sae <model>-l<n>). "
                        "Optional if lastSaeLayer is in session memory."
                    ),
                },
            },
            "required": ["prompt"],
        },
    },
})
def run_full_inspection(args: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import get_active_model_id, get_loaded_model, load_model, require_sae_layer, resolve_model_id
    from aquin.compute.causal_trace import run_chat, run_logit_lens
    from aquin.compute.feature_analysis import run_feature_analysis_unlabeled
    from aquin.engine.session_memory_store import load_local_memory, save_local_memory
 
    session_id: str = ctx.get("session_id") or ""
    state: dict[str, Any] = ctx.get("state", {})
 
    prompt: str = args.get("prompt", "") or state.get("lastPrompt") or "Hello"
    model_id: str = get_active_model_id() or state.get("activeModelId") or "llama-3.2-1b"
    layer_raw = args.get("layer")
    layer: int | None = int(layer_raw) if layer_raw is not None else None
 
    try:
        model_id = resolve_model_id(model_id)
    except Exception:
        return {"error": f"Unknown model '{model_id}'. Run: aquin load --model <model-id>"}
 
    try:
        sae_layer = _resolve_inspection_sae_layer(
            model_id,
            layer_raw,
            memory=state.get("memory") or {},
            command="trace",
        )
    except ValueError as e:
        return {"error": str(e)}
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": f"Could not load model: {e}"}
 
    # 1. Generate response
    try:
        response = run_chat(prompt, model_id=model_id, max_new_tokens=200, temperature=0.7)
    except Exception as e:
        return {"error": f"Generation failed: {e}"}
 
    # 2. SAE feature analysis
    try:
        feat_result = run_feature_analysis_unlabeled(
            prompt, response, model, model_id=model_id, layer=sae_layer,
        )
    except Exception as e:
        return {"error": f"Feature analysis failed: {e}"}
 
    actual_layer = feat_result.get("sae_layer")
    if actual_layer is not None and int(actual_layer) != int(sae_layer):
        return {
            "error": (
                f"SAE layer mismatch: requested layer {sae_layer}, "
                f"feature analysis ran at layer {actual_layer}."
            ),
        }
 
    prompt_tokens: list[str] = feat_result.get("prompt_tokens", [])
    response_tokens: list[str] = feat_result.get("response_tokens", [])
    top_features: list[dict] = feat_result.get("top_response_features", [])
    attribution: list[dict] = feat_result.get("attribution", [])
    sae_layer: int = feat_result.get("sae_layer", sae_layer)
 
    # Causal labels for displayed features (cached; uses last prompt context)
    try:
        from aquin.compute.feature_analysis import label_inspection_features
        from aquin.compute.openai_client import get_openai_client
        label_inspection_features(
            feat_result,
            prompt=prompt,
            model=model,
            client=get_openai_client(ctx),
            model_id=model_id,
            layer=sae_layer,
        )
        top_features = feat_result.get("top_response_features", [])
        attribution = feat_result.get("attribution", [])
    except Exception as e:
        print(f"[inspect] feature labeling failed: {e}", flush=True)
 
    # 3. Logit lens
    try:
        logit_lens = run_logit_lens(prompt, model_id=model_id, top_k=5)
    except Exception as e:
        print(f"[inspect] logit-lens failed: {e}", flush=True)
        logit_lens = []
 
    color_map = _build_color_map(attribution)
    trace_results = _build_trace_results(logit_lens)
    trace_target = response_tokens[0].strip() if response_tokens else ""
 
    panel_card: dict[str, Any] = {
        "type": "inspectionFull",
        "data": {
            "prompt": prompt,
            "response": response,
            "modelId": model_id,
            "topFeatures": top_features,
            "promptTokens": prompt_tokens,
            "responseTokens": response_tokens,
            "attribution": attribution,
            "logitLensResults": logit_lens,
            "traceResults": trace_results,
            "traceTarget": trace_target,
            "colorMap": color_map,
            "saeLayer": sae_layer,
        },
    }
 
    mem = load_local_memory(session_id or "local")
    mem.update({
        "lastPrompt": prompt,
        "lastResponse": response,
        "lastTopFeatures": top_features[:5],
        "lastSaeLayer": sae_layer,
    })
    save_local_memory(session_id or "local", mem)
    state.setdefault("memory", {})
    state["memory"].update(mem)
 
    return {
        "content": {
            "prompt": prompt,
            "response": response,
            "model_id": model_id,
            "top_features": top_features,
            "sae_layer": sae_layer,
        },
        "card": panel_card,
    }
 
 
@register("get_feature_logits", {
    "type": "function",
    "function": {
        "name": "get_feature_logits",
        "description": (
            "Get top tokens boosted and suppressed by a specific SAE feature. "
            "Use the same layer as the last inspection (or pass layer explicitly)."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "feature_idx": {"type": "number", "description": "SAE feature index"},
                "layer": {
                    "type": "number",
                    "description": "SAE layer (defaults to lastSaeLayer from session memory).",
                },
                "top_k": {"type": "number", "description": "Number of tokens to return (default 10)."},
            },
            "required": ["feature_idx"],
        },
    },
})
def get_feature_logits(args: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import (
        get_loaded_llm_id,
        get_loaded_model,
        load_model,
        resolve_model_id,
    )
    from aquin.compute.feature_analysis import get_feature_logits as _get_logits
    from aquin.compute.model_runtime import resident_from_cache
 
    state: dict[str, Any] = ctx.get("state", {}) if isinstance(ctx.get("state"), dict) else {}
    resident, _ = resident_from_cache()
    model_id = (
        state.get("activeModelId")
        or args.get("model_id")
        or get_loaded_llm_id()
        or resident
    )
    if not model_id:
        return {
            "error": "No model loaded. Ask the user to load one via the Model picker.",
        }
    feature_idx: int = int(args.get("feature_idx", 0))
    layer_raw = args.get("layer")
    try:
        top_k = int(args.get("top_k") or 10)
    except (TypeError, ValueError):
        top_k = 10
 
    try:
        model_id = resolve_model_id(str(model_id))
    except Exception:
        return {"error": f"Unknown model '{model_id}'"}
 
    try:
        sae_layer = _resolve_inspection_sae_layer(
            model_id,
            layer_raw,
            memory=state.get("memory") or {},
            command="feature logit",
        )
    except ValueError as e:
        return {"error": str(e)}
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
 
    try:
        from aquin.compute.feature_analysis import enrich_feature_tool_result, prompt_for_labeling
        from aquin.compute.openai_client import get_openai_client
        result = _get_logits(feature_idx, model, model_id=model_id, layer=sae_layer, top_k=top_k)
        enrich_feature_tool_result(
            result,
            prompt=prompt_for_labeling(ctx, args),
            model=model,
            client=get_openai_client(ctx),
            model_id=model_id,
            layer=sae_layer,
        )
        return {"content": result}
    except Exception as e:
        from aquin.user_errors import friendly_message
        return {"error": friendly_message(e)}
 
 
@register("get_feature_neighbors", {
    "type": "function",
    "function": {
        "name": "get_feature_neighbors",
        "description": (
            "Find SAE features nearest to a given feature by cosine similarity. "
            "Use the same layer as the last inspection (or pass layer explicitly)."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "feature_idx": {"type": "number", "description": "SAE feature index"},
                "layer": {
                    "type": "number",
                    "description": "SAE layer (defaults to lastSaeLayer from session memory).",
                },
                "top_k": {"type": "number", "description": "Number of neighbors (default 8)."},
            },
            "required": ["feature_idx"],
        },
    },
})
def get_feature_neighbors(args: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import (
        get_loaded_llm_id,
        get_loaded_model,
        load_model,
        resolve_model_id,
    )
    from aquin.compute.feature_analysis import (
        enrich_feature_tool_result,
        get_feature_neighbors as _get_neighbors,
        prompt_for_labeling,
    )
    from aquin.compute.openai_client import get_openai_client
    from aquin.compute.model_runtime import resident_from_cache
 
    state: dict[str, Any] = ctx.get("state", {}) if isinstance(ctx.get("state"), dict) else {}
    resident, _ = resident_from_cache()
    model_id = (
        state.get("activeModelId")
        or args.get("model_id")
        or get_loaded_llm_id()
        or resident
    )
    if not model_id:
        return {
            "error": "No model loaded. Ask the user to load one via the Model picker.",
        }
    feature_idx: int = int(args.get("feature_idx", 0))
    layer_raw = args.get("layer")
    try:
        top_k = int(args.get("top_k") or 8)
    except (TypeError, ValueError):
        top_k = 8
 
    try:
        model_id = resolve_model_id(str(model_id))
    except Exception:
        return {"error": f"Unknown model '{model_id}'"}
 
    try:
        sae_layer = _resolve_inspection_sae_layer(
            model_id,
            layer_raw,
            memory=state.get("memory") or {},
            command="feature neighbor",
        )
    except ValueError as e:
        return {"error": str(e)}
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
 
    try:
        result = _get_neighbors(feature_idx, model_id=model_id, layer=sae_layer, top_k=top_k)
        enrich_feature_tool_result(
            result,
            prompt=prompt_for_labeling(ctx, args),
            model=model,
            client=get_openai_client(ctx),
            model_id=model_id,
            layer=sae_layer,
            label_neighbors=True,
        )
        return {"content": result}
    except Exception as e:
        from aquin.user_errors import friendly_message
        return {"error": friendly_message(e)}