stack/tests/llm/test_classify.py

"""llm.classify — one-of-N choice from the local model, never free text."""

from __future__ import annotations

from unittest.mock import MagicMock

from llm.classify import build_prompt, closed_vocab_classifier, parse_choice


class TestParseChoice:
    def test_exact(self):
        assert parse_choice("consent", ["consent", "24-7-access"]) == "consent"

    def test_quoted_and_cased(self):
        assert (
            parse_choice('  "24-7-Access".\n', ["consent", "24-7-access"])
            == "24-7-access"
        )

    def test_none_and_garbage(self):
        assert parse_choice("none", ["consent"]) is None
        assert (
            parse_choice(
                "I think it is about consent and access", ["consent", "24-7-access"]
            )
            is None
        )


class TestPrompt:
    def test_prompt_lists_choices_and_none(self):
        msgs = build_prompt(
            "Provide 24/7 access for urgent needs", ["consent", "24-7-access"]
        )
        assert msgs[0]["role"] == "system"
        user = msgs[1]["content"]
        assert "consent" in user and "24-7-access" in user and "none" in user
        assert "Provide 24/7 access" in user


class _Pool:
    def __init__(self):
        self.checked = []

    def check(self, model):
        self.checked.append(model)
        return ["http://h"]

    def vram(self, host):
        return 0.0

    def serves(self, host, model):
        return True

    class _Ctx:
        def __enter__(self):
            return "http://h"

        def __exit__(self, *a):
            return False

    def acquire_generation(self):
        return self._Ctx()


class _Cfg:
    instruct_model = "qwen2.5:14b"
    instruct_model_large = ""
    large_min_vram_gb = 20.0
    chat_num_ctx = 8192
    host_vram = {}


class TestClassifier:
    def test_returns_choice_from_model_reply(self):
        calls = []

        def post(url, json):
            calls.append((url, json))
            return {"message": {"content": "24-7-access"}}

        classify = closed_vocab_classifier(_Cfg(), _Pool(), post=post)
        assert (
            classify("Provide 24/7 access for urgent needs", ["consent", "24-7-access"])
            == "24-7-access"
        )
        url, body = calls[0]
        assert url == "http://h/api/chat" and body["stream"] is False
        assert body["options"]["temperature"] == 0

    def test_unparseable_reply_is_none(self):
        classify = closed_vocab_classifier(
            _Cfg(),
            _Pool(),
            post=lambda url, json: {"message": {"content": "maybe consent?"}},
        )
        assert classify("x", ["consent"]) is None

    def test_pool_check_runs_once_not_per_call(self):
        # I5: pool.check is a host-liveness probe, done once when the
        # classifier closure is built — not per classify() call.
        pool = _Pool()
        classify = closed_vocab_classifier(
            _Cfg(), pool, post=lambda url, json: {"message": {"content": "consent"}}
        )
        classify("a", ["consent"])
        classify("b", ["consent"])
        assert pool.checked == ["qwen2.5:14b"]


class TestDefaultPost:
    def test_default_post_used_when_none_given(self, monkeypatch):
        """No ``post`` kwarg — ``_default_post`` opens its own httpx.Client,
        posts, raises for status, and returns the parsed JSON body."""
        resp = MagicMock()
        resp.raise_for_status.return_value = None
        resp.json.return_value = {"message": {"content": "consent"}}
        client = MagicMock()
        client.post.return_value = resp
        client_cls = MagicMock()
        client_cls.return_value.__enter__.return_value = client
        client_cls.return_value.__exit__.return_value = False
        monkeypatch.setattr("llm.classify.httpx.Client", client_cls)

        classify = closed_vocab_classifier(_Cfg(), _Pool())
        assert classify("x", ["consent"]) == "consent"

        url, kwargs = client.post.call_args.args[0], client.post.call_args.kwargs
        assert url == "http://h/api/chat"
        assert kwargs["json"]["model"] == "qwen2.5:14b"
        resp.raise_for_status.assert_called_once()