session_mode.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.
 
"""LLM session mode — controls which tools exist in a CLI session."""
from __future__ import annotations
 
from typing import Literal
 
SessionMode = Literal["llm"]
 
# ── Tool registry names ───────────────────────────────────────────────────────
 
SHARED_TOOLS: frozenset[str] = frozenset({
    "write_session_memory",
    "read_session_memory",
    "run_custom_eval",
    "run_sae_stats",
})
 
# Tools that never touch the loaded LLM weights — skip ensure_model in dispatch.
MODEL_FREE_TOOLS: frozenset[str] = frozenset({
    "read_session_memory",
    "write_session_memory",
    "list_simulation_runs",
    "load_simulation_run",
    "compare_simulations",
})
 
 
def tool_requires_model(tool_name: str) -> bool:
    return tool_name not in MODEL_FREE_TOOLS
 
LLM_TOOLS: frozenset[str] = frozenset({
    "run_full_inspection",
    "run_benchmarks_on_top_feature",
    "run_find_feature",
    "run_audit",
    "run_consistency_eval",
    "run_suppression_eval",
    "run_boundary_eval",
    "get_feature_logits",
    "get_feature_neighbors",
    "run_steer_and_show",
    "extract_steer_vector",
    "run_multi_steer",
    "ensure_umap_loaded",
    "run_layer_analysis",
    "run_confidence_analysis",
    "run_weight_diff",
    "run_merge_analysis",
    "run_trajectory_analysis",
    "run_residual_drift",
    "run_sae_stats",
    "run_perturbation_sensitivity",
    "run_attention_routing",
    "dataset_generate",
    "run_simulation",
    "list_simulation_runs",
    "load_simulation_run",
    "compare_simulations",
    "run_red_team",
    "check_weights",
})
 
# CLI-only shortcuts (not in tool registry)
LLM_CLI_ONLY: frozenset[str] = frozenset({"trace", "feature logit", "feature neighbor"})
 
# ── CLI verb → tool registry ─────────────────────────────────────────────────
 
SHARED_CLI_VERBS: dict[str, str] = {
    "feature locate": "run_find_feature",
    "feature logit": "get_feature_logits",
    "feature neighbor": "get_feature_neighbors",
    "sae-stats": "run_sae_stats",
    "diff weight": "run_merge_analysis",
    "check trajectory": "run_trajectory_analysis",
    "diff residue": "run_residual_drift",
    "confidence-analysis": "run_confidence_analysis",
    "eval": "run_custom_eval",
    "list simulation": "list_simulation_runs",
    "replay simulation": "load_simulation_run",
    "compare simulation": "compare_simulations",
    "list-runs": "list_simulation_runs",
    "load-run": "load_simulation_run",
    "compare-runs": "compare_simulations",
}
 
LLM_CLI_VERBS: dict[str, str] = {
    "benchmark": "run_benchmarks_on_top_feature",
    "benchmarks": "run_benchmarks_on_top_feature",
    "consistency-eval": "run_consistency_eval",
    "suppression-eval": "run_suppression_eval",
    "boundary-eval": "run_boundary_eval",
    "steer": "run_steer_and_show",
    "multi-steer": "run_multi_steer",
    "layer-analysis": "run_layer_analysis",
    "perturbation": "run_perturbation_sensitivity",
    "attention": "run_attention_routing",
    "check-weights": "check_weights",
    "simulate": "run_simulation",
    "red-team": "run_red_team",
}
 
 
def resolve_legacy_verb(verb: str) -> str:
    return verb
 
 
def cli_verbs_for_mode(mode: SessionMode | None) -> dict[str, str]:
    """CLI verb map for the active session mode."""
    return {**SHARED_CLI_VERBS, **LLM_CLI_VERBS}
 
 
def verb_known(verb: str) -> bool:
    verb = resolve_legacy_verb(verb)
    return verb in SHARED_CLI_VERBS or verb in LLM_CLI_VERBS
 
 
def normalize_mode(raw: str | None) -> SessionMode:
    return "llm"
 
 
def tools_for_mode(mode: SessionMode) -> frozenset[str]:
    return LLM_TOOLS | SHARED_TOOLS
 
 
def is_tool_allowed(tool_name: str, mode: SessionMode) -> bool:
    return tool_name in tools_for_mode(mode)
 
 
def filter_cli_tool_map(full_map: dict[str, str], mode: SessionMode) -> dict[str, str]:
    allowed = tools_for_mode(mode)
    return {verb: tool for verb, tool in full_map.items() if tool in allowed}
 
 
def get_filtered_tool_schemas(mode: SessionMode) -> list[dict]:
    from aquin.engine.tools.registry import get_tool_schemas
 
    allowed = tools_for_mode(mode)
    return [s for s in get_tool_schemas() if s.get("function", {}).get("name") in allowed]
 
 
def mode_label(mode: SessionMode) -> str:
    return "LLM"
 
 
def mode_for_model_id(model_id: str) -> SessionMode:
    """Validate model id and return LLM session mode."""
    from aquin.compute.model_loader import resolve_model_id
 
    resolve_model_id(model_id)  # raises if unknown LLM
    return "llm"
 
 
def mode_for_active_model() -> SessionMode | None:
    """Infer mode from the model last loaded via `aquin load` (~/.aquin/active_model.txt)."""
    from aquin.compute.model_loader import get_active_model_id
 
    model_id = get_active_model_id()
    if not model_id:
        return None
    try:
        from aquin.compute.model_loader import resolve_model_id
        resolve_model_id(model_id)
        return "llm"
    except ValueError:
        return None