NexusOS/synapse/hardware.py

"""Hardware detection + a curated model catalog for the Models page.

Detects system RAM (reliable, psutil) and GPU VRAM (best-effort: nvidia-smi,
then Linux AMD sysfs). VRAM is None when it can't be read (e.g. AMD/Vulkan on
Windows) — the UI then recommends by RAM + model size alone.
"""
from __future__ import annotations

import glob
import subprocess
from typing import Any, Dict, List, Optional

# Curated local-friendly models with their Q4 on-disk sizes (GB) and role.
# Sizes are approximate default-quant download sizes. `required: True` marks
# the two models NexusOS itself depends on (memory curator + embeddings) —
# see DEFAULT_MEMORY_MODEL / DEFAULT_EMBED_MODEL in nexus_config.py, the
# single source of truth these two entries must stay in sync with. The
# Models page surfaces required models first and gates the rest behind them.
CATALOG: List[Dict[str, Any]] = [
    {"name": "mistral:latest",  "size_gb": 4.1, "params": "7B",   "role": "memory",     "note": "Memory curator: extracts facts, titles conversations", "required": True},
    {"name": "nomic-embed-text","size_gb": 0.27,"params": "—",    "role": "embeddings", "note": "Powers semantic recall of past conversations",          "required": True},
    {"name": "gemma2:2b",       "size_gb": 1.6, "params": "2B",   "role": "chat",       "note": "Smallest; fast on any GPU"},
    {"name": "llama3.2:3b",     "size_gb": 2.0, "params": "3B",   "role": "chat",       "note": "Small Llama, fits 4GB GPU"},
    {"name": "qwen3:4b",        "size_gb": 2.5, "params": "4B",   "role": "chat",       "note": "Reasoning (Think toggle); great on a 4GB card"},
    {"name": "phi3:mini",       "size_gb": 2.2, "params": "3.8B", "role": "chat",       "note": "Strong for its size"},
    {"name": "qwen2.5:7b",      "size_gb": 4.7, "params": "7B",   "role": "chat",       "note": "Capable 7B"},
    {"name": "llama3.1:8b",     "size_gb": 4.7, "params": "8B",   "role": "chat",       "note": "Higher quality; CPU or 6GB+ GPU"},
    {"name": "qwen2.5:14b",     "size_gb": 9.0, "params": "14B",  "role": "chat",       "note": "High-end; needs a 12-16GB GPU or 24GB+ RAM"},
]


def _ram_gb() -> Optional[float]:
    try:
        import psutil
        return round(psutil.virtual_memory().total / 1024 ** 3, 1)
    except Exception:
        return None


def _vram_gb() -> Optional[float]:
    # NVIDIA
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=3,
        )
        if r.returncode == 0 and r.stdout.strip():
            mb = max(int(x) for x in r.stdout.split())
            return round(mb / 1024, 1)
    except Exception:
        pass
    # Linux AMD/Intel via DRM sysfs
    try:
        best = 0
        for p in glob.glob("/sys/class/drm/card*/device/mem_info_vram_total"):
            try:
                best = max(best, int(open(p).read().strip()))
            except Exception:
                pass
        if best:
            return round(best / 1024 ** 3, 1)
    except Exception:
        pass
    return None  # unknown (e.g. AMD/Vulkan on Windows) -> recommend by RAM


def detect() -> Dict[str, Any]:
    return {"ram_gb": _ram_gb(), "vram_gb": _vram_gb()}


def _fit(size_gb: float, vram: Optional[float], ram: Optional[float]) -> str:
    """Where a model can run: 'gpu' (fully resident), 'ram' (CPU), or 'no'."""
    if vram and size_gb + 1.0 <= vram:      # ~1GB headroom for KV cache/context
        return "gpu"
    if ram and size_gb + 2.0 <= ram:        # ~2GB headroom for the OS/app
        return "ram"
    return "no"


def recommend() -> Dict[str, Any]:
    hw = detect()
    models = [
        {**m, "fit": _fit(m["size_gb"], hw["vram_gb"], hw["ram_gb"])}
        for m in CATALOG
    ]
    return {"hardware": hw, "models": models}