status_display.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin status` — one bullet per fact, nothing else."""
 
from __future__ import annotations
 
import json
from dataclasses import dataclass, field
from typing import Any
 
 
@dataclass
class StatusSnapshot:
    logged_in: bool = False
    account: str | None = None
    active_id: str | None = None
    token_masked: str | None = None
    other_accounts: list[str] = field(default_factory=list)
    model_id: str | None = None
    mode_label: str | None = None
    device: str | None = None
    location_label: str | None = None
    gpu_line: str | None = None
    engine_line: str = ""
    license_line: str = ""
    live_note: str | None = None
 
 
def collect_status_snapshot(
    *,
    account: str | None = None,
    active_id: str | None = None,
    token_masked: str | None = None,
    other_accounts: list[str] | None = None,
    live_note: str | None = None,
    logged_in: bool = True,
) -> StatusSnapshot:
    """Gather local state only (no network). Auth is not required."""
    _ = account, active_id, token_masked, other_accounts, live_note, logged_in
 
    from aquin.compute.model_loader import get_active_model_id
    from aquin.engine.main import _load_state
    from aquin.license_acceptance import is_license_accepted
 
    state = _load_state()
    model_id = get_active_model_id(allow_daemon=False)
    mode_label_str = None
    if model_id:
        try:
            from aquin.session_mode import mode_for_model_id, mode_label as mode_label_fn
 
            mode_label_str = mode_label_fn(mode_for_model_id(model_id))
        except Exception:
            pass
 
    loc = state.get("location") or {}
    location_label = (
        loc.get("location_label") if isinstance(loc, dict) else None
    )
 
    return StatusSnapshot(
        logged_in=True,
        account=None,
        active_id=None,
        token_masked=None,
        other_accounts=[],
        model_id=model_id,
        mode_label=mode_label_str,
        device=state.get("device"),
        location_label=location_label,
        gpu_line=_gpu_line_from_state(state),
        engine_line=_engine_line_fast(),
        license_line="accepted" if is_license_accepted() else "not accepted",
        live_note=None,
    )
 
 
def _status_rows(snapshot: StatusSnapshot) -> list[tuple[str, str]]:
    """Small set of rows for the default human status output."""
    rows: list[tuple[str, str]] = []
 
    if snapshot.model_id:
        model = snapshot.model_id
        if snapshot.mode_label:
            model = f"{model} ({snapshot.mode_label})"
        rows.append(("model", model))
    else:
        rows.append(("model", "none - aquin load model <id>"))
 
    rows.append(("engine", snapshot.engine_line or "off"))
    rows.append(("gpu", snapshot.gpu_line or "unknown"))
    if snapshot.live_note:
        rows.append(("note", snapshot.live_note))
    rows.append(("license", snapshot.license_line or "?"))
    return rows
 
 
def status_payload(snapshot: StatusSnapshot) -> dict[str, Any]:
    """Machine-readable status for Ink / scripts (`aquin status --json`)."""
    return {
        "loggedIn": bool(snapshot.logged_in),
        "rows": [{"label": label, "value": value} for label, value in _status_rows(snapshot)],
    }
 
def render_status(snapshot: StatusSnapshot, *, plain: bool = False) -> None:
    """Print a small static status summary."""
    _ = plain
    rows = _status_rows(snapshot)
    label_w = max((len(label) for label, _ in rows), default=8)
    for label, value in rows:
        print(f"{label.ljust(label_w)}  {value}")
 
 
# Back-compat name
def render_status_ascii(snapshot: StatusSnapshot) -> None:
    render_status(snapshot, plain=True)
 
 
def collect_and_render_status(
    *,
    account: str | None,
    active_id: str | None,
    token_masked: str | None,
    other_accounts: list[str],
    live_note: str | None = None,
) -> None:
    snap = collect_status_snapshot(
        account=account,
        active_id=active_id,
        token_masked=token_masked,
        other_accounts=other_accounts,
        live_note=live_note,
        logged_in=True,
    )
    render_status(snap)
 
 
def _gpu_line_from_state(state: dict[str, Any]) -> str | None:
    gpu_info = state.get("gpu_info")
    parsed: dict[str, Any] | None = None
    if gpu_info:
        try:
            parsed = json.loads(gpu_info) if isinstance(gpu_info, str) else gpu_info
        except Exception:
            parsed = None
 
    if parsed:
        gpus = parsed.get("gpus") or []
        backend = parsed.get("backend") or parsed.get("selected_device")
        if gpus:
            g = gpus[0]
            name = g.get("name", "?")
            vram = g.get("vram_gb", "?")
            line = f"{name} ({vram} GB)"
            if backend and backend not in ("cuda",):
                line += f" - {backend}"
            return line
        if parsed.get("mps_available") or parsed.get("selected_device") == "mps":
            return "Apple Metal (MPS)"
        if parsed.get("accelerator_available") is False:
            return "no GPU accelerator"
 
    try:
        from aquin.compute.device import probe_backend
 
        probe = probe_backend()
        if probe.get("gpus"):
            g = probe["gpus"][0]
            return f"{g.get('name', '?')} ({g.get('vram_gb', '?')} GB) - {probe.get('backend')}"
        if probe.get("selected") == "mps":
            return "Apple Metal (MPS)"
        if not probe.get("accelerator_available"):
            return "no GPU accelerator"
        return probe.get("summary")
    except Exception:
        return None
 
 
def _engine_line_fast() -> str:
    try:
        from aquin.engine import model_daemon
 
        h = model_daemon.health_fast()
        if h:
            resident = h.get("model_id") or "none"
            status = h.get("model_status") or "idle"
            if status == "loading":
                target = h.get("loading_model_id") or resident
                elapsed = h.get("elapsed_s")
                bit = f"loading: {target}"
                if elapsed is not None:
                    bit = f"{bit} ({elapsed}s)"
                return bit
            vram = h.get("vram")
            if resident == "none":
                return "on, no model"
            line = f"ready: {resident}"
            if vram and resident != "none":
                line = f"{line} ({vram})"
            return line
    except Exception:
        pass
    return "off"