diff --git a/src/openjarvis/cli/model.py b/src/openjarvis/cli/model.py index 4a6fc9d1..d3243311 100644 --- a/src/openjarvis/cli/model.py +++ b/src/openjarvis/cli/model.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import subprocess import sys import click @@ -15,6 +16,7 @@ from openjarvis.core.config import load_config from openjarvis.core.registry import ModelRegistry from openjarvis.engine import discover_engines, discover_models from openjarvis.intelligence import merge_discovered_models, register_builtin_models +from openjarvis.intelligence.model_catalog import BUILTIN_MODELS @click.group() @@ -149,18 +151,8 @@ def info(model_name: str) -> None: console.print(Panel("\n".join(lines), title=spec.name, border_style="blue")) -@model.command() -@click.argument("model_name") -def pull(model_name: str) -> None: - """Download a model (Ollama only).""" - console = Console() - config = load_config() - host = ( - config.engine.ollama_host - or os.environ.get("OLLAMA_HOST") - or "http://localhost:11434" - ).rstrip("/") - +def ollama_pull(host: str, model_name: str, console: Console) -> bool: + """Pull a model via Ollama API. Returns True on success.""" console.print(f"Pulling [cyan]{model_name}[/cyan] via Ollama...") try: with httpx.stream( @@ -171,7 +163,6 @@ def pull(model_name: str) -> None: ) as resp: resp.raise_for_status() import json - for line in resp.iter_lines(): if not line.strip(): continue @@ -188,9 +179,90 @@ def pull(model_name: str) -> None: elif status: console.print(f" {status}") console.print(f"\n[green]Successfully pulled {model_name}[/green]") + return True except httpx.ConnectError: console.print("[red]Cannot connect to Ollama.[/red] Is it running?") - sys.exit(1) + return False except httpx.HTTPStatusError as exc: console.print(f"[red]Ollama error:[/red] {exc.response.status_code}") - sys.exit(1) + return False + + +def find_model_spec(model_name: str): + """Look up a model in the builtin catalog. Returns None if not found.""" + for spec in BUILTIN_MODELS: + if spec.model_id == model_name: + return spec + return None + + +def hf_download(repo: str, filename: str | None, console: Console) -> bool: + """Download from HuggingFace via huggingface-cli. Returns True on success.""" + cmd = ["huggingface-cli", "download", repo] + if filename: + cmd.append(filename) + try: + subprocess.run(cmd, check=True) + console.print("[green]Download complete.[/green]") + return True + except FileNotFoundError: + console.print( + "[red]huggingface-cli not found.[/red]\n" + "Install it: [cyan]pip install huggingface_hub[/cyan]\n" + f"Or download manually: https://huggingface.co/{repo}" + ) + return False + except subprocess.CalledProcessError: + console.print("[red]Download failed.[/red]") + return False + + +@model.command() +@click.argument("model_name") +@click.option("--engine", default=None, help="Engine to download for (ollama, llamacpp, mlx).") +def pull(model_name: str, engine: str | None) -> None: + """Download a model.""" + console = Console() + config = load_config() + engine = engine or config.engine.default or "ollama" + + if engine == "ollama": + host = ( + config.engine.ollama_host + or os.environ.get("OLLAMA_HOST") + or "http://localhost:11434" + ).rstrip("/") + if not ollama_pull(host, model_name, console): + sys.exit(1) + elif engine in ("llamacpp", "mlx"): + spec = find_model_spec(model_name) + if not spec: + console.print(f"[red]Model not in catalog:[/red] {model_name}") + sys.exit(1) + if engine == "llamacpp": + repo = spec.metadata.get("hf_repo", "") + gguf = spec.metadata.get("gguf_file", "") + if not repo or not gguf: + console.print(f"[red]No GGUF download info for {model_name}[/red]") + sys.exit(1) + console.print(f"Downloading [cyan]{gguf}[/cyan] from {repo}...") + if not hf_download(repo, gguf, console): + sys.exit(1) + else: # mlx + mlx_repo = spec.metadata.get("mlx_repo", "") + if not mlx_repo: + console.print(f"[red]No MLX repo info for {model_name}[/red]") + sys.exit(1) + console.print(f"Downloading [cyan]{mlx_repo}[/cyan]...") + if not hf_download(mlx_repo, None, console): + sys.exit(1) + elif engine in ("vllm", "sglang"): + console.print( + f"[cyan]{model_name}[/cyan] will download automatically when " + f"{engine} starts serving it." + ) + else: + console.print( + f"Manual download required for engine [cyan]{engine}[/cyan].\n" + f"Check the engine documentation for instructions." + ) diff --git a/tests/cli/test_model_pull.py b/tests/cli/test_model_pull.py new file mode 100644 index 00000000..a87a55ae --- /dev/null +++ b/tests/cli/test_model_pull.py @@ -0,0 +1,108 @@ +"""Tests for ``jarvis model pull`` multi-engine support.""" + +from __future__ import annotations + +from unittest import mock + +import pytest +from click.testing import CliRunner +from rich.console import Console + +from openjarvis.cli.model import ollama_pull + + +class TestOllamaPull: + """Test the extracted ollama_pull helper.""" + + def test_ollama_pull_success(self) -> None: + import io + console = Console(file=io.StringIO()) + mock_lines = [ + '{"status": "pulling manifest"}', + '{"status": "downloading", "total": 100, "completed": 100}', + '{"status": "success"}', + ] + mock_resp = mock.MagicMock() + mock_resp.raise_for_status = mock.MagicMock() + mock_resp.iter_lines.return_value = iter(mock_lines) + mock_resp.__enter__ = mock.MagicMock(return_value=mock_resp) + mock_resp.__exit__ = mock.MagicMock(return_value=False) + + with mock.patch("httpx.stream", return_value=mock_resp): + result = ollama_pull("http://localhost:11434", "qwen3.5:3b", console) + assert result is True + + def test_ollama_pull_connect_error(self) -> None: + import io + import httpx + + console = Console(file=io.StringIO()) + with mock.patch("httpx.stream", side_effect=httpx.ConnectError("refused")): + result = ollama_pull("http://localhost:11434", "qwen3.5:3b", console) + assert result is False + + +class TestPullCliMultiEngine: + """Test the pull CLI command dispatches to correct engine.""" + + def test_pull_llamacpp_uses_huggingface_cli(self) -> None: + from openjarvis.cli import cli + + runner = CliRunner() + with ( + mock.patch("openjarvis.cli.model.load_config") as mock_cfg, + mock.patch("subprocess.run") as mock_run, + ): + mock_cfg.return_value.engine.default = "llamacpp" + mock_cfg.return_value.engine.ollama_host = None + mock_run.return_value = mock.MagicMock(returncode=0) + + result = runner.invoke( + cli, ["model", "pull", "qwen3.5:8b", "--engine", "llamacpp"] + ) + + assert result.exit_code == 0 + mock_run.assert_called_once() + call_args = mock_run.call_args[0][0] + assert "huggingface-cli" in call_args + assert "qwen3.5-8b-q4_k_m.gguf" in call_args + + def test_pull_mlx_uses_huggingface_cli(self) -> None: + from openjarvis.cli import cli + + runner = CliRunner() + with ( + mock.patch("openjarvis.cli.model.load_config") as mock_cfg, + mock.patch("subprocess.run") as mock_run, + ): + mock_cfg.return_value.engine.default = "mlx" + mock_cfg.return_value.engine.ollama_host = None + mock_run.return_value = mock.MagicMock(returncode=0) + + result = runner.invoke( + cli, ["model", "pull", "qwen3.5:8b", "--engine", "mlx"] + ) + + assert result.exit_code == 0 + mock_run.assert_called_once() + call_args = mock_run.call_args[0][0] + assert "huggingface-cli" in call_args + assert "mlx-community/Qwen3.5-8B-4bit" in call_args + + def test_pull_llamacpp_huggingface_cli_not_found(self) -> None: + from openjarvis.cli import cli + + runner = CliRunner() + with ( + mock.patch("openjarvis.cli.model.load_config") as mock_cfg, + mock.patch("subprocess.run", side_effect=FileNotFoundError), + ): + mock_cfg.return_value.engine.default = "llamacpp" + mock_cfg.return_value.engine.ollama_host = None + + result = runner.invoke( + cli, ["model", "pull", "qwen3.5:8b", "--engine", "llamacpp"] + ) + + assert result.exit_code != 0 + assert "huggingface_hub" in result.output or "pip install" in result.output