stubs.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.
 
"""
Step 35: every tool in tools.ts wired to bridge.call via execute_tool.
 
Tools already wired with full handlers (not re-registered here):
  inspect.py      — run_full_inspection, get_feature_logits,
                    get_feature_neighbors
  session_memory.py — write_session_memory, read_session_memory
"""
from __future__ import annotations
from typing import Any
from .registry import register
from .executor import execute_tool
 
 
def _wire(name: str, schema: dict[str, Any]) -> None:
    @register(name, schema)
    def _handler(args: dict[str, Any], ctx: dict[str, Any]) -> Any:
        return execute_tool(name, args, ctx)
    _handler.__name__ = name
 
 
# ── Inspection ──
 
_wire("run_find_feature", {
    "type": "function",
    "function": {
        "name": "run_find_feature",
        "description": (
            "Rank SAE features that separate honest vs deceptive probes. "
            "Default behavior mode generates completions and buckets by observed output. "
            "Returns top features by activation delta; optionally persists canonical feature id."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "scorer": {"type": "string", "description": "Scorer name (default: deception)."},
                "prompts": {"type": "string", "description": "Path to JSON/JSONL honest/deceptive probes."},
                "layer": {"type": "number", "description": "SAE layer (default: model default)."},
                "checkpoint": {"type": "string", "description": "Optional fine-tuned checkpoint .pt path."},
                "top_k": {"type": "number", "description": "Number of ranked features to return (default 20)."},
                "direction": {
                    "type": "string",
                    "description": "both (|delta|, default), deceptive (delta>0 only), honest (delta<0 only).",
                },
                "conditioning": {
                    "type": "string",
                    "description": "behavior (default): generate + classify output; prompt: static probe text only.",
                },
                "benchmark_top": {"type": "number", "description": "Re-rank top K with InterpScore (needs OpenAI)."},
                "persist": {"type": "string", "description": "Session/experiment key for chosen feature id."},
            },
            "required": [],
        },
    },
})
 
_wire("run_benchmarks_on_top_feature", {
    "type": "function",
    "function": {
        "name": "run_benchmarks_on_top_feature",
        "description": (
            "Feature benchmark: InterpScore, FeaturePurity, and MUI on one SAE feature. "
            "Use the same SAE layer as the inspection that surfaced the feature."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "feature_idx": {"type": "number"},
                "layer": {
                    "type": "number",
                    "description": "SAE layer (defaults to lastSaeLayer from session memory).",
                },
                "prompt": {
                    "type": "string",
                    "description": "Probe prompt for labeling (defaults to last inspected prompt).",
                },
            },
            "required": ["feature_idx"],
        },
    },
})
 
_wire("run_audit", {
    "type": "function",
    "function": {
        "name": "run_audit",
        "description": "Run a full audit: fact-check, bias detection, and censor audit in parallel.",
        "parameters": {"type": "object", "properties": {}, "required": []},
    },
})
 
_wire("run_consistency_eval", {
    "type": "function",
    "function": {
        "name": "run_consistency_eval",
        "description": "Measure output stability across paraphrases.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "templates": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["query", "templates"],
        },
    },
})
 
_wire("run_suppression_eval", {
    "type": "function",
    "function": {
        "name": "run_suppression_eval",
        "description": "Detect topics where the model's behavior differs from neutral baseline.",
        "parameters": {
            "type": "object",
            "properties": {
                "topics": {"type": "object", "additionalProperties": {"type": "array", "items": {"type": "string"}}},
            },
            "required": ["topics"],
        },
    },
})
 
_wire("run_boundary_eval", {
    "type": "function",
    "function": {
        "name": "run_boundary_eval",
        "description": "Map robustness to surface-level input corruptions.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompts": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["prompts"],
        },
    },
})
 
_wire("run_weight_diff", {
    "type": "function",
    "function": {
        "name": "run_weight_diff",
        "description": (
            "Per-layer weight delta analysis: base vs fine-tuned checkpoint. "
            "Q/K/V/O/MLP ΔW or LoRA effective B@A."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "checkpoint": {"type": "string", "description": "Path to fine-tuned .pt or HF directory."},
                "name": {"type": "string", "description": "Checkpoint label for the report."},
                "save": {"type": "string", "description": "Optional JSON export path."},
                "model_id": {"type": "string"},
            },
            "required": ["checkpoint"],
        },
    },
})
 
