llama-cpp-server/tests/test_generate_models_preset.py

import re
import unittest
from unittest.mock import patch
from pathlib import Path

from generate_models_preset import (
    ModelEntry,
    ModelMetadata,
    build_parser,
    build_preset_values,
    classify_model_family,
    make_alias,
    render_ini,
)


def make_entry(
    *,
    filename: str = "model.gguf",
    architecture: str | None = "llama",
    context_length: int | None = 32768,
    block_count: int | None = 40,
    expert_count: int | None = None,
    expert_used_count: int | None = None,
    mmproj: bool = False,
) -> ModelEntry:
    return ModelEntry(
        publisher="publisher",
        model_dir="model-dir",
        model_path=Path(r"C:\models") / filename,
        mmproj_path=(Path(r"C:\models") / "mmproj.gguf") if mmproj else None,
        quant_key="Q4_K_M",
        metadata=ModelMetadata(
            architecture=architecture,
            context_length=context_length,
            block_count=block_count,
            expert_count=expert_count,
            expert_used_count=expert_used_count,
            size_label=None,
            general_name=None,
            file_size_gib=24.0,
        ),
    )


class ClassificationTests(unittest.TestCase):
    def test_moe_detection_uses_metadata_not_filename(self):
        entry = make_entry(filename="mixtral-8x7b.gguf", architecture="llama")
        self.assertEqual(classify_model_family(entry), "dense")

    def test_moe_detection_accepts_moe_architecture_metadata(self):
        entry = make_entry(architecture="qwen35moe")
        self.assertEqual(classify_model_family(entry), "moe")

    def test_moe_metadata_takes_priority_over_mmproj(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="qwen35moe",
            expert_count=128,
            expert_used_count=8,
            mmproj=True,
        )
        self.assertEqual(classify_model_family(entry), "moe")


class PresetValueTests(unittest.TestCase):
    def test_dense_values_keep_gpu_layers_field(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="llama",
            context_length=32768,
            block_count=40,
        )

        with patch("generate_models_preset.os.cpu_count", return_value=20):
            values = build_preset_values(entry, "smart")

        self.assertEqual(values["ctx-size"], "32768")
        self.assertEqual(values["n-gpu-layers"], "40")
        self.assertEqual(values["parallel"], "4")
        self.assertEqual(values["threads"], "14")

    def test_moe_values_include_broader_runtime_fields(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="qwen35moe",
            context_length=131072,
            block_count=99,
            expert_count=128,
            expert_used_count=8,
        )

        with patch("generate_models_preset.os.cpu_count", return_value=20):
            values = build_preset_values(entry, "smart")

        self.assertEqual(values["ctx-size"], "43008")
        self.assertEqual(values["n-gpu-layers"], "99")
        self.assertEqual(values["parallel"], "4")
        self.assertEqual(values["batch-size"], "2048")
        self.assertEqual(values["ubatch-size"], "512")
        self.assertEqual(values["n-cpu-moe"], "0")
        self.assertEqual(values["cache-type-k"], "q8_0")
        self.assertEqual(values["cache-type-v"], "q8_0")
        self.assertEqual(values["mmap"], "true")
        self.assertEqual(values["threads"], "14")
        self.assertEqual(values["temp"], "0.7")
        self.assertEqual(values["top-k"], "40")
        self.assertEqual(values["top-p"], "0.95")
        self.assertEqual(values["min-p"], "0.0")
        self.assertEqual(values["repeat-penalty"], "1.0")

    def test_moe_with_mmproj_still_gets_n_cpu_moe(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="qwen35moe",
            context_length=131072,
            block_count=99,
            expert_count=128,
            expert_used_count=8,
            mmproj=True,
        )

        with patch("generate_models_preset.os.cpu_count", return_value=20):
            values = build_preset_values(entry, "smart")

        self.assertEqual(values["n-cpu-moe"], "0")
        self.assertEqual(values["batch-size"], "2048")

    def test_dense_models_get_sampling_defaults(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="llama",
            context_length=32768,
            block_count=40,
        )

        with patch("generate_models_preset.os.cpu_count", return_value=20):
            values = build_preset_values(entry, "smart")

        self.assertEqual(values["temp"], "0.8")
        self.assertEqual(values["top-k"], "40")
        self.assertEqual(values["top-p"], "0.95")
        self.assertEqual(values["min-p"], "0.05")
        self.assertEqual(values["repeat-penalty"], "1.0")


class RenderIniTests(unittest.TestCase):
    def test_render_ini_emits_range_comments_next_to_numeric_values(self):
        entry = make_entry(
            filename="plain-model.gguf",
            architecture="qwen35moe",
            context_length=131072,
            block_count=99,
            expert_count=128,
            expert_used_count=8,
        )

        with patch("generate_models_preset.os.cpu_count", return_value=20):
            text = render_ini(
                entries=[entry],
                models_dir=Path(r"C:\models-root"),
                select_mode="all",
                profile="smart",
                alias_style="section",
                backup_path=None,
            )

        self.assertRegex(
            text,
            re.compile(
                r"; ctx-size range = 512\.\.131072; chosen = 43008\s+ctx-size = 43008",
                re.MULTILINE,
            ),
        )
        self.assertRegex(
            text,
            re.compile(
                r"; n-gpu-layers range = 0\.\.99; chosen = 99\s+n-gpu-layers = 99",
                re.MULTILINE,
            ),
        )
        self.assertRegex(
            text,
            re.compile(
                r"; n-cpu-moe range = 0\.\.99; chosen = 0\s+n-cpu-moe = 0",
                re.MULTILINE,
            ),
        )
        self.assertRegex(
            text,
            re.compile(
                r"; threads range = 1\.\.20; chosen = 14\s+threads = 14",
                re.MULTILINE,
            ),
        )
        self.assertRegex(
            text,
            re.compile(
                r"; parallel range = 1\.\.4; chosen = 4\s+parallel = 4",
                re.MULTILINE,
            ),
        )
        self.assertRegex(
            text,
            re.compile(
                r"temp = 0\.7\s+top-k = 40\s+top-p = 0\.95\s+min-p = 0\.0\s+repeat-penalty = 1\.0",
                re.MULTILINE,
            ),
        )
        self.assertIn(
            "; Heuristic sampling defaults / 经验采样默认值, not GGUF metadata-derived",
            text,
        )


class ThreadHeuristicsTests(unittest.TestCase):
    def test_threads_use_seventy_percent_of_logical_cpu_count(self):
        entry = make_entry()
        with patch("generate_models_preset.os.cpu_count", return_value=20):
            values = build_preset_values(entry, "smart")

        self.assertEqual(values["threads"], "14")


class CliTests(unittest.TestCase):
    def test_backup_is_disabled_by_default(self):
        parser = build_parser()
        args = parser.parse_args(["--models-dir", r"C:\models-root"])
        self.assertFalse(args.backup)


class AliasTests(unittest.TestCase):
    def test_section_alias_matches_router_model_identifier_shape(self):
        alias = make_alias(
            publisher="ggml-org",
            model_dir="Qwen3-8B-GGUF",
            model_file="Qwen3-8B-Q4_K_M.gguf",
            alias_style="section",
            used_aliases=set(),
        )
        self.assertEqual(alias, "ggml-org/Qwen3-8B-GGUF:Q4_K_M")


if __name__ == "__main__":
    unittest.main()