tests/test_local_llm.py

import json
from typing import Any

from local_llm import (
    LocalChatClient,
    LocalCompletionClient,
    LocalJsonChatClient,
    MLXChatClient,
    MiniCPMTransformersChatClient,
    NemotronTransformersChatClient,
    chat_messages,
    chat_payload,
    completion_payload,
    decode_generated_text,
    json_chat_payload,
    minicpm_text_prompt,
    nemotron_prompt,
    parse_chat_response,
    parse_completion_response,
    tensor_to_model_device,
    token_length,
    transformers_model_kwargs,
)


class FakeResponse:
    # Initialize a fake HTTP response body.
    def __init__(self, raw: dict[str, Any]) -> None:
        self.raw = raw

    # Enter the fake response context.
    def __enter__(self) -> "FakeResponse":
        return self

    # Exit the fake response context.
    def __exit__(self, *_: object) -> None:
        return None

    # Return encoded response bytes.
    def read(self) -> bytes:
        return json.dumps(self.raw).encode("utf-8")


class FakeInputs(list):
    # Initialize tensor-like fake inputs.
    def __init__(self) -> None:
        super().__init__([[1, 2]])
        self.shape = (1, 2)
        self.moved_to: str | None = None

    # Capture device moves.
    def to(self, device: str) -> "FakeInputs":
        self.moved_to = device
        return self


class FakeTokenizer:
    # Initialize a fake tokenizer.
    def __init__(self) -> None:
        self.messages: list[dict[str, str]] = []

    # Capture chat template messages.
    def apply_chat_template(self, messages: list[dict[str, str]], **_: Any) -> FakeInputs:
        self.messages = messages
        return FakeInputs()

    # Decode generated token ids.
    def decode(self, tokens: list[int], **_: Any) -> str:
        return f"decoded:{tokens}"


class FakeModel:
    device = "cuda"

    # No-op device move (ZeroGPU does the real one inside the @gpu call).
    def to(self, device: str) -> "FakeModel":
        return self

    # Return a generated token sequence with prompt prefix included.
    def generate(self, inputs: FakeInputs, **kwargs: Any) -> list[list[int]]:
        self.inputs = inputs
        self.kwargs = kwargs
        return [[1, 2, 3, 4]]


class FakeMiniCPMModel:
    # No-op device move (ZeroGPU does the real one inside the @gpu call).
    def to(self, device: str) -> "FakeMiniCPMModel":
        return self

    # Capture MiniCPM chat kwargs.
    def chat(self, **kwargs: Any) -> str:
        self.kwargs = kwargs
        return "mini"


class FakeMLXGenerate:
    # Capture MLX generation arguments.
    def __call__(self, *args: Any, **kwargs: Any) -> str:
        self.args = args
        self.kwargs = kwargs
        return "mlx"


class FakeTorch:
    bfloat16 = "bf16"

    class cuda:
        # Report no CUDA for dtype fallback tests.
        @staticmethod
        def is_available() -> bool:
            return False


class FakeCudaTorch:
    bfloat16 = "bf16"

    class cuda:
        # Report CUDA for dtype selection tests.
        @staticmethod
        def is_available() -> bool:
            return True


# Verify chat payload shape matches OpenAI-compatible endpoints.
def test_chat_payload() -> None:
    payload = chat_payload("local", "sys", "user", 0.5)
    assert payload["model"] == "local"
    assert payload["messages"][0]["role"] == "system"
    assert payload["temperature"] == 0.5
    assert "max_tokens" not in payload
    assert chat_payload("local", "sys", "user", 0.5, 44)["max_tokens"] == 44
    assert chat_payload("local", "sys", "user", 0.5, enable_thinking=False)["chat_template_kwargs"] == {"enable_thinking": False}


# Verify completion payload shape matches OpenAI-compatible endpoints.
def test_completion_payload() -> None:
    payload = completion_payload("local", "prompt", 0.4, 99)
    assert payload == {"model": "local", "prompt": "prompt", "temperature": 0.4, "max_tokens": 99}


# Verify JSON chat payload requests JSON object output.
def test_json_chat_payload() -> None:
    payload = json_chat_payload("local", "sys", "user", 0.0, 55)
    assert payload["response_format"] == {"type": "json_object"}
    assert payload["max_tokens"] == 55


# Verify chat message payloads use standard roles.
def test_chat_messages() -> None:
    assert chat_messages("sys", "user") == [{"role": "system", "content": "sys"}, {"role": "user", "content": "user"}]


# Verify Nemotron prompt markers match the documented template.
def test_nemotron_prompt() -> None:
    assert nemotron_prompt("sys", "user") == "<extra_id_0>System\nsys\n\n<extra_id_1>User\nuser\n<extra_id_1>Assistant\n"


# Verify MiniCPM receives system and user text in one chat prompt.
def test_minicpm_text_prompt() -> None:
    assert minicpm_text_prompt("sys", "user") == "System:\nsys\n\nUser:\nuser"


# Verify chat response parsing returns assistant content.
def test_parse_chat_response() -> None:
    assert parse_chat_response({"choices": [{"message": {"content": "ok"}}]}) == "ok"


# Verify completion response parsing supports common local shapes.
def test_parse_completion_response() -> None:
    assert parse_completion_response({"content": "llama"}) == "llama"
    assert parse_completion_response({"choices": [{"text": "done"}]}) == "done"
    assert parse_completion_response({"choices": [{"message": {"content": "chat"}}]}) == "chat"


