"""Tests for ``recommend_model()`` hardware-aware model recommendation.""" from __future__ import annotations from openjarvis.core.config import GpuInfo, HardwareInfo, recommend_model class TestRecommendModelGpu: """GPU-based model recommendation.""" def test_24gb_gpu_picks_qwen35_35b(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=64.0, gpu=GpuInfo(vendor="nvidia", name="RTX 4090", vram_gb=24.0, count=1), ) result = recommend_model(hw, "ollama") # 35B * 0.5 * 1.1 = 19.25 GB; available = 24 * 0.9 = 21.6 → fits assert result == "qwen3.5:35b" def test_8gb_gpu_picks_qwen35_14b(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=32.0, gpu=GpuInfo(vendor="nvidia", name="RTX 3070", vram_gb=8.0, count=1), ) result = recommend_model(hw, "ollama") # 14B * 0.5 * 1.1 = 7.7 GB; available = 8 * 0.9 = 7.2 → too big # 8B * 0.5 * 1.1 = 4.4 GB; available = 7.2 → fits assert result == "qwen3.5:8b" def test_4gb_gpu_picks_qwen35_4b(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=16.0, gpu=GpuInfo(vendor="nvidia", name="GTX 1650", vram_gb=4.0, count=1), ) result = recommend_model(hw, "ollama") # 4B * 0.5 * 1.1 = 2.2 GB; available = 4 * 0.9 = 3.6 → fits assert result == "qwen3.5:4b" def test_2gb_gpu_picks_qwen35_3b(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=8.0, gpu=GpuInfo(vendor="nvidia", name="GTX 750", vram_gb=2.0, count=1), ) result = recommend_model(hw, "ollama") # 3B * 0.5 * 1.1 = 1.65 GB; available = 2 * 0.9 = 1.8 → fits assert result == "qwen3.5:3b" def test_multi_gpu_picks_larger_model(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=256.0, gpu=GpuInfo(vendor="nvidia", name="A100", vram_gb=80.0, count=2), ) result = recommend_model(hw, "vllm") # available = 80 * 2 * 0.9 = 144 GB # 397B * 0.5 * 1.1 = 218.35 → too big # 122B * 0.5 * 1.1 = 67.1 → fits assert result == "qwen3.5:122b" def test_huge_vram_picks_397b(self) -> None: hw = HardwareInfo( platform="linux", ram_gb=512.0, gpu=GpuInfo(vendor="nvidia", name="H100", vram_gb=80.0, count=4), ) result = recommend_model(hw, "vllm") # available = 80 * 4 * 0.9 = 288 GB # 397B * 0.5 * 1.1 = 218.35 → fits assert result == "qwen3.5:397b" class TestRecommendModelCpuOnly: """CPU-only model recommendation.""" def test_cpu_only_16gb_ram(self) -> None: hw = HardwareInfo(platform="linux", ram_gb=16.0, gpu=None) result = recommend_model(hw, "llamacpp") # available = (16 - 4) * 0.8 = 9.6 GB # 14B * 0.5 * 1.1 = 7.7 → fits assert result == "qwen3.5:14b" def test_cpu_only_8gb_ram(self) -> None: hw = HardwareInfo(platform="linux", ram_gb=8.0, gpu=None) result = recommend_model(hw, "llamacpp") # available = (8 - 4) * 0.8 = 3.2 GB # 8B * 0.5 * 1.1 = 4.4 → too big # 4B * 0.5 * 1.1 = 2.2 → fits assert result == "qwen3.5:4b" def test_cpu_only_4gb_ram(self) -> None: hw = HardwareInfo(platform="linux", ram_gb=4.0, gpu=None) result = recommend_model(hw, "llamacpp") # available = (4 - 4) * 0.8 = 0 → nothing fits assert result == "" class TestRecommendModelEdgeCases: """Edge cases.""" def test_no_ram_no_gpu(self) -> None: hw = HardwareInfo(platform="linux", ram_gb=0.0, gpu=None) assert recommend_model(hw, "ollama") == "" def test_engine_filter(self) -> None: """397b is not supported on ollama, only vllm/sglang.""" hw = HardwareInfo( platform="linux", ram_gb=512.0, gpu=GpuInfo(vendor="nvidia", name="H100", vram_gb=80.0, count=4), ) # With ollama, 397b is excluded (only vllm, sglang) result = recommend_model(hw, "ollama") assert result == "qwen3.5:122b" class TestRecommendModelMlx: """Apple Silicon (MLX) model recommendation.""" def test_apple_silicon_8gb_mlx(self) -> None: hw = HardwareInfo( platform="darwin", ram_gb=8.0, gpu=GpuInfo(vendor="apple", name="Apple M1", vram_gb=8.0, count=1), ) result = recommend_model(hw, "mlx") assert result == "qwen3.5:8b" def test_apple_silicon_16gb_mlx(self) -> None: hw = HardwareInfo( platform="darwin", ram_gb=16.0, gpu=GpuInfo(vendor="apple", name="Apple M2", vram_gb=16.0, count=1), ) result = recommend_model(hw, "mlx") assert result == "qwen3.5:14b" def test_apple_silicon_32gb_mlx(self) -> None: hw = HardwareInfo( platform="darwin", ram_gb=32.0, gpu=GpuInfo(vendor="apple", name="Apple M2 Pro", vram_gb=32.0, count=1), ) result = recommend_model(hw, "mlx") assert result == "qwen3.5:14b" def test_apple_silicon_64gb_mlx(self) -> None: hw = HardwareInfo( platform="darwin", ram_gb=64.0, gpu=GpuInfo(vendor="apple", name="Apple M2 Max", vram_gb=64.0, count=1), ) result = recommend_model(hw, "mlx") assert result == "qwen3.5:14b"