mita-code 0.1.0__py3-none-any.whl

mita/models/hardware.py

@@ -0,0 +1,289 @@
+"""Hardware detection: RAM, VRAM, CPU, GPU."""
+
+from __future__ import annotations
+
+import platform
+import subprocess
+
+import psutil
+from pydantic import BaseModel, Field
+
+
+class GPUInfo(BaseModel):
+    """Information about a detected GPU."""
+
+    name: str
+    vram_gb: float
+    vendor: str  # "nvidia", "amd", "apple"
+
+
+class HardwareInfo(BaseModel):
+    """Detected hardware capabilities."""
+
+    ram_gb: float
+    cpu_cores: int
+    cpu_name: str
+    gpus: list[GPUInfo] = Field(default_factory=list)
+    os: str  # "darwin", "linux", "windows"
+    apple_silicon: bool = False
+    unified_memory: bool = False
+
+    @property
+    def available_vram_gb(self) -> float:
+        """Effective VRAM available for models.
+
+        On Apple Silicon with unified memory, most of RAM is usable as VRAM.
+        On discrete GPUs, use the GPU VRAM.
+        Falls back to a conservative fraction of RAM for CPU-only inference.
+        """
+        if self.unified_memory:
+            # Apple Silicon can use ~75% of unified memory for ML
+            return self.ram_gb * 0.75
+        if self.gpus:
+            # Use the largest single GPU — Ollama loads a model into one GPU
+            return max(gpu.vram_gb for gpu in self.gpus)
+        # CPU-only: models load into RAM, leave room for OS
+        return max(0, self.ram_gb - 4)
+
+
+def detect_hardware() -> HardwareInfo:
+    """Detect hardware capabilities of the current machine."""
+    system = platform.system().lower()
+    os_name = {"darwin": "darwin", "linux": "linux", "windows": "windows"}.get(system, system)
+
+    ram_gb = round(psutil.virtual_memory().total / (1024**3), 1)
+    cpu_cores = psutil.cpu_count(logical=False) or psutil.cpu_count() or 1
+    cpu_name = _detect_cpu_name()
+
+    apple_silicon = False
+    unified_memory = False
+    gpus: list[GPUInfo] = []
+
+    if os_name == "darwin":
+        apple_silicon, unified_memory, gpus = _detect_macos_gpu(ram_gb)
+    elif os_name == "linux":
+        gpus = _detect_linux_gpu()
+
+    return HardwareInfo(
+        ram_gb=ram_gb,
+        cpu_cores=cpu_cores,
+        cpu_name=cpu_name,
+        gpus=gpus,
+        os=os_name,
+        apple_silicon=apple_silicon,
+        unified_memory=unified_memory,
+    )
+
+
+def _detect_cpu_name() -> str:
+    """Detect the CPU model name."""
+    system = platform.system().lower()
+
+    if system == "darwin":
+        try:
+            result = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.brand_string"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                return result.stdout.strip()
+        except (subprocess.SubprocessError, FileNotFoundError):
+            pass
+        # Apple Silicon doesn't have brand_string, use chip name
+        try:
+            result = subprocess.run(
+                ["sysctl", "-n", "hw.chip"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                return result.stdout.strip()
+        except (subprocess.SubprocessError, FileNotFoundError):
+            pass
+
+    elif system == "linux":
+        try:
+            with open("/proc/cpuinfo") as f:
+                for line in f:
+                    if line.startswith("model name"):
+                        return line.split(":", 1)[1].strip()
+        except (OSError, IndexError):
+            pass
+
+    return platform.processor() or "Unknown CPU"
+
+
+def _is_apple_silicon() -> bool:
+    """Detect Apple Silicon, even under Rosetta emulation."""
+    machine = platform.machine().lower()
+    if machine in ("arm64", "aarch64"):
+        return True
+    # Under Rosetta, platform.machine() returns x86_64.
+    # Check sysctl for the actual hardware.
+    try:
+        result = subprocess.run(
+            ["sysctl", "-n", "hw.optional.arm64"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0 and result.stdout.strip() == "1":
+            return True
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    return False
+
+
+def _detect_macos_gpu(ram_gb: float) -> tuple[bool, bool, list[GPUInfo]]:
+    """Detect GPU on macOS. Returns (apple_silicon, unified_memory, gpus)."""
+    apple_silicon = _is_apple_silicon()
+
+    if apple_silicon:
+        # Apple Silicon has unified memory — RAM is shared with GPU
+        chip_name = _get_apple_chip_name()
+        return (
+            True,
+            True,
+            [
+                GPUInfo(
+                    name=chip_name,
+                    vram_gb=ram_gb,  # Unified — all RAM is addressable
+                    vendor="apple",
+                )
+            ],
+        )
+
+    # Intel Mac — try to detect discrete GPU via system_profiler
+    gpus = _detect_macos_discrete_gpu()
+    return False, False, gpus
+
+
+def _get_apple_chip_name() -> str:
+    """Get the Apple Silicon chip name (e.g., 'Apple M2 Pro')."""
+    try:
+        result = subprocess.run(
+            ["sysctl", "-n", "machdep.cpu.brand_string"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0 and result.stdout.strip():
+            return result.stdout.strip()
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    return "Apple Silicon"
+
+
+def _detect_macos_discrete_gpu() -> list[GPUInfo]:
+    """Detect discrete GPUs on Intel Macs via system_profiler."""
+    try:
+        result = subprocess.run(
+            ["system_profiler", "SPDisplaysDataType", "-detailLevel", "basic"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode != 0:
+            return []
+
+        gpus: list[GPUInfo] = []
+        current_name = ""
+        for line in result.stdout.splitlines():
+            stripped = line.strip()
+            if stripped.startswith("Chipset Model:"):
+                current_name = stripped.split(":", 1)[1].strip()
+            elif stripped.startswith("VRAM") and current_name:
+                vram_str = stripped.split(":", 1)[1].strip()
+                vram_gb = _parse_vram_string(vram_str)
+                vendor = _guess_gpu_vendor(current_name)
+                gpus.append(GPUInfo(name=current_name, vram_gb=vram_gb, vendor=vendor))
+                current_name = ""
+        return gpus
+    except (subprocess.SubprocessError, FileNotFoundError):
+        return []
+
+
+def _detect_linux_gpu() -> list[GPUInfo]:
+    """Detect NVIDIA GPUs on Linux via nvidia-smi."""
+    gpus: list[GPUInfo] = []
+
+    # Try nvidia-smi
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=name,memory.total",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode == 0:
+            for line in result.stdout.strip().splitlines():
+                parts = line.split(",")
+                if len(parts) >= 2:
+                    name = parts[0].strip()
+                    vram_mb = float(parts[1].strip())
+                    gpus.append(
+                        GPUInfo(name=name, vram_gb=round(vram_mb / 1024, 1), vendor="nvidia")
+                    )
+    except (subprocess.SubprocessError, FileNotFoundError, ValueError):
+        pass
+
+    # Try AMD ROCm
+    if not gpus:
+        try:
+            result = subprocess.run(
+                ["rocm-smi", "--showmeminfo", "vram", "--csv"],
+                capture_output=True,
+                text=True,
+                timeout=10,
+            )
+            if result.returncode == 0:
+                for line in result.stdout.strip().splitlines()[1:]:  # skip header
+                    parts = line.split(",")
+                    if len(parts) >= 2:
+                        try:
+                            vram_bytes = int(parts[1].strip())
+                            gpus.append(
+                                GPUInfo(
+                                    name="AMD GPU",
+                                    vram_gb=round(vram_bytes / (1024**3), 1),
+                                    vendor="amd",
+                                )
+                            )
+                        except ValueError:
+                            pass
+        except (subprocess.SubprocessError, FileNotFoundError):
+            pass
+
+    return gpus
+
+
+def _parse_vram_string(vram_str: str) -> float:
+    """Parse VRAM strings like '8 GB', '4096 MB' into GB."""
+    vram_str = vram_str.strip().upper()
+    try:
+        if "GB" in vram_str:
+            return float(vram_str.replace("GB", "").strip())
+        if "MB" in vram_str:
+            return round(float(vram_str.replace("MB", "").strip()) / 1024, 1)
+    except ValueError:
+        pass
+    return 0.0
+
+
+def _guess_gpu_vendor(name: str) -> str:
+    """Guess GPU vendor from the chipset name."""
+    name_lower = name.lower()
+    if "nvidia" in name_lower or "geforce" in name_lower or "quadro" in name_lower:
+        return "nvidia"
+    if "amd" in name_lower or "radeon" in name_lower:
+        return "amd"
+    if "intel" in name_lower:
+        return "intel"
+    return "unknown"

mita/models/manager.py

@@ -0,0 +1,268 @@
+"""CLI command handlers for model management."""
+
+from __future__ import annotations
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import BarColumn, DownloadColumn, Progress, TextColumn, TransferSpeedColumn
+from rich.table import Table
+
+from mita.config.loader import load_config
+from mita.models.hardware import detect_hardware
+from mita.models.ollama_client import OllamaClient
+from mita.models.recommender import recommend_models
+from mita.models.registry import find_model
+
+console = Console()
+
+
+def _get_client() -> OllamaClient:
+    """Get an OllamaClient from the current config."""
+    cfg = load_config()
+    return OllamaClient(host=cfg.ollama.host, timeout=cfg.ollama.timeout)
+
+
+def _check_ollama(client: OllamaClient) -> bool:
+    """Check if Ollama is running, auto-start if configured."""
+    if client.is_running():
+        return True
+
+    # Try auto-start if configured
+    cfg = load_config()
+    if cfg.ollama.auto_manage:
+        from mita.models.server import ensure_server
+
+        if ensure_server(host=cfg.ollama.host, auto_manage=True, console=console):
+            return True
+
+    console.print(
+        "[red]Cannot connect to Ollama.[/red]\n"
+        "Make sure Ollama is running: [bold]ollama serve[/bold]"
+    )
+    return False
+
+
+def list_models() -> None:
+    """List all installed Ollama models."""
+    client = _get_client()
+    if not _check_ollama(client):
+        return
+
+    models = client.list_models()
+    if not models:
+        console.print(
+            "[dim]No models installed. "
+            "Run [bold]mita models pull <name>[/bold] to get started.[/dim]"
+        )
+        return
+
+    table = Table(title="Installed Models")
+    table.add_column("Name", style="bold")
+    table.add_column("Size", justify="right")
+    table.add_column("Parameters", justify="right")
+    table.add_column("Quantization")
+    table.add_column("Family")
+
+    for m in models:
+        table.add_row(
+            m.name,
+            f"{m.size_gb:.1f} GB",
+            m.parameter_size,
+            m.quantization,
+            m.family,
+        )
+
+    console.print(table)
+
+
+def pull_model(name: str) -> None:
+    """Pull a model from the Ollama registry."""
+    import ollama as ollama_lib
+
+    client = _get_client()
+    if not _check_ollama(client):
+        return
+
+    console.print(f"Pulling [bold]{name}[/bold]...")
+
+    try:
+        with Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            DownloadColumn(),
+            TransferSpeedColumn(),
+            console=console,
+        ) as progress:
+            task = progress.add_task("Downloading", total=None)
+
+            for update in client.pull(name, stream=True):
+                status = update.get("status", "")
+                completed = update.get("completed", 0)
+                total = update.get("total", 0)
+
+                if total > 0:
+                    progress.update(task, completed=completed, total=total, description=status)
+                else:
+                    progress.update(task, description=status)
+
+            progress.update(task, description="Complete")
+
+        console.print(f"[green]Successfully pulled {name}[/green]")
+    except ollama_lib.ResponseError as e:
+        console.print(f"[red]Failed to pull {name}: {e}[/red]")
+
+
+def remove_model(name: str) -> None:
+    """Remove an installed model."""
+    client = _get_client()
+    if not _check_ollama(client):
+        return
+
+    import ollama as ollama_lib
+
+    try:
+        client.remove(name)
+        console.print(f"[green]Removed {name}[/green]")
+    except ollama_lib.ResponseError as e:
+        console.print(f"[red]Failed to remove {name}: {e}[/red]")
+
+
+def show_model_info(name: str) -> None:
+    """Show details about a model."""
+    # Check registry first
+    card = find_model(name)
+    if card:
+        table = Table(title=f"Registry: {name}", show_header=False)
+        table.add_column("Field", style="bold")
+        table.add_column("Value")
+        table.add_row("Family", card.family)
+        table.add_row("Parameters", card.param_count)
+        table.add_row("Context Window", f"{card.context_window:,}")
+        table.add_row("Min RAM", f"{card.min_ram_gb:.0f} GB")
+        table.add_row("Min VRAM", f"{card.min_vram_gb:.0f} GB")
+        table.add_row("Tool Calls", "Yes" if card.tool_call_support else "No")
+        table.add_row("Quantization", card.quantization)
+        table.add_row("Description", card.description)
+        table.add_row("Tags", ", ".join(card.tags))
+        console.print(table)
+        console.print()
+
+    # Also show Ollama details if installed
+    import ollama as ollama_lib
+
+    client = _get_client()
+    if not _check_ollama(client):
+        return
+
+    try:
+        info = client.show(name)
+        details = info.get("details", {})
+        if isinstance(details, dict) and details:
+            table = Table(title=f"Ollama: {name}", show_header=False)
+            table.add_column("Field", style="bold")
+            table.add_column("Value")
+            for k, v in details.items():
+                table.add_row(str(k), str(v))
+            console.print(table)
+        elif not card:
+            console.print(f"[dim]Model {name} is installed but not in the curated registry.[/dim]")
+    except ollama_lib.ResponseError as e:
+        if not card:
+            console.print(f"[red]Model {name} not found in registry or Ollama: {e}[/red]")
+
+
+def show_recommendations() -> None:
+    """Detect hardware and recommend models."""
+    hw = detect_hardware()
+
+    # Show hardware summary
+    hw_panel = Panel(
+        f"[bold]CPU:[/bold] {hw.cpu_name} ({hw.cpu_cores} cores)\n"
+        f"[bold]RAM:[/bold] {hw.ram_gb:.1f} GB\n"
+        f"[bold]GPU:[/bold] {hw.gpus[0].name if hw.gpus else 'None detected'}\n"
+        f"[bold]VRAM:[/bold] {hw.available_vram_gb:.1f} GB effective\n"
+        f"[bold]OS:[/bold] {hw.os}"
+        + ("\n[bold]Unified Memory:[/bold] Yes" if hw.unified_memory else ""),
+        title="Hardware Detected",
+        border_style="blue",
+    )
+    console.print(hw_panel)
+    console.print()
+
+    # Show recommendations
+    recs = recommend_models(hw)
+    if not recs:
+        console.print(
+            "[yellow]No models in the registry fit your hardware.[/yellow]\n"
+            "You may need more RAM/VRAM, or try a smaller model manually."
+        )
+        return
+
+    table = Table(title="Recommended Models")
+    table.add_column("Model", style="bold")
+    table.add_column("Params", justify="right")
+    table.add_column("Context", justify="right")
+    table.add_column("Fit", justify="right")
+    table.add_column("Notes")
+
+    for rec in recs:
+        # Color-code fit score
+        score = rec.fit_score
+        if score >= 0.8:
+            fit_style = "green"
+        elif score >= 0.6:
+            fit_style = "yellow"
+        else:
+            fit_style = "red"
+
+        table.add_row(
+            rec.model.name,
+            rec.model.param_count,
+            f"{rec.model.context_window:,}",
+            f"[{fit_style}]{score:.0%}[/{fit_style}]",
+            rec.notes,
+        )
+
+    console.print(table)
+
+
+def set_default_model(name: str) -> None:
+    """Set the default model (prints instruction since config is TOML-based)."""
+    card = find_model(name)
+    if card:
+        console.print(f"[green]Model [bold]{name}[/bold] is in the registry.[/green]")
+    else:
+        console.print(f"[yellow]Model [bold]{name}[/bold] is not in the curated registry.[/yellow]")
+
+    console.print(
+        f"\nTo set as default, add to your config:\n\n"
+        f'  [bold]mita config set model.default "{name}"[/bold]\n\n'
+        f"Or edit your config file directly:\n\n"
+        f'  [dim][model]\n  default = "{name}"[/dim]'
+    )
+
+
+def show_hardware() -> None:
+    """Show detected hardware information."""
+    hw = detect_hardware()
+
+    table = Table(title="Hardware Information", show_header=False)
+    table.add_column("Field", style="bold")
+    table.add_column("Value")
+    table.add_row("OS", hw.os)
+    table.add_row("CPU", hw.cpu_name)
+    table.add_row("CPU Cores", str(hw.cpu_cores))
+    table.add_row("RAM", f"{hw.ram_gb:.1f} GB")
+    table.add_row("Apple Silicon", "Yes" if hw.apple_silicon else "No")
+    table.add_row("Unified Memory", "Yes" if hw.unified_memory else "No")
+
+    for i, gpu in enumerate(hw.gpus):
+        prefix = f"GPU {i}" if len(hw.gpus) > 1 else "GPU"
+        table.add_row(prefix, f"{gpu.name} ({gpu.vram_gb:.1f} GB, {gpu.vendor})")
+
+    if not hw.gpus:
+        table.add_row("GPU", "None detected")
+
+    table.add_row("Effective VRAM", f"{hw.available_vram_gb:.1f} GB")
+
+    console.print(table)

mita/models/ollama_client.py

@@ -0,0 +1,104 @@
+"""Wrapper around the Ollama Python client."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+import ollama
+from pydantic import BaseModel
+
+
+class OllamaModelInfo(BaseModel):
+    """Info about an installed Ollama model."""
+
+    name: str
+    size_gb: float
+    parameter_size: str = ""
+    quantization: str = ""
+    family: str = ""
+    modified_at: str = ""
+
+
+class OllamaClient:
+    """Thin wrapper around the ollama Python client."""
+
+    def __init__(self, host: str = "http://localhost:11434", timeout: int = 120) -> None:
+        self._client = ollama.Client(host=host, timeout=timeout)
+        self._host = host
+
+    def is_running(self) -> bool:
+        """Check if the Ollama server is reachable."""
+        try:
+            self._client.list()
+            return True
+        except (ConnectionError, OSError, ollama.ResponseError, TimeoutError):
+            return False
+
+    def list_models(self) -> list[OllamaModelInfo]:
+        """List all installed models."""
+        response = self._client.list()
+        models: list[OllamaModelInfo] = []
+        for m in response.models:
+            details: Any = m.details or {}
+            size_gb = round((m.size or 0) / (1024**3), 2)
+
+            # Handle details as either dict or object
+            if isinstance(details, dict):
+                param_size = details.get("parameter_size", "")
+                quant = details.get("quantization_level", "")
+                family = details.get("family", "")
+            else:
+                param_size = getattr(details, "parameter_size", "") or ""
+                quant = getattr(details, "quantization_level", "") or ""
+                family = getattr(details, "family", "") or ""
+
+            models.append(
+                OllamaModelInfo(
+                    name=m.model or "",
+                    size_gb=size_gb,
+                    parameter_size=param_size,
+                    quantization=quant,
+                    family=family,
+                    modified_at=str(m.modified_at or ""),
+                )
+            )
+        return models
+
+    def pull(self, model_name: str, stream: bool = True) -> Iterator[dict[str, Any]]:
+        """Pull a model from the Ollama registry.
+
+        Yields progress dicts with keys like 'status', 'completed', 'total'.
+        Raises ollama.ResponseError on failure (e.g., model not found).
+        """
+        if stream:
+            response = self._client.pull(model_name, stream=True)
+            for chunk in response:
+                if isinstance(chunk, dict):
+                    yield chunk
+                else:
+                    yield {"status": getattr(chunk, "status", str(chunk))}
+        else:
+            self._client.pull(model_name, stream=False)
+            yield {"status": "success"}
+
+    def remove(self, model_name: str) -> None:
+        """Remove an installed model.
+
+        Raises ollama.ResponseError if the model is not found.
+        """
+        self._client.delete(model_name)
+
+    def show(self, model_name: str) -> dict[str, Any]:
+        """Show details about a model."""
+        response = self._client.show(model_name)
+        if isinstance(response, dict):
+            return response
+        # Convert response object to dict
+        return {
+            "modelfile": getattr(response, "modelfile", ""),
+            "parameters": getattr(response, "parameters", ""),
+            "template": getattr(response, "template", ""),
+            "details": getattr(response, "details", {}),
+            "model_info": getattr(response, "model_info", {}),
+        }