neuron-tai/tests/test_gguf_backend.py

"""Tests for the GGUF backend adapter and recipe-gated startup seam."""

from __future__ import annotations

from types import SimpleNamespace

from meshnet_node.gguf_backend import GgufNodeBackend, build_gguf_backend
from meshnet_node.model_backend import TailTokenResult, TensorPayload
from meshnet_node.recipe_manifest import DEFAULT_RECIPE_ID, load_recipe_manifest
from meshnet_node.startup import _gguf_backend_for_recipe


class _RecordingTransport:
    def __init__(self) -> None:
        self.calls: list[tuple[str, tuple, dict]] = []

    def encode_prompt(self, prompt: str, session_id: str | None = None):
        self.calls.append(("encode_prompt", (prompt, session_id), {}))
        return TensorPayload(
            body=b"\x00" * 16,
            shape=[1, 2, 4],
            attention_mask_header=None,
            position_ids_header=None,
        )

    def encode_next_token(self, token_id: int, session_id: str):
        self.calls.append(("encode_next_token", (token_id, session_id), {}))
        return TensorPayload(
            body=b"\x00" * 8,
            shape=[1, 1, 4],
            attention_mask_header=None,
            position_ids_header=None,
            past_len=2,
        )

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        *,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ):
        self.calls.append(
            (
                "forward_bytes",
                (body, tuple(shape), attention_mask_header, position_ids_header),
                {
                    "start_layer": start_layer,
                    "session_id": session_id,
                    "cache_mode": cache_mode,
                    "past_len": past_len,
                },
            )
        )
        if cache_mode == "decode":
            return TailTokenResult(text=" done", token_id=17)
        return TensorPayload(
            body=b"\x00" * 16,
            shape=[1, 2, 4],
            attention_mask_header=attention_mask_header,
            position_ids_header=position_ids_header,
            past_len=past_len,
        )

    def decode_tail_token(self, hidden_states):
        self.calls.append(("decode_tail_token", (hidden_states.shape,), {}))
        return TailTokenResult(text=" tail", token_id=19)

    def generate_text(self, messages, max_new_tokens=5120, temperature=1.0, top_p=1.0):
        self.calls.append(("generate_text", (tuple(messages), max_new_tokens, temperature, top_p), {}))
        return "ok"

    def generate_text_streaming(self, messages, max_new_tokens=5120, temperature=1.0, top_p=1.0):
        self.calls.append(("generate_text_streaming", (tuple(messages), max_new_tokens, temperature, top_p), {}))
        yield "ok"

    def count_prompt_tokens(self, messages):
        self.calls.append(("count_prompt_tokens", (tuple(messages),), {}))
        return 3

    def count_text_tokens(self, text):
        self.calls.append(("count_text_tokens", (text,), {}))
        return 2

    def eos_token_ids(self):
        self.calls.append(("eos_token_ids", (), {}))
        return [19]

    def release_session(self, session_id: str) -> None:
        self.calls.append(("release_session", (session_id,), {}))


def test_build_gguf_backend_delegates_to_transport():
    transport = _RecordingTransport()
    backend = build_gguf_backend(
        model_id="meshnet/native-model",
        shard_start=0,
        shard_end=1,
        quantization="bfloat16",
        transport=transport,
        total_layers=2,
        device_type="cpu",
    )

    assert isinstance(backend, GgufNodeBackend)
    assert backend.backend_id == "llama.cpp"
    assert backend.is_head is True
    assert backend.is_tail is True
    assert backend.model.config.to_dict()["architecture_adapter"] == "dense-llama"
    assert backend.loaded_tensor_names[0] == "blk.0.weight"

    prompt = backend.encode_prompt("hello", session_id="session-1")
    assert prompt.shape == [1, 2, 4]

    decode = backend.forward_bytes(
        b"\x00" * 16,
        [1, 2, 4],
        None,
        None,
        session_id="session-1",
        cache_mode="decode",
        past_len=2,
    )
    assert isinstance(decode, TailTokenResult)
    assert decode.token_id == 17

    backend.release_session("session-1")

    assert [call[0] for call in transport.calls] == [
        "encode_prompt",
        "forward_bytes",
        "release_session",
    ]
    assert transport.calls[0][1] == ("hello", "session-1")
    assert transport.calls[1][2]["cache_mode"] == "decode"
    assert transport.calls[1][2]["past_len"] == 2


def test_recipe_gates_native_backend_selection(monkeypatch):
    manifest = load_recipe_manifest()
    torch_recipe = manifest.require(DEFAULT_RECIPE_ID)
    native_recipe = manifest.require("llama-cpp-native")

    sentinel_backend = object()
    calls: list[dict] = []

    def fake_build_gguf_backend(**kwargs):
        calls.append(kwargs)
        return sentinel_backend

    monkeypatch.setattr(
        "meshnet_node.startup.build_gguf_backend",
        fake_build_gguf_backend,
    )

    assert _gguf_backend_for_recipe(
        torch_recipe,
        model_id="meshnet/native-model",
        shard_start=0,
        shard_end=1,
        quantization="bfloat16",
        total_layers=2,
        device="cpu",
    ) is None

    backend = _gguf_backend_for_recipe(
        native_recipe,
        model_id="meshnet/native-model",
        shard_start=0,
        shard_end=1,
        quantization="bfloat16",
        total_layers=2,
        device="cpu",
    )

    assert backend is sentinel_backend
    assert calls[0]["model_id"] == "meshnet/native-model"
    assert calls[0]["shard_start"] == 0
    assert calls[0]["shard_end"] == 1
    assert calls[0]["quantization"] == "bfloat16"
    assert calls[0]["total_layers"] == 2