command_log.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Local command tracking — ~/.aquin/commands/"""
from __future__ import annotations
 
import json
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_COMMANDS_DIR = Path.home() / ".aquin" / "commands"
_INDEX_PATH = _COMMANDS_DIR / "index.json"
 
_CATEGORY_BY_TOOL: dict[str, str] = {
    "run_full_inspection": "inspection",
    "get_feature_logits": "inspection",
    "get_feature_neighbors": "inspection",
    "run_attention_routing": "inspection",
    "run_layer_analysis": "inspection",
    "run_perturbation_sensitivity": "inspection",
    "check_weights": "inspection",
    "run_simulation": "simulation",
    "run_sae_stats": "sae",
    "run_find_feature": "sae",
    "run_confidence_analysis": "sae",
    "run_sae_diff": "sae",
    "run_sae_align": "sae",
    "run_sae_train": "sae",
    "run_capture_activations": "capture",
    "run_weight_diff": "checkpoint",
    "run_merge_analysis": "checkpoint",
    "run_trajectory_analysis": "checkpoint",
    "run_residual_drift": "checkpoint",
    "ensure_umap_loaded": "inspection",
    "run_steer_and_show": "sae",
    "run_multi_steer": "sae",
    "extract_steer_vector": "sae",
}
 
_TOOL_TO_CLI: dict[str, str] = {
    "run_full_inspection": "trace",
    "get_feature_logits": "feature logit",
    "get_feature_neighbors": "feature neighbor",
    "run_attention_routing": "check attention",
    "run_layer_analysis": "check layer",
    "run_perturbation_sensitivity": "check perturbation",
    "check_weights": "check weight",
    "run_simulation": "simulate",
    "run_merge_analysis": "diff weight",
    "run_trajectory_analysis": "check trajectory",
    "run_residual_drift": "diff residue",
    "run_sae_stats": "sae-stats",
    "run_find_feature": "feature locate",
    "run_confidence_analysis": "check confidence",
    "run_capture_activations": "activations capture",
    "run_custom_eval": "eval custom",
    "run_consistency_eval": "eval consistency",
    "run_suppression_eval": "eval suppress",
    "run_boundary_eval": "eval boundary",
}
 
 
def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
 
 
def _ensure_dir() -> None:
    _COMMANDS_DIR.mkdir(parents=True, exist_ok=True)
 
 
def _load_index() -> list[dict[str, Any]]:
    _ensure_dir()
    if not _INDEX_PATH.exists():
        return []
    try:
        data = json.loads(_INDEX_PATH.read_text())
        return data if isinstance(data, list) else []
    except Exception:
        return []
 
 
def _save_index(rows: list[dict[str, Any]]) -> None:
    _ensure_dir()
    _INDEX_PATH.write_text(json.dumps(rows, indent=2))
 
 
def category_for_tool(tool_name: str) -> str:
    return _CATEGORY_BY_TOOL.get(tool_name, "other")
 
 
def cli_name_for_tool(tool_name: str) -> str:
    return _TOOL_TO_CLI.get(tool_name, tool_name.replace("run_", "").replace("_", "-"))
 
 
def _build_summary(
    *,
    command: str,
    tool_name: str,
    model_id: str | None,
    result: Any,
    card: dict[str, Any] | None,
) -> str:
    parts = [command]
    if model_id:
        parts.append(model_id)
    hint = ""
    if isinstance(result, dict):
        if result.get("error"):
            hint = "error"
        elif isinstance(result.get("top_features"), list) and result["top_features"]:
            f0 = result["top_features"][0]
            if isinstance(f0, dict) and "feature_idx" in f0:
                hint = f"F{f0['feature_idx']}"
        elif result.get("mergeVerdict"):
            hint = str(result["mergeVerdict"])
        elif result.get("nChanged") is not None:
            hint = f"{result['nChanged']} changed"
    elif card and isinstance(card.get("type"), str):
        hint = card["type"]
    if hint:
        parts.append(hint)
    return " · ".join(parts)
 
 
def append_record(
    *,
    tool_name: str,
    command: str | None = None,
    args: dict[str, Any] | None = None,
    result: Any = None,
    card: dict[str, Any] | None = None,
    model_id: str | None = None,
    duration_ms: int | None = None,
    exit_ok: bool = True,
) -> str:
    """Persist a command run locally. Returns record id."""
    record_id = str(uuid.uuid4())
    cmd = command or cli_name_for_tool(tool_name)
    category = category_for_tool(tool_name)
 
    if not model_id:
        try:
            from aquin.compute.model_loader import get_active_model_id
 
            model_id = get_active_model_id() or None
        except Exception:
            model_id = None
 
    record: dict[str, Any] = {
        "id": record_id,
        "ts": _now_iso(),
        "command": cmd,
        "tool_name": tool_name,
        "category": category,
        "model_id": model_id,
        "args": args or {},
        "result": result,
        "card": card,
        "summary": _build_summary(
            command=cmd,
            tool_name=tool_name,
            model_id=model_id,
            result=result,
            card=card,
        ),
        "duration_ms": duration_ms,
        "exit_ok": exit_ok,
    }
 
    _ensure_dir()
    (_COMMANDS_DIR / f"{record_id}.json").write_text(json.dumps(record, indent=2, default=str))
 
    index = _load_index()
    index.insert(0, {
        "id": record_id,
        "ts": record["ts"],
        "command": cmd,
        "tool_name": tool_name,
        "category": category,
        "model_id": model_id,
        "summary": record["summary"],
        "exit_ok": exit_ok,
    })
    _save_index(index[:5000])
    return record_id
 
 
def load_record(record_id: str) -> dict[str, Any] | None:
    path = _COMMANDS_DIR / f"{record_id}.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text())
    except Exception:
        return None
 
 
def _parse_since(value: str) -> datetime | None:
    value = value.strip().lower()
    m = re.fullmatch(r"(\d+)(m|h|d)", value)
    if m:
        n, unit = int(m.group(1)), m.group(2)
        from datetime import timedelta
 
        delta = {"m": timedelta(minutes=n), "h": timedelta(hours=n), "d": timedelta(days=n)}[unit]
        return datetime.now(timezone.utc) - delta
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except Exception:
        return None
 
 
def list_records(
    *,
    last: int | None = None,
    since: str | None = None,
    category: str | None = None,
    model_id: str | None = None,
    command: str | None = None,
) -> list[dict[str, Any]]:
    rows = _load_index()
    since_dt = _parse_since(since) if since else None
    out: list[dict[str, Any]] = []
    for row in rows:
        if category and row.get("category") != category:
            continue
        if model_id and row.get("model_id") != model_id:
            continue
        if command and row.get("command") != command:
            continue
        if since_dt:
            try:
                ts = datetime.fromisoformat(str(row.get("ts", "")).replace("Z", "+00:00"))
                if ts < since_dt:
                    continue
            except Exception:
                pass
        out.append(row)
        if last is not None and len(out) >= last:
            break
    return out
 
 
def clear_records(*, before: str | None = None) -> int:
    rows = _load_index()
    if not before:
        count = len(rows)
        for row in rows:
            rid = row.get("id")
            if rid:
                p = _COMMANDS_DIR / f"{rid}.json"
                if p.exists():
                    p.unlink()
        _save_index([])
        return count
 
    cutoff = _parse_since(before)
    if cutoff is None:
        return 0
 
    kept: list[dict[str, Any]] = []
    removed = 0
    for row in rows:
        try:
            ts = datetime.fromisoformat(str(row.get("ts", "")).replace("Z", "+00:00"))
        except Exception:
            kept.append(row)
            continue
        if ts < cutoff:
            rid = row.get("id")
            if rid:
                p = _COMMANDS_DIR / f"{rid}.json"
                if p.exists():
                    p.unlink()
            removed += 1
        else:
            kept.append(row)
    _save_index(kept)
    return removed