caudate-cli 0.1.0__py3-none-any.whl

api/forge_system_routes.py

@@ -0,0 +1,294 @@
+"""Forge system + provider routes — hardware detection, provider scan.
+
+Mirrors LocalForge's /api/system/hardware and /api/providers/scan
+under Cognos's /forge prefix. Lazy detection + on-disk cache so the
+hot path is fast on every platform.
+
+Endpoints:
+  GET  /forge/system/hardware?refresh=1   CPU/RAM/GPU info
+  GET  /forge/providers/scan              probe Ollama + LM-Studio
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import httpx
+from fastapi import APIRouter
+
+logger = logging.getLogger(__name__)
+
+_HARDWARE_CACHE_PATH = Path(__file__).resolve().parents[1] / "data" / "hardware-cache.json"
+_memory_cache: dict[str, Any] | None = None
+
+
+# ─────────────────────────── Hardware detection ────────────────────────
+
+
+def _read_disk_cache() -> dict[str, Any] | None:
+    try:
+        if not _HARDWARE_CACHE_PATH.exists():
+            return None
+        data = json.loads(_HARDWARE_CACHE_PATH.read_text())
+        if data.get("platform") != sys.platform:
+            return None
+        return data
+    except Exception:
+        return None
+
+
+def _write_disk_cache(value: dict[str, Any]) -> None:
+    try:
+        _HARDWARE_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
+        _HARDWARE_CACHE_PATH.write_text(json.dumps(value, indent=2))
+    except Exception as e:
+        logger.debug(f"hardware cache write failed: {e}")
+
+
+def _detect_apple_silicon() -> bool:
+    return sys.platform == "darwin" and platform.machine() in ("arm64", "aarch64")
+
+
+def _detect_total_ram_mb() -> int:
+    try:
+        import psutil
+        return round(psutil.virtual_memory().total / 1024 / 1024)
+    except Exception:
+        # Fallback for hosts without psutil
+        try:
+            with open("/proc/meminfo") as f:
+                for line in f:
+                    if line.startswith("MemTotal:"):
+                        kb = int(line.split()[1])
+                        return round(kb / 1024)
+        except Exception:
+            pass
+    return 0
+
+
+def _detect_nvidia_gpus() -> list[dict[str, Any]]:
+    """Return one dict per NVIDIA GPU via nvidia-smi."""
+    if shutil.which("nvidia-smi") is None:
+        return []
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi",
+             "--query-gpu=name,memory.total",
+             "--format=csv,noheader,nounits"],
+            stderr=subprocess.DEVNULL, timeout=3,
+        ).decode().strip()
+    except Exception:
+        return []
+    gpus: list[dict[str, Any]] = []
+    for line in out.splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) >= 2:
+            try:
+                vram = int(float(parts[1]))
+            except ValueError:
+                continue
+            gpus.append({
+                "vendor": "NVIDIA",
+                "model": parts[0],
+                "vram_mb": vram,
+            })
+    return gpus
+
+
+def _detect_amd_gpus() -> list[dict[str, Any]]:
+    """Best-effort AMD ROCm detection via rocm-smi."""
+    if shutil.which("rocm-smi") is None:
+        return []
+    try:
+        out = subprocess.check_output(
+            ["rocm-smi", "--showmeminfo", "vram", "--csv"],
+            stderr=subprocess.DEVNULL, timeout=3,
+        ).decode()
+    except Exception:
+        return []
+    gpus: list[dict[str, Any]] = []
+    # Parse CSV; columns vary by version. Look for vram total bytes.
+    for line in out.splitlines():
+        if "card" not in line.lower():
+            continue
+        cols = [c.strip() for c in line.split(",")]
+        if len(cols) >= 2:
+            try:
+                vram_bytes = int(cols[-1])
+                gpus.append({
+                    "vendor": "AMD",
+                    "model": cols[0],
+                    "vram_mb": vram_bytes // (1024 * 1024),
+                })
+            except ValueError:
+                pass
+    return gpus
+
+
+def _detect_apple_gpu_budget(total_ram_mb: int) -> dict[str, Any]:
+    """Apple Silicon: GPU shares unified memory. Budget ~75% of RAM."""
+    return {
+        "vendor": "Apple",
+        "model": "Apple Silicon",
+        "vram_mb": round(total_ram_mb * 0.75),
+    }
+
+
+def _detect_hardware() -> dict[str, Any]:
+    """Snapshot the local machine's CPU/RAM/GPU. Conservative — failures
+    return ``{available: false, reason: ...}`` so the UI can render a
+    helpful message instead of a stack trace."""
+    plat = sys.platform
+    total_ram = _detect_total_ram_mb()
+    is_apple = _detect_apple_silicon()
+    cpu_count = os.cpu_count() or 0
+
+    if is_apple:
+        gpu = _detect_apple_gpu_budget(total_ram)
+        return {
+            "available": True,
+            "platform": plat,
+            "is_apple_silicon": True,
+            "cpu_count": cpu_count,
+            "gpus": [gpu],
+            "total_vram_mb": gpu["vram_mb"],
+            "total_system_ram_mb": total_ram,
+            "detected_at": time.time(),
+        }
+
+    gpus = _detect_nvidia_gpus() + _detect_amd_gpus()
+    if not gpus:
+        return {
+            "available": False,
+            "platform": plat,
+            "is_apple_silicon": False,
+            "cpu_count": cpu_count,
+            "total_system_ram_mb": total_ram,
+            "reason": "No discrete GPU with reportable VRAM detected "
+                      "(checked nvidia-smi, rocm-smi).",
+            "detected_at": time.time(),
+        }
+    gpus.sort(key=lambda g: g["vram_mb"], reverse=True)
+    return {
+        "available": True,
+        "platform": plat,
+        "is_apple_silicon": False,
+        "cpu_count": cpu_count,
+        "gpus": gpus,
+        "total_vram_mb": max(g["vram_mb"] for g in gpus),
+        "total_system_ram_mb": total_ram,
+        "detected_at": time.time(),
+    }
+
+
+# ─────────────────────────── Provider scan ─────────────────────────────
+
+
+_PROVIDERS = [
+    {
+        "id": "ollama",
+        "label": "Ollama",
+        "default_base_url": "http://127.0.0.1:11434",
+        "tags_path": "/api/tags",
+    },
+    {
+        "id": "lm-studio",
+        "label": "LM Studio",
+        "default_base_url": "http://127.0.0.1:1234",
+        "tags_path": "/v1/models",
+    },
+]
+
+
+async def _probe_ollama(base_url: str, timeout: float = 1.5) -> int | None:
+    """Return model count if Ollama responds with the expected shape."""
+    url = base_url.rstrip("/") + "/api/tags"
+    try:
+        async with httpx.AsyncClient(timeout=timeout) as cx:
+            r = await cx.get(url)
+            if r.status_code != 200:
+                return None
+            data = r.json()
+            if not isinstance(data, dict) or not isinstance(data.get("models"), list):
+                return None
+            return len(data["models"])
+    except Exception:
+        return None
+
+
+async def _probe_lm_studio(base_url: str, timeout: float = 1.5) -> int | None:
+    """Return model count if LM-Studio's OpenAI-compat endpoint answers."""
+    url = base_url.rstrip("/") + "/v1/models"
+    try:
+        async with httpx.AsyncClient(timeout=timeout) as cx:
+            r = await cx.get(url)
+            if r.status_code != 200:
+                return None
+            data = r.json()
+            if not isinstance(data, dict) or not isinstance(data.get("data"), list):
+                return None
+            return len(data["data"])
+    except Exception:
+        return None
+
+
+async def _scan_providers() -> list[dict[str, Any]]:
+    hits: list[dict[str, Any]] = []
+    probes = [
+        (_PROVIDERS[0], _probe_ollama(_PROVIDERS[0]["default_base_url"])),
+        (_PROVIDERS[1], _probe_lm_studio(_PROVIDERS[1]["default_base_url"])),
+    ]
+    for desc, coro in probes:
+        try:
+            count = await coro
+        except Exception:
+            count = None
+        if count is not None:
+            hits.append({
+                "provider_id": desc["id"],
+                "label": desc["label"],
+                "url": desc["default_base_url"],
+                "model_count": count,
+            })
+    return hits
+
+
+# ─────────────────────────────── Router ────────────────────────────────
+
+
+def build_router() -> APIRouter:
+    router = APIRouter(prefix="/forge", tags=["forge-system"])
+
+    @router.get("/system/hardware")
+    async def hardware(refresh: int = 0) -> dict[str, Any]:
+        global _memory_cache
+        if not refresh and _memory_cache is not None:
+            return _memory_cache
+        if not refresh:
+            disk = _read_disk_cache()
+            if disk is not None:
+                _memory_cache = disk
+                return disk
+        result = await asyncio.get_running_loop().run_in_executor(
+            None, _detect_hardware,
+        )
+        _memory_cache = result
+        if result.get("available"):
+            _write_disk_cache(result)
+        return result
+
+    @router.get("/providers/scan")
+    async def providers_scan() -> dict[str, Any]:
+        return {"hits": await _scan_providers()}
+
+    return router