lmstudiofixed/tests/test_lmstudio_wrapper_generator.py

import importlib.util
import json
import struct
import tempfile
import unittest
from unittest import mock
from pathlib import Path


MODULE_PATH = Path(__file__).resolve().parent.parent / "tools" / "lmstudio" / "wrapper_generator.py"
if MODULE_PATH.exists():
    spec = importlib.util.spec_from_file_location("wrapper_generator", MODULE_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    spec.loader.exec_module(module)
    gen = module
else:  # pragma: no cover
    gen = None


class WrapperGeneratorTests(unittest.TestCase):
    def setUp(self):
        if gen is None:
            self.fail("tools.lmstudio.wrapper_generator is not available")

    def _write_gguf_metadata(self, path: Path, metadata: dict[str, object]) -> None:
        type_string = 8

        def write_string(value: str) -> bytes:
            raw = value.encode("utf-8")
            return struct.pack("<Q", len(raw)) + raw

        def write_value(value: object) -> tuple[int, bytes]:
            if isinstance(value, str):
                return type_string, write_string(value)
            raise TypeError(f"Unsupported GGUF test metadata value: {value!r}")

        payload = bytearray()
        payload.extend(b"GGUF")
        payload.extend(struct.pack("<I", 3))
        payload.extend(struct.pack("<Q", 0))
        payload.extend(struct.pack("<Q", len(metadata)))

        for key, value in metadata.items():
            value_type, raw_value = write_value(value)
            payload.extend(write_string(key))
            payload.extend(struct.pack("<I", value_type))
            payload.extend(raw_value)

        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(payload)

    def test_detect_capabilities_from_cached_metadata(self):
        metadata = {
            "general.architecture": "qwen35moe",
            "general.name": "Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive",
            "general.basename": "Qwen3.6-35B-A3B",
            "tokenizer.chat_template": (
                "{%- if tools %}<tools>{%- endif %}"
                "{%- if add_generation_prompt %}"
                "{%- if enable_thinking is defined and enable_thinking is false %}"
                "{{- '<think>\\n\\n</think>\\n\\n' }}"
                "{%- else %}{{- '<think>\\n' }}{%- endif %}{%- endif %}"
                "<|vision_start|><|image_pad|><|vision_end|>"
            ),
            "qwen2.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)

        self.assertEqual(profile["architecture"], "qwen35moe")
        self.assertEqual(profile["basename_slug"], "qwen3.6-35b-a3b")
        self.assertEqual(profile["params_string"], "35B")
        self.assertEqual(profile["context_length"], 262144)
        self.assertTrue(profile["capabilities"]["reasoning"])
        self.assertTrue(profile["capabilities"]["vision"])
        self.assertTrue(profile["capabilities"]["tools"])
        self.assertTrue(profile["needs_enable_thinking"])
        self.assertTrue(profile["needs_preserve_thinking"])

    def test_prefers_internal_cache_before_touching_gguf(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            gguf_path = root / "models" / "publisher-a" / "Fancy-Model-GGUF" / "model.gguf"
            self._write_gguf_metadata(
                gguf_path,
                {
                    "general.architecture": "qwen35",
                    "general.name": "Wrong Fallback Name",
                    "general.basename": "Wrong Fallback Name",
                    "tokenizer.chat_template": "",
                    "qwen2.context_length": "4096",
                },
            )
            cache_path = root / ".internal" / "gguf-metadata-cache.json"
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            cache_payload = {
                "json": {
                    "map": [
                        [
                            gguf_path.resolve().as_posix(),
                            {
                                "mtimeMs": gguf_path.stat().st_mtime * 1000,
                                "metadata": {
                                    "arch": "qwen35moe",
                                    "name": "Cached Fancy Model",
                                    "chatTemplate": "<tools>{{ '<think>\\n' }}",
                                    "parameters": "35B-A3B",
                                    "contextLength": 262144,
                                },
                            },
                        ]
                    ]
                }
            }
            cache_path.write_text(json.dumps(cache_payload, ensure_ascii=False), encoding="utf-8")

            with mock.patch.object(gen, "read_gguf_metadata_streaming", side_effect=AssertionError("stream parser should not run on cache hit")):
                metadata = gen.resolve_gguf_metadata(
                    gguf_path=gguf_path,
                    cache_index=gen.load_gguf_metadata_cache(cache_path),
                )

            self.assertEqual(metadata["general.architecture"], "qwen35moe")
            self.assertEqual(metadata["general.name"], "Cached Fancy Model")
            self.assertEqual(metadata["general.parameter_count_label"], "35B-A3B")
            self.assertEqual(metadata["general.context_length"], "262144")

    def test_stream_parser_is_used_when_cache_missing(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            gguf_path = Path(tmpdir) / "model.gguf"
            self._write_gguf_metadata(
                gguf_path,
                {
                    "general.architecture": "gemma4",
                    "general.name": "Streaming Model",
                    "general.basename": "Streaming Model",
                    "tokenizer.chat_template": "{{ '<think>\\n' }}",
                    "qwen2.context_length": "8192",
                },
            )

            with mock.patch.object(Path, "read_bytes", side_effect=AssertionError("full file read should not be used")):
                metadata = gen.resolve_gguf_metadata(
                    gguf_path=gguf_path,
                    cache_index={},
                )

            self.assertEqual(metadata["general.architecture"], "gemma4")
            self.assertEqual(metadata["general.name"], "Streaming Model")

    def test_skips_mmproj_when_filename_contains_mmproj_not_only_prefix(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            gguf_path = (
                models_root
                / "publisher-a"
                / "Vision-Model-GGUF"
                / "Vision-Model-it.mmproj-Q8_0.gguf"
            )
            self._write_gguf_metadata(
                gguf_path,
                {
                    "general.architecture": "qwen3vl",
                    "general.name": "Vision Model",
                    "general.basename": "Vision Model",
                    "tokenizer.chat_template": "<|vision_start|><|image_pad|><|vision_end|>",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=True,
            )

            self.assertEqual(summary["generated"], [])
            self.assertEqual(len(summary["skipped"]), 1)
            self.assertEqual(summary["skipped"][0]["reason"], "mmproj")

    def test_sibling_mmproj_marks_model_as_vision_capable(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"

            model_dir = models_root / "publisher-a" / "Gemma4-Model-GGUF"
            self._write_gguf_metadata(
                model_dir / "Gemma4-Model-Q4_K_P.gguf",
                {
                    "general.architecture": "gemma4",
                    "general.name": "Gemma4 Model",
                    "general.basename": "Gemma4 Model",
                    "tokenizer.chat_template": "{% if enable_thinking is defined and enable_thinking is false %}{{ '' }}{% endif %}",
                    "general.context_length": "262144",
                },
            )
            self._write_gguf_metadata(
                model_dir / "mmproj-Gemma4-Model-f16.gguf",
                {
                    "general.architecture": "clip",
                    "general.name": "Gemma4 mmproj",
                    "general.basename": "Gemma4 mmproj",
                    "tokenizer.chat_template": "",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=True,
            )

            self.assertEqual(len(summary["generated"]), 1)
            generated = summary["generated"][0]
            self.assertTrue(generated["capabilities"]["vision"])

    def test_dry_run_scans_models_root_and_skips_lmstudio_community(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"

            self._write_gguf_metadata(
                models_root / "publisher-a" / "Fancy-Model-GGUF" / "model.gguf",
                {
                    "general.architecture": "qwen35",
                    "general.name": "Fancy Model",
                    "general.basename": "Fancy Model",
                    "tokenizer.chat_template": "<tools>{% if enable_thinking is defined and enable_thinking is false %}{{ '<think>\\n\\n</think>\\n\\n' }}{% else %}{{ '<think>\\n' }}{% endif %}<|vision_start|><|image_pad|><|vision_end|>",
                    "qwen2.context_length": "262144",
                },
            )
            self._write_gguf_metadata(
                models_root / "lmstudio-community" / "Qwen3.5-9B-GGUF" / "model.gguf",
                {
                    "general.architecture": "qwen35",
                    "general.name": "Qwen3.5-9B",
                    "general.basename": "Qwen3.5-9B",
                    "tokenizer.chat_template": "{{ '<think>\\n' }}",
                    "qwen2.context_length": "262144",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=True,
            )

            self.assertTrue(summary["dry_run"])
            self.assertEqual(len(summary["generated"]), 1)
            self.assertEqual(len(summary["skipped"]), 1)

            generated = summary["generated"][0]
            self.assertEqual(generated["publisher"], "publisher-a")
            self.assertEqual(generated["model_key"], "publisher-a/Fancy-Model-GGUF/model.gguf")
            self.assertEqual(generated["output_dir"], str(output_root / "publisher-a" / "fancy-model"))
            self.assertTrue(generated["capabilities"]["reasoning"])
            self.assertTrue(generated["capabilities"]["vision"])
            self.assertTrue(generated["capabilities"]["tools"])

            skipped = summary["skipped"][0]
            self.assertEqual(skipped["reason"], "lmstudio-community")

    def test_generate_virtual_model_writes_official_style_files(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            model_index_cache = root / ".internal" / "model-index-cache.json"
            model_index_cache.parent.mkdir(parents=True, exist_ok=True)
            model_index_cache.write_text("cached", encoding="utf-8")

            self._write_gguf_metadata(
                models_root / "publisher-b" / "Fancy-Moe-Model-GGUF" / "Fancy-Moe-Model-Q4_K_M.gguf",
                {
                    "general.architecture": "qwen35moe",
                    "general.name": "Fancy Moe Model",
                    "general.basename": "Fancy Moe Model",
                    "tokenizer.chat_template": (
                        "{%- if tools %}<tools>{%- endif %}"
                        "{%- if add_generation_prompt %}"
                        "{%- if enable_thinking is defined and enable_thinking is false %}"
                        "{{- '<think>\\n\\n</think>\\n\\n' }}"
                        "{%- else %}{{- '<think>\\n' }}{%- endif %}{%- endif %}"
                        "<|vision_start|><|image_pad|><|vision_end|>"
                    ),
                    "qwen2.context_length": "262144",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=False,
                model_index_cache_path=model_index_cache,
            )

            self.assertEqual(len(summary["generated"]), 1)
            output_dir = output_root / "publisher-b" / "fancy-moe-model"
            self.assertTrue(output_dir.exists())

            model_yaml = (output_dir / "model.yaml").read_text(encoding="utf-8")
            manifest = json.loads((output_dir / "manifest.json").read_text(encoding="utf-8"))

            self.assertIn("model: publisher-b/fancy-moe-model", model_yaml)
            self.assertIn("base:", model_yaml)
            self.assertIn("key: publisher-b/Fancy-Moe-Model-GGUF/Fancy-Moe-Model-Q4_K_M.gguf", model_yaml)
            self.assertIn("sources:", model_yaml)
            self.assertIn("reasoning: true", model_yaml)
            self.assertIn("vision: true", model_yaml)
            self.assertIn("trainedForToolUse: true", model_yaml)
            self.assertIn("key: enableThinking", model_yaml)
            self.assertIn("key: preserveThinking", model_yaml)

            self.assertEqual(manifest["owner"], "publisher-b")
            self.assertEqual(manifest["name"], "fancy-moe-model")
            self.assertEqual(
                manifest["dependencies"][0]["modelKeys"],
                ["publisher-b/Fancy-Moe-Model-GGUF/Fancy-Moe-Model-Q4_K_M.gguf"],
            )
            self.assertIn("sources", manifest["dependencies"][0])
            self.assertEqual(manifest["dependencies"][0]["sources"], [])

    def test_dry_run_does_not_delete_model_index_cache(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            model_index_cache = root / ".internal" / "model-index-cache.json"
            model_index_cache.parent.mkdir(parents=True, exist_ok=True)
            model_index_cache.write_text("cached", encoding="utf-8")

            self._write_gguf_metadata(
                models_root / "publisher-a" / "Fancy-Model-GGUF" / "model.gguf",
                {
                    "general.architecture": "qwen35",
                    "general.name": "Fancy Model",
                    "general.basename": "Fancy Model",
                    "tokenizer.chat_template": "{% if enable_thinking is defined and enable_thinking is false %}{{ '' }}{% endif %}",
                    "qwen2.context_length": "262144",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=True,
                model_index_cache_path=model_index_cache,
            )

            self.assertFalse(summary["deleted_model_index_cache"])
            self.assertTrue(model_index_cache.exists())

    def test_generation_deletes_internal_model_index_cache_after_write(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            model_index_cache = root / ".internal" / "model-index-cache.json"
            model_index_cache.parent.mkdir(parents=True, exist_ok=True)
            model_index_cache.write_text("cached", encoding="utf-8")

            self._write_gguf_metadata(
                models_root / "publisher-a" / "Fancy-Model-GGUF" / "model.gguf",
                {
                    "general.architecture": "qwen35",
                    "general.name": "Fancy Model",
                    "general.basename": "Fancy Model",
                    "tokenizer.chat_template": "{% if enable_thinking is defined and enable_thinking is false %}{{ '' }}{% endif %}",
                    "qwen2.context_length": "262144",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=False,
                model_index_cache_path=model_index_cache,
            )

            self.assertTrue(summary["deleted_model_index_cache"])
            self.assertEqual(summary["model_index_cache_path"], str(model_index_cache))
            self.assertFalse(model_index_cache.exists())

    def test_owner_is_normalized_for_lmstudio_manifest_and_model_id(self):
        metadata = {
            "general.architecture": "qwen35moe",
            "general.name": "Fancy Moe Model",
            "general.basename": "Fancy Moe Model",
            "tokenizer.chat_template": (
                "{% if enable_thinking is defined and enable_thinking is false %}"
                "{{ '<think>\\n\\n</think>\\n\\n' }}"
                "{% endif %}"
            ),
            "qwen2.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)
        files = gen.build_virtual_model_files(
            publisher="HauhauCS",
            slug="fancy-moe-model",
            model_key="HauhauCS/Fancy-Moe-Model-GGUF/Fancy-Moe-Model-Q4_K_M.gguf",
            profile=profile,
        )
        manifest = json.loads(files["manifest.json"])

        self.assertIn("model: hauhaucs/fancy-moe-model", files["model.yaml"])
        self.assertEqual(manifest["owner"], "hauhaucs")
        self.assertEqual(manifest["dependencies"][0]["sources"], [])

    def test_collect_job_uses_normalized_owner_for_output_dir_and_model_id(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            gguf_path = (
                models_root
                / "HauhauCS"
                / "Fancy-Moe-Model-GGUF"
                / "Fancy-Moe-Model-Q4_K_M.gguf"
            )
            self._write_gguf_metadata(
                gguf_path,
                {
                    "general.architecture": "qwen35moe",
                    "general.name": "Fancy Moe Model",
                    "general.basename": "Fancy Moe Model",
                    "tokenizer.chat_template": (
                        "{% if enable_thinking is defined and enable_thinking is false %}"
                        "{{ '<think>\\n\\n</think>\\n\\n' }}"
                        "{% endif %}"
                    ),
                    "qwen2.context_length": "262144",
                },
            )

            job = gen.collect_virtual_model_job(
                models_root=models_root,
                output_root=output_root,
                gguf_path=gguf_path,
                cache_index={},
            )

            self.assertEqual(job["publisher"], "hauhaucs")
            self.assertEqual(
                job["output_dir"],
                str(output_root / "hauhaucs" / "fancy-moe-model"),
            )

    def test_nothink_variant_disables_reasoning_and_custom_field(self):
        metadata = {
            "general.architecture": "qwen35",
            "general.name": "Generic Model NoThink",
            "general.basename": "Generic Model",
            "tokenizer.chat_template": "<tools><|vision_start|><|image_pad|><|vision_end|>",
            "qwen2.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)
        files = gen.build_virtual_model_files(
            publisher="publisher-c",
            slug="generic-model-nothink",
            model_key="publisher-c/Generic-Model-GGUF/Generic-Model-NoThink-Q4_K_M.gguf",
            profile=profile,
        )

        self.assertIn("reasoning: false", files["model.yaml"])
        self.assertNotIn("key: enableThinking", files["model.yaml"])
        self.assertNotIn("key: preserveThinking", files["model.yaml"])

    def test_nothink_directory_forces_reasoning_off_even_if_template_contains_think(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            models_root = root / "models"
            output_root = root / "hub" / "models"
            model_index_cache = root / ".internal" / "model-index-cache.json"
            model_index_cache.parent.mkdir(parents=True, exist_ok=True)
            model_index_cache.write_text("cached", encoding="utf-8")
            gguf_path = (
                models_root
                / "publisher-b"
                / "Generic-Model-NoThink-GGUF"
                / "Generic-Model-NoThink-Q4_K_M.gguf"
            )
            self._write_gguf_metadata(
                gguf_path,
                {
                    "general.architecture": "qwen35",
                    "general.name": "Generic Model NoThink",
                    "general.basename": "Generic Model",
                    "tokenizer.chat_template": "{{ '<think>\\n\\n</think>\\n\\n' }}<tools>",
                    "qwen2.context_length": "262144",
                },
            )

            summary = gen.generate_virtual_models_for_directory(
                models_root=models_root,
                output_root=output_root,
                dry_run=False,
                model_index_cache_path=model_index_cache,
            )

            self.assertEqual(len(summary["generated"]), 1)
            output_dir = output_root / "publisher-b" / "generic-model-nothink"
            model_yaml = (output_dir / "model.yaml").read_text(encoding="utf-8")
            self.assertIn("reasoning: false", model_yaml)
            self.assertNotIn("key: enableThinking", model_yaml)

    def test_static_think_block_without_enable_thinking_is_not_reasoning_capable(self):
        metadata = {
            "general.architecture": "qwen35",
            "general.name": "Static Think Model",
            "general.basename": "Static Think Model",
            "tokenizer.chat_template": "{{ '<think>\\n\\n</think>\\n\\n' }}<tools>",
            "qwen2.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)

        self.assertFalse(profile["capabilities"]["reasoning"])
        self.assertFalse(profile["needs_enable_thinking"])

    def test_preserve_thinking_without_enable_thinking_is_not_reasoning_capable(self):
        metadata = {
            "general.architecture": "qwen35moe",
            "general.name": "Preserve Only Model",
            "general.basename": "Preserve Only Model",
            "tokenizer.chat_template": (
                "{%- if preserve_thinking is defined and preserve_thinking is true %}"
                "{{ '<think>\\nold reasoning\\n</think>\\n\\n' }}"
                "{%- endif %}<tools>"
            ),
            "qwen2.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)
        files = gen.build_virtual_model_files(
            publisher="publisher-d",
            slug="preserve-only-model",
            model_key="publisher-d/Preserve-Only-Model-GGUF/model.gguf",
            profile=profile,
        )

        self.assertFalse(profile["capabilities"]["reasoning"])
        self.assertFalse(profile["needs_enable_thinking"])
        self.assertFalse(profile["needs_preserve_thinking"])
        self.assertIn("reasoning: false", files["model.yaml"])
        self.assertNotIn("key: enableThinking", files["model.yaml"])
        self.assertNotIn("key: preserveThinking", files["model.yaml"])

    def test_gemma4_pipe_style_tool_tokens_are_detected(self):
        metadata = {
            "general.architecture": "gemma4",
            "general.name": "Gemma4 Tool Model",
            "general.basename": "Gemma4 Tool Model",
            "tokenizer.chat_template": (
                "{% if tools %}{{ '<|tool>' }}{% endif %}"
                "{{ '<|tool_call>' }}"
                "{{ '<|tool_response>' }}"
            ),
            "general.context_length": "262144",
        }

        profile = gen.build_virtual_model_profile(metadata)

        self.assertTrue(profile["capabilities"]["tools"])

    def test_profile_inference_uses_metadata_and_directory_augments_vision(self):
        metadata = {
            "general.architecture": "gemma4",
            "general.name": "Augmented Vision Model",
            "general.basename": "Augmented Vision Model",
            "tokenizer.chat_template": "{% if enable_thinking is defined and enable_thinking is false %}{{ '' }}{% endif %}",
            "general.context_length": "262144",
        }

        profile_without_mmproj = gen.build_virtual_model_profile(metadata)
        profile_with_mmproj = gen.build_virtual_model_profile(
            metadata,
            has_sibling_mmproj=True,
        )

        self.assertFalse(profile_without_mmproj["capabilities"]["vision"])
        self.assertTrue(profile_with_mmproj["capabilities"]["vision"])
        self.assertEqual(
            profile_without_mmproj["capabilities"]["tools"],
            profile_with_mmproj["capabilities"]["tools"],
        )


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