_wire("run_merge_analysis", {
    "type": "function",
    "function": {
        "name": "run_merge_analysis",
        "description": (
            "Pre-merge LoRA gate: weight-diff + rank/collapse signals + optional behavioral model-diff (LLM). "
            "Returns mergeVerdict pass|warn|fail."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "checkpoint": {"type": "string", "description": "Path to adapter or merged .pt checkpoint."},
                "name": {"type": "string", "description": "Checkpoint label for the report."},
                "prompts": {"type": "string", "description": "Optional JSON/JSONL probes for behavioral diff."},
                "no_behavioral": {"type": "boolean", "description": "Skip generation-based behavioral scores."},
                "save": {"type": "string", "description": "Optional JSON export path."},
                "model_id": {"type": "string"},
            },
            "required": ["checkpoint"],
        },
    },
})
 
_wire("run_trajectory_analysis", {
    "type": "function",
    "function": {
        "name": "run_trajectory_analysis",
        "description": (
            "Training trajectory: weight-diff summary per checkpoint vs base, sorted by training step."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "checkpoints": {"type": "string", "description": "Glob of .pt checkpoint paths."},
                "dir": {"type": "string", "description": "Directory to scan recursively for *.pt."},
                "name": {"type": "string", "description": "Optional label prefix per step."},
                "save": {"type": "string", "description": "Optional JSON export path."},
                "model_id": {"type": "string"},
            },
        },
    },
})
 
_wire("run_residual_drift", {
    "type": "function",
    "function": {
        "name": "run_residual_drift",
        "description": (
            "Per-layer activation drift: catalog base vs fine-tuned checkpoint on the same probes. "
            "Last-token hook_resid_post cosine distance."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "checkpoint": {"type": "string", "description": "Path to fine-tuned .pt or HF directory."},
                "prompts": {"type": "string", "description": "Optional JSON/JSONL probe file."},
                "name": {"type": "string", "description": "Checkpoint label for the report."},
                "save": {"type": "string", "description": "Optional JSON export path."},
                "model_id": {"type": "string"},
            },
            "required": ["checkpoint"],
        },
    },
})
 
_wire("run_confidence_analysis", {
    "type": "function",
    "function": {
        "name": "run_confidence_analysis",
        "description": (
            "Per-probe confidence analysis over probes: token logits (mean confidence, entropy, ECE proxy). "
            "Optional SAE join for confidence ↔ feature ↔ layer on the same prompts."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "prompts": {"type": "string", "description": "Path to JSON/JSONL probe file."},
                "threshold": {"type": "number", "description": "Low-confidence cutoff (default 0.40)."},
                "join_sae": {"type": "boolean", "description": "Attach SAE mean L0 + top feature per probe."},
                "layer": {"type": "number", "description": "SAE layer for join (default model sae_layer)."},
                "save": {"type": "string", "description": "Optional JSON export path."},
                "model_id": {"type": "string"},
            },
            "required": ["prompts"],
        },
    },
})
 
_wire("extract_steer_vector", {
    "type": "function",
    "function": {
        "name": "extract_steer_vector",
        "description": "Export an SAE decoder steering vector (LAT) to a JSON file for reuse across sessions and checkpoints.",
        "parameters": {
            "type": "object",
            "properties": {
                "feature_idx": {"type": "number", "description": "SAE feature index to export."},
                "layer": {"type": "number", "description": "SAE layer (default: model default)."},
                "save": {"type": "string", "description": "Output JSON path."},
                "feature_label": {"type": "string", "description": "Optional human-readable label stored in the file."},
                "probe_id": {"type": "string", "description": "Optional experiment key from feature locate --persist (default: deception_feature)."},
            },
            "required": ["feature_idx", "save"],
        },
    },
})
 
_wire("run_steer_and_show", {
    "type": "function",
    "function": {
        "name": "run_steer_and_show",
        "description": "Run feature steering AND open the steer panel with the result injected. Use feature_idx or a saved LAT vector file. Pass eval=true to score baseline vs steered on a probe suite (default: deception probes).",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "Optional single-prompt demo. Omit with eval=true to run the probe suite only."},
                "feature_idx": {"type": "number", "description": "SAE feature index (omit when using vector)."},
                "vector": {"type": "string", "description": "Path to a saved LAT JSON from steer --save."},
                "feature_label": {"type": "string", "description": "Optional — auto-resolved from causal labeling if omitted"},
                "strength": {"type": "number"},
                "layer": {"type": "number"},
                "max_new_tokens": {"type": "number"},
                "eval": {"type": "boolean", "description": "If true, score the same probe set under baseline vs steered and report pass rates."},
                "prompts": {"type": "string", "description": "Probe file (json/jsonl) or JSON array."},
                "reference_answers": {"type": "string", "description": "For custom eval mode: JSON array of references matching prompts (keyword-overlap pass)."},
                "threshold": {"type": "number", "description": "Custom-mode pass threshold (default 0.5)."},
            },
            "required": [],
        },
    },
})
 
