tests/unit/adapters/test_llama_cpp.py

"""Tests for the llama.cpp adapter (MiniCPM-V 4.6)."""

from __future__ import annotations

import io
import os
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
from PIL import Image

from vivamais.adapters.models.llama_cpp import (
    VARIANT,
    LlamaCppBackend,
    LlamaCppVisionModel,
    _configure_cuda_libs,
    _flash_attn,
    _n_gpu_layers,
    _resolve_path,
)


def _dummy_jpeg_bytes() -> bytes:
    img = Image.new("RGB", (64, 64), color=(255, 255, 255))
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    return buf.getvalue()


class TestVariant:
    def test_defaults_to_vivamais_finetune_gguf(self) -> None:
        assert VARIANT.repo == "marinarosa/minicpmv4.6-vivamais-v1-GGUF"
        assert VARIANT.model_file == "MiniCPM-V-4.6-PTBR-v4-Q4_K_M.gguf"


class TestResolvePath:
    def test_uses_hint(self) -> None:
        assert _resolve_path("/custom/path.gguf", "X", "r", "f") == "/custom/path.gguf"

    def test_uses_env_var(self) -> None:
        with patch.dict(os.environ, {"TEST_VAR": "/env/path.gguf"}, clear=False):
            assert _resolve_path(None, "TEST_VAR", "r", "f") == "/env/path.gguf"

    def test_huggingface_download(self) -> None:
        with patch("vivamais.adapters.models.llama_cpp.hf_hub_download") as mock:
            mock.return_value = "/cache/model.gguf"
            result = _resolve_path(None, "MISSING_VAR", "repo", "file.gguf")
            mock.assert_called_once_with(repo_id="repo", filename="file.gguf")
            assert result == "/cache/model.gguf"


class TestCudaLibConfig:
    def test_configure_cuda_libs_noop_without_nvidia_packages(self) -> None:
        with patch("vivamais.adapters.models.llama_cpp._cuda_lib_dirs", return_value=[]):
            _configure_cuda_libs()

    def test_configure_cuda_libs_prepends_ld_library_path(self) -> None:
        with patch(
            "vivamais.adapters.models.llama_cpp._cuda_lib_dirs",
            return_value=["/fake/cuda/lib"],
        ):
            with patch.dict(os.environ, {}, clear=True):
                _configure_cuda_libs()
                assert os.environ["LD_LIBRARY_PATH"] == "/fake/cuda/lib"


class TestNGpuLayers:
    def test_cpu_when_build_lacks_gpu_backend(self) -> None:
        with patch.dict(os.environ, {}, clear=True):
            with patch(
                "vivamais.adapters.models.llama_cpp._gpu_offload_supported",
                return_value=False,
            ):
                assert _n_gpu_layers() == 0

    def test_full_offload_when_build_has_gpu_backend(self) -> None:
        with patch.dict(os.environ, {}, clear=True):
            with patch(
                "vivamais.adapters.models.llama_cpp._gpu_offload_supported",
                return_value=True,
            ):
                assert _n_gpu_layers() == -1

    def test_env_override(self) -> None:
        with patch.dict(os.environ, {"VIVAMAIS_N_GPU_LAYERS": "32"}, clear=False):
            assert _n_gpu_layers() == 32

    def test_scoped_env_wins_over_global(self) -> None:
        env = {"VIVAMAIS_N_GPU_LAYERS": "32", "VIVAMAIS_VISION_N_GPU_LAYERS": "0"}
        with patch.dict(os.environ, env, clear=False):
            assert _n_gpu_layers("VIVAMAIS_VISION_N_GPU_LAYERS") == 0

    def test_zero_gpu_forces_cpu_even_with_env_override(self) -> None:
        env = {"SPACES_ZERO_GPU": "true", "VIVAMAIS_N_GPU_LAYERS": "-1"}
        with patch.dict(os.environ, env, clear=False):
            assert _n_gpu_layers() == 0


class TestFlashAttn:
    def test_on_when_offloaded(self) -> None:
        with patch.dict(os.environ, {}, clear=True):
            assert _flash_attn(-1) is True

    def test_off_on_cpu(self) -> None:
        with patch.dict(os.environ, {}, clear=True):
            assert _flash_attn(0) is False

    def test_env_override(self) -> None:
        with patch.dict(os.environ, {"VIVAMAIS_FLASH_ATTN": "0"}, clear=False):
            assert _flash_attn(-1) is False


class TestBuildMessages:
    def test_build_messages_includes_system_prompt(self) -> None:
        backend = LlamaCppBackend()
        messages = backend._build_messages(b"img", "task")
        assert messages[0]["role"] == "system"
        assert messages[1]["role"] == "user"

    def test_chat_forwards_messages_and_response_format(self) -> None:
        backend = LlamaCppBackend()
        backend._ensure_loaded = MagicMock()  # type: ignore[method-assign]
        fake = MagicMock()
        fake.create_chat_completion.return_value = {
            "choices": [{"message": {"content": "resposta"}}]
        }
        backend._llama = fake
        messages = [{"role": "user", "content": "oi"}]

        out = backend.chat(messages, {"type": "json_object"})

        assert out == "resposta"
        _, kwargs = fake.create_chat_completion.call_args
        assert kwargs["messages"] == messages
        assert kwargs["response_format"] == {"type": "json_object"}