# Verify local chat client posts JSON and returns content.
def test_local_chat_client(monkeypatch: Any) -> None:
    calls: list[Any] = []

    # Capture the request and return a fake response.
    def fake_urlopen(req: Any, timeout: int) -> FakeResponse:
        calls.append((req, timeout))
        return FakeResponse({"choices": [{"message": {"content": "done"}}]})

    monkeypatch.setattr("local_llm.request.urlopen", fake_urlopen)
    client = LocalChatClient("http://localhost/v1/chat/completions", "local", timeout_seconds=3)
    assert client.complete("sys", "user") == "done"
    assert calls[0][1] == 3
    payload = json.loads(calls[0][0].data.decode("utf-8"))
    assert payload["model"] == "local"
    assert payload["max_tokens"] == 256
    assert "chat_template_kwargs" not in payload


# Verify local JSON chat client posts JSON mode payloads.
def test_local_json_chat_client(monkeypatch: Any) -> None:
    calls: list[Any] = []

    # Capture the request and return a fake response.
    def fake_urlopen(req: Any, timeout: int) -> FakeResponse:
        calls.append((req, timeout))
        return FakeResponse({"choices": [{"message": {"content": "{\"ok\": true}"}}]})

    monkeypatch.setattr("local_llm.request.urlopen", fake_urlopen)
    client = LocalJsonChatClient("http://localhost/v1/chat/completions", "local", timeout_seconds=5, max_tokens=22)
    assert client.complete("sys", "user") == "{\"ok\": true}"
    payload = json.loads(calls[0][0].data.decode("utf-8"))
    assert payload["response_format"] == {"type": "json_object"}
    assert payload["max_tokens"] == 22
    assert calls[0][1] == 5


# Verify local completion client posts a rendered prompt and returns content.
def test_local_completion_client(monkeypatch: Any) -> None:
    calls: list[Any] = []

    # Capture the request and return a fake response.
    def fake_urlopen(req: Any, timeout: int) -> FakeResponse:
        calls.append((req, timeout))
        return FakeResponse({"choices": [{"text": "done"}]})

    monkeypatch.setattr("local_llm.request.urlopen", fake_urlopen)
    client = LocalCompletionClient("http://localhost/v1/completions", "local", nemotron_prompt, timeout_seconds=4, max_tokens=12)
    assert client.complete("sys", "user") == "done"
    payload = json.loads(calls[0][0].data.decode("utf-8"))
    assert payload["prompt"] == nemotron_prompt("sys", "user")
    assert payload["max_tokens"] == 12
    assert calls[0][1] == 4


# Verify tensor device moves are optional.
def test_tensor_to_model_device() -> None:
    inputs = FakeInputs()
    assert tensor_to_model_device(inputs, FakeModel()).moved_to == "cuda"
    assert tensor_to_model_device([1, 2], object()) == [1, 2]


# Verify token length handles tensor-like and list inputs.
def test_token_length() -> None:
    assert token_length(FakeInputs()) == 2
    assert token_length([[1, 2, 3]]) == 3
    assert token_length([1, 2]) == 2


# Verify generated decoding removes the prompt prefix.
def test_decode_generated_text() -> None:
    assert decode_generated_text(FakeTokenizer(), FakeInputs(), [[1, 2, 3]]) == "decoded:[3]"


# Verify Nemotron Transformers client uses tokenizer chat templates.
def test_nemotron_transformers_chat_client() -> None:
    tokenizer = FakeTokenizer()
    model = FakeModel()
    client = NemotronTransformersChatClient(model, tokenizer, max_new_tokens=7, temperature=0.0)
    assert client.complete("sys", "user") == "decoded:[3, 4]"
    assert tokenizer.messages[0]["role"] == "system"
    assert model.inputs.moved_to == "cuda"
    assert model.kwargs["max_new_tokens"] == 7
    assert model.kwargs["do_sample"] is False


# Verify MLX chat client renders prompts and calls generate.
def test_mlx_chat_client() -> None:
    generate = FakeMLXGenerate()
    client = MLXChatClient("model", "tok", nemotron_prompt, generate, max_tokens=11)
    assert client.complete("sys", "user") == "mlx"
    assert generate.args == ("model", "tok", nemotron_prompt("sys", "user"))
    assert generate.kwargs == {"verbose": False, "max_tokens": 11}


# Verify MiniCPM Transformers client uses model.chat with no image.
def test_minicpm_transformers_chat_client() -> None:
    model = FakeMiniCPMModel()
    client = MiniCPMTransformersChatClient(model, tokenizer="tok", max_new_tokens=77, temperature=0.7)
    assert client.complete("sys", "user") == "mini"
    assert model.kwargs["image"] is None
    assert model.kwargs["tokenizer"] == "tok"
    assert model.kwargs["system_prompt"] == "sys"
    assert model.kwargs["sampling"] is True
    assert model.kwargs["temperature"] == 0.7
    assert model.kwargs["max_new_tokens"] == 77
    assert model.kwargs["msgs"][0]["content"] == "user"
    # temperature 0 falls back to deterministic decoding
    MiniCPMTransformersChatClient(model, tokenizer="tok", temperature=0.0).complete("s", "u")
    assert model.kwargs["sampling"] is False


# Verify Transformers model kwargs use auto dtype.
def test_transformers_model_kwargs(monkeypatch: Any) -> None:
    monkeypatch.setenv("TABRAS_MODEL_OFFLOAD", "/tmp/custom-offload")
    assert transformers_model_kwargs()["torch_dtype"] == "auto"
    if "device_map" in transformers_model_kwargs():
        assert transformers_model_kwargs()["offload_folder"] == "/tmp/custom-offload"


# Verify MiniCPM dtype follows CUDA availability.
def test_local_torch_dtype() -> None:
    from local_llm import local_torch_dtype

    assert local_torch_dtype(FakeTorch) == "auto"
    assert local_torch_dtype(FakeCudaTorch) == "bf16"