_wire("run_multi_steer", {
    "type": "function",
    "function": {
        "name": "run_multi_steer",
        "description": "Steer multiple SAE features simultaneously in a single forward pass.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string"},
                "features": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "feature_idx": {"type": "number"},
                            "strength": {"type": "number"},
                            "label": {"type": "string", "description": "Optional — auto-resolved if omitted"},
                        },
                        "required": ["feature_idx", "strength"],
                    },
                },
                "max_new_tokens": {"type": "number"},
            },
            "required": ["prompt", "features"],
        },
    },
})
 
_wire("ensure_umap_loaded", {
    "type": "function",
    "function": {
        "name": "ensure_umap_loaded",
        "description": "Load the UMAP projection and open the UMAP Explorer panel. CLI equivalent: add --umap to a feature command (e.g. aquin feature locate --umap).",
        "parameters": {"type": "object", "properties": {}, "required": []},
    },
})
 
_wire("run_layer_analysis", {
    "type": "function",
    "function": {
        "name": "run_layer_analysis",
        "description": (
            "Run activation stability and OOD similarity together. CLI: aquin check layer. "
            "Pass localize=true to also rank layers where honest vs deceptive representations "
            "peak and collapse (optional stressor_prompts / feature_idx / vector)."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "prompts": {
                    "type": "string",
                    "description": "Stability prompts (JSON array) or, with localize, path to deception probe file.",
                },
                "top_k": {"type": "number"},
                "in_domain_prompts": {"type": "array", "items": {"type": "string"}},
                "ood_prompts": {"type": "array", "items": {"type": "string"}},
                "localize": {
                    "type": "boolean",
                    "description": "If true, rank layers by honest vs deceptive representation strength and report collapse locus.",
                },
                "stressor_prompts": {
                    "type": "string",
                    "description": "Optional second deception probe file for stressor collapse deltas.",
                },
                "feature_idx": {
                    "type": "number",
                    "description": "Optional SAE feature to project at its layer.",
                },
                "vector": {
                    "type": "string",
                    "description": "Optional saved LAT vector path for direction projection.",
                },
                "layer": {"type": "number", "description": "Override SAE/vector layer."},
            },
            "required": [],
        },
    },
})
 
_wire("run_sae_stats", {
    "type": "function",
    "function": {
        "name": "run_sae_stats",
        "description": (
            "Export multi-layer SAE statistics over a probe dataset. "
            "Returns per-layer sparsity, top features, and probe×layer heatmap data."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "prompts": {
                    "type": "string",
                    "description": "Path to JSON/JSONL probe file (text + optional stressor/lang/quant_run_id tags).",
                },
                "layers": {
                    "type": "string",
                    "description": "Comma-separated layer indices or 'all' (default: all layers with SAE checkpoints).",
                },
                "top_k": {"type": "number", "description": "Top features per layer (default 10)."},
                "save": {"type": "string", "description": "Optional path to write full JSON export."},
                "model_id": {"type": "string", "description": "Override active model."},
            },
            "required": ["prompts"],
        },
    },
})
 
_wire("run_perturbation_sensitivity", {
    "type": "function",
    "function": {
        "name": "run_perturbation_sensitivity",
        "description": "Zero out hidden channels and measure KL divergence vs clean output. CLI: aquin check perturbation.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string"},
                "n_channels": {"type": "number"},
                "method": {"type": "string", "enum": ["dropout", "gaussian"]},
            },
            "required": ["prompt"],
        },
    },
})
 
_wire("run_attention_routing", {
    "type": "function",
    "function": {
        "name": "run_attention_routing",
        "description": "Extract per-head attention patterns and score each head. CLI: aquin check attention.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string"},
                "top_k": {"type": "number"},
            },
            "required": ["prompt"],
        },
    },
})
 
# ── Training: dataset ──
 
_wire("dataset_generate", {
    "type": "function",
    "function": {
        "name": "dataset_generate",
        "description": "Generate instruction/response pairs for a topic and write a .jsonl file in the working directory.",
        "parameters": {
            "type": "object",
            "properties": {
                "topic": {"type": "string"},
                "count": {"type": "number"},
                "save": {"type": "string", "description": "Filename in cwd (default: <topic>_dataset.jsonl)"},
            },
            "required": ["topic"],
        },
    },
})
 
