tests_python/test_local_model.py

from langchain_core.messages import HumanMessage, SystemMessage

from arena.local_model import (
    CachedLocalChatModel,
    _bucket_cache_key,
    clear_local_model_cache,
    local_model_cache_size,
    _cached_pipeline,
    _local_generation_duration,
    _model_load_sources,
    _messages_to_prompt,
)


class FakePipeline:
    tokenizer = None

    def __init__(self) -> None:
        self.calls = 0

    def __call__(self, prompt, **kwargs):
        self.calls += 1
        return [{"generated_text": f"ok:{prompt[-10:]}"}]


def test_cached_pipeline_reuses_loaded_model(monkeypatch) -> None:
    clear_local_model_cache()
    loads = []

    def fake_load(model_id: str):
        loads.append(model_id)
        return FakePipeline()

    monkeypatch.setattr("arena.local_model._load_pipeline", fake_load)

    first = _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800)
    second = _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800)

    assert first is second
    assert loads == ["local-7b"]
    assert local_model_cache_size() == 1
    clear_local_model_cache()


def test_cached_pipeline_refresh_clears_model(monkeypatch) -> None:
    clear_local_model_cache()
    monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline())

    _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800)
    clear_local_model_cache()

    assert local_model_cache_size() == 0


def test_cached_local_chat_model_generates_with_cached_pipeline(monkeypatch) -> None:
    clear_local_model_cache()
    monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline())

    model = CachedLocalChatModel(model_id="local-7b", max_new_tokens=16)
    result = model.invoke([HumanMessage(content="Build a tiny app")])

    assert result.content.startswith("ok:")
    assert local_model_cache_size() == 1
    clear_local_model_cache()


def test_local_gpu_duration_uses_cold_then_warm_cache(monkeypatch) -> None:
    clear_local_model_cache()
    monkeypatch.setenv("LOCAL_MODEL_GPU_COLD_DURATION_SECONDS", "700")
    monkeypatch.setenv("LOCAL_MODEL_GPU_WARM_DURATION_SECONDS", "90")
    monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline())

    cold = _local_generation_duration(
        model_id="local-7b",
        messages=[],
        temperature=0.2,
        max_new_tokens=16,
        stop=None,
        cache_ttl_seconds=1800,
    )
    _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800)
    warm = _local_generation_duration(
        model_id="local-7b",
        messages=[],
        temperature=0.2,
        max_new_tokens=16,
        stop=None,
        cache_ttl_seconds=1800,
    )

    assert cold == 700
    assert warm == 90
    clear_local_model_cache()


def test_local_gpu_duration_does_not_sync_bucket(monkeypatch) -> None:
    clear_local_model_cache()
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", "hf://buckets/user/model-cache/gemma")

    def fail_sync(*args, **kwargs):
        raise AssertionError("duration lookup should not sync bucket files")

    monkeypatch.setattr("huggingface_hub.sync_bucket", fail_sync)

    duration = _local_generation_duration(
        model_id="local-7b",
        messages=[],
        temperature=0.2,
        max_new_tokens=16,
        stop=None,
        cache_ttl_seconds=1800,
    )

    assert duration > 0


def test_model_load_sources_use_bucket_gguf(monkeypatch, tmp_path) -> None:
    bucket_uri = "hf://buckets/user/model-cache/gemma"
    cache_dir = tmp_path / "bucket-cache"
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", bucket_uri)
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_CACHE_DIR", str(cache_dir))
    monkeypatch.setenv("LOCAL_MODEL_GGUF_FILE", "model.gguf")

    def fake_sync(source, target):
        assert source == bucket_uri
        target_path = tmp_path / "bucket-cache" / _bucket_cache_key(bucket_uri)
        assert str(target_path) == target
        target_path.mkdir(parents=True, exist_ok=True)
        (target_path / "model.gguf").write_text("fake", encoding="utf-8")

    monkeypatch.setattr("huggingface_hub.sync_bucket", fake_sync)

    sources = _model_load_sources("unsloth/gemma-4-12B-it-qat-GGUF")

    assert sources.model_source == "unsloth/gemma-4-12B-it-qat-GGUF"
    assert sources.tokenizer_source == "unsloth/gemma-4-12B-it-qat-GGUF"
    assert sources.gguf_file.endswith("model.gguf")
    assert sources.bucket_uri == bucket_uri


def test_model_load_sources_use_bucket_transformers_dir(monkeypatch, tmp_path) -> None:
    bucket_uri = "hf://buckets/user/model-cache/safetensors"
    cache_dir = tmp_path / "bucket-cache"
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", bucket_uri)
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_CACHE_DIR", str(cache_dir))

    def fake_sync(source, target):
        assert source == bucket_uri
        target_path = tmp_path / "bucket-cache" / _bucket_cache_key(bucket_uri)
        assert str(target_path) == target
        target_path.mkdir(parents=True, exist_ok=True)
        (target_path / "config.json").write_text("{}", encoding="utf-8")
        (target_path / "tokenizer.json").write_text("{}", encoding="utf-8")

    monkeypatch.setattr("huggingface_hub.sync_bucket", fake_sync)

    sources = _model_load_sources("local-bucket-model")

    assert sources.model_source == str(cache_dir / _bucket_cache_key(bucket_uri))
    assert sources.tokenizer_source == str(cache_dir / _bucket_cache_key(bucket_uri))
    assert sources.gguf_file == ""


def test_model_load_sources_prefer_mounted_model_dir(monkeypatch, tmp_path) -> None:
    mounted_dir = tmp_path / "models"
    mounted_dir.mkdir()
    (mounted_dir / "model.gguf").write_text("fake", encoding="utf-8")
    monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", "hf://buckets/user/model-cache")
    monkeypatch.setenv("LOCAL_MODEL_MOUNT_DIR", str(mounted_dir))
    monkeypatch.setenv("LOCAL_MODEL_GGUF_FILE", "model.gguf")

    def fail_sync(*args, **kwargs):
        raise AssertionError("mounted model dir should not sync bucket files")

    monkeypatch.setattr("huggingface_hub.sync_bucket", fail_sync)

    sources = _model_load_sources("unsloth/gemma-4-12B-it-qat-GGUF")

    assert sources.model_source == "unsloth/gemma-4-12B-it-qat-GGUF"
    assert sources.tokenizer_source == "unsloth/gemma-4-12B-it-qat-GGUF"
    assert sources.gguf_file == str(mounted_dir / "model.gguf")


def test_messages_to_prompt_falls_back_without_chat_template() -> None:
    prompt = _messages_to_prompt(
        [
            SystemMessage(content="You are concise."),
            HumanMessage(content="Build it."),
        ],
        tokenizer=None,
    )

    assert "system: You are concise." in prompt
    assert "user: Build it." in prompt
    assert prompt.endswith("assistant:")