def _model_integration_enabled() -> bool:
    if os.environ.get("VIVAMAIS_RUN_MODEL_TESTS") != "1":
        return False
    return (
        os.environ.get("VIVAMAIS_MODEL_PATH") is not None
        or os.environ.get("VIVAMAIS_MMPROJ_PATH") is not None
        or (
            Path.home() / ".cache" / "huggingface" / "hub" / "models--openbmb--MiniCPM-V-4.6-gguf"
        ).exists()
    )


@pytest.mark.skipif(
    not _model_integration_enabled(),
    reason="Set VIVAMAIS_RUN_MODEL_TESTS=1 to run local GGUF integration tests",
)
class TestLlamaCppVisionModelIntegration:
    def test_extract_ocr(self) -> None:
        from vivamais.adapters.models.llama_cpp import LlamaCppTextModel

        backend = LlamaCppBackend(verbose=False)
        model = LlamaCppVisionModel(backend)
        result = model.extract_ocr(_dummy_jpeg_bytes())
        assert isinstance(result, str)
        _ = LlamaCppTextModel(backend)


class TestMiniCPM5TextModel:
    def _model_with_fake_llama(self, content: str):
        from vivamais.adapters.models.llama_cpp import MiniCPM5TextModel

        model = MiniCPM5TextModel()
        fake = MagicMock()
        fake.create_chat_completion.return_value = {"choices": [{"message": {"content": content}}]}
        model._llama = fake
        return model, fake

    def test_generate_json_constrains_grammar_and_temperature(self) -> None:
        schema = {"type": "object", "properties": {"ok": {"type": "boolean"}}}
        model, fake = self._model_with_fake_llama('{"ok": true}')
        out = model.generate_json([{"role": "user", "content": "x"}], schema=schema)
        assert out == '{"ok": true}'
        _, kwargs = fake.create_chat_completion.call_args
        assert kwargs["temperature"] == 0.0
        assert kwargs["response_format"] == {
            "type": "json_object",
            "schema": schema,
        }

    def test_generate_json_without_schema_uses_json_mode(self) -> None:
        model, fake = self._model_with_fake_llama("{}")
        model.generate_json([{"role": "user", "content": "x"}])
        _, kwargs = fake.create_chat_completion.call_args
        assert kwargs["response_format"] == {"type": "json_object"}

    def test_generate_strips_reasoning_block(self) -> None:
        model, _ = self._model_with_fake_llama("<think>plan</think>Ola!")
        assert model.generate([{"role": "user", "content": "q"}]) == "Ola!"

    def test_generate_uses_low_temperature_with_room_to_think(self) -> None:
        model, fake = self._model_with_fake_llama("Ola!")
        model.generate([{"role": "user", "content": "q"}])
        _, kwargs = fake.create_chat_completion.call_args
        assert kwargs["temperature"] == 0.2
        assert kwargs["max_tokens"] == 2048

    def test_default_context_window_is_8k(self) -> None:
        from vivamais.adapters.models.llama_cpp import MiniCPM5TextModel

        with patch.dict(os.environ, {}, clear=True):
            assert MiniCPM5TextModel()._n_ctx == 8192

    def test_context_window_env_override(self) -> None:
        from vivamais.adapters.models.llama_cpp import MiniCPM5TextModel

        with patch.dict(os.environ, {"VIVAMAIS_QA_N_CTX": "4096"}, clear=False):
            assert MiniCPM5TextModel()._n_ctx == 4096

    def _model_with_fake_stream(self, pieces: list[str]):
        from vivamais.adapters.models.llama_cpp import MiniCPM5TextModel

        model = MiniCPM5TextModel()
        fake = MagicMock()
        fake.create_chat_completion.return_value = iter(
            [{"choices": [{"delta": {"content": p}}]} for p in pieces]
        )
        model._llama = fake
        return model

    def test_stream_filters_think_block_split_across_chunks(self) -> None:
        model = self._model_with_fake_stream(
            ["<th", "ink>plano ", "secreto</th", "ink>Ola", " Maria"]
        )
        chunks = list(model.stream([{"role": "user", "content": "q"}]))
        assert "".join(chunks) == "Ola Maria"
        assert all("think" not in c for c in chunks)

    def test_stream_passes_through_plain_answer(self) -> None:
        model = self._model_with_fake_stream(["Ola ", "Maria"])
        assert "".join(model.stream([{"role": "user", "content": "q"}])) == "Ola Maria"