96 lines
2.6 KiB
Python
96 lines
2.6 KiB
Python
"""llm.classify — one-of-N choice from the local model, never free text."""
|
|
|
|
from __future__ import annotations
|
|
|
|
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
|