# ── Training: algo run ──
 
_LORA_PARAMS = {
    "rank": {"type": "number"},
    "alpha": {"type": "number"},
    "lr": {"type": "number"},
    "epochs": {"type": "number"},
    "dropout": {"type": "number"},
    "targetModules": {"type": "array", "items": {"type": "string"}},
    "warmupSteps": {"type": "number"},
    "gradClip": {"type": "number"},
    "weightDecay": {"type": "number"},
    "gradAccumSteps": {"type": "number"},
    "optimizer": {"type": "string", "enum": ["adamw", "sgd", "lion"]},
    "scheduler": {"type": "string", "enum": ["cosine", "linear", "constant"]},
    "maxSeqLen": {"type": "number"},
    "useQlora": {"type": "boolean"},
    "rows": {"type": "array", "items": {"type": "object"}},
}
 
_wire("run_simulation", {
    "type": "function",
    "function": {
        "name": "run_simulation",
        "description": "Simulate the user's existing training setup. Pass paths to their dataset and pipeline (script or config) — no Aquin-specific formats.",
        "parameters": {
            "type": "object",
            "properties": {
                "topic": {"type": "string", "description": "Quick probe when no dataset path yet."},
                "dataset": {"type": "string", "description": "Path to the user's dataset file."},
                "algo": {"type": "string", "description": "Path to the user's training script or config."},
                "count": {"type": "number"},
                **{k: v for k, v in _LORA_PARAMS.items() if k != "rows"},
                "use_rlhf": {"type": "boolean"},
                "rlhf_beta": {"type": "number"},
            },
            "required": [],
        },
    },
})
 
_wire("list_simulation_runs", {
    "type": "function",
    "function": {
        "name": "list_simulation_runs",
        "description": "List the user's saved simulation runs.",
        "parameters": {"type": "object", "properties": {}, "required": []},
    },
})
 
_wire("load_simulation_run", {
    "type": "function",
    "function": {
        "name": "load_simulation_run",
        "description": "Load a saved simulation run by ID. CLI: aquin replay simulation --run_id <id>.",
        "parameters": {
            "type": "object",
            "properties": {
                "run_id": {"type": "string"},
            },
            "required": ["run_id"],
        },
    },
})
 
_wire("compare_simulations", {
    "type": "function",
    "function": {
        "name": "compare_simulations",
        "description": "Compare two saved simulation runs side-by-side, including attack-surface metrics (consistency, suppression, robustness) from model diff.",
        "parameters": {
            "type": "object",
            "properties": {
                "run_id_a": {"type": "string", "description": "First simulation run ID (before)."},
                "run_id_b": {"type": "string", "description": "Second simulation run ID (after)."},
                "run_id_before": {"type": "string", "description": "Alias for run_id_a."},
                "run_id_after": {"type": "string", "description": "Alias for run_id_b."},
                "label_a": {"type": "string"},
                "label_b": {"type": "string"},
            },
            "required": ["run_id_a", "run_id_b"],
        },
    },
})
 
# ── Interactive cards ──
 
_wire("run_red_team", {
    "type": "function",
    "function": {
        "name": "run_red_team",
        "description": "Full adversarial red-team suite (all six vectors). Always runs every vector — do not ask the user to pick vectors. call with {}.",
        "parameters": {
            "type": "object",
            "properties": {}, "required": [],
        },
    },
})
 
_wire("check_weights", {
    "type": "function",
    "function": {
        "name": "check_weights",
        "description": "Run a combined weight health check: trojan/backdoor signatures and SVD rank analysis. CLI: aquin check weight.",
        "parameters": {
            "type": "object",
            "properties": {
                "collapse_threshold": {"type": "number"},
            },
            "required": [],
        },
    },
})
 
_wire("run_custom_eval", {
    "type": "function",
    "function": {
        "name": "run_custom_eval",
        "description": "Run prompts through the model and score each response against a reference answer (keyword overlap). Prefer built-in eval tools over inventing a code runner.",
        "parameters": {
            "type": "object",
            "properties": {
                "name": {"type": "string"},
                "description": {"type": "string"},
                "prompts": {"type": "array", "items": {"type": "string"}},
                "reference_answers": {"type": "array", "items": {"type": "string"}},
                "threshold": {"type": "number"},
                "max_tokens": {"type": "number"},
                "temperature": {"type": "number"},
            },
            "required": ["name", "prompts", "reference_answers"],
        },
    },
})
 
 
# ── Code execution removed (never shipped) ──