feat(llm): closed-vocabulary classifier on the local model pool (refs #685)

src/llm/classify.py

@@ -0,0 +1,81 @@
+"""Closed-vocabulary classification on the local model pool.
+
+Given a text and a list of allowed slugs, the model must answer with
+exactly one slug or ``none``. Anything else is treated as ``None`` so a
+caller can queue the text for human review instead of inventing a value.
+Self-hosted inference only (Ollama ``/api/chat`` on the largest live
+host, the same route ``llm.rag.stream_answer`` takes).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Sequence
+
+import httpx
+
+from llm.pool import pick_model
+
+log = logging.getLogger(__name__)
+
+Classifier = Callable[[str, Sequence[str]], str | None]
+
+_SYSTEM = (
+    "You label short healthcare-policy text with exactly one slug from a "
+    "closed list. Reply with the slug only — no words, no punctuation. If no "
+    "slug fits, reply: none"
+)
+_TIMEOUT = httpx.Timeout(120.0, connect=10.0)
+
+
+def build_prompt(text: str, choices: Sequence[str]) -> list[dict]:
+    listing = "\n".join(f"- {c}" for c in choices)
+    return [
+        {"role": "system", "content": _SYSTEM},
+        {
+            "role": "user",
+            "content": f"Allowed slugs:\n{listing}\n- none\n\nText:\n{text.strip()}\n\nSlug:",
+        },
+    ]
+
+
+def parse_choice(reply: str, choices: Sequence[str]) -> str | None:
+    token = (reply or "").strip().strip("\"'`.").strip().lower()
+    if token == "none":
+        return None
+    return token if token in {c.lower() for c in choices} else None
+
+
+def _default_post(url: str, json: dict) -> dict:
+    with httpx.Client(timeout=_TIMEOUT) as client:
+        resp = client.post(url, json=json)
+        resp.raise_for_status()
+        return resp.json()
+
+
+def closed_vocab_classifier(
+    cfg: Any, pool: Any, *, post: Callable[..., dict] | None = None
+) -> Classifier:
+    """Return ``classify(text, choices) -> slug | None`` bound to *pool*."""
+    send = post or _default_post
+
+    def classify(text: str, choices: Sequence[str]) -> str | None:
+        pool.check(cfg.instruct_model)
+        with pool.acquire_generation() as host:
+            model = pick_model(cfg, pool, host)
+            data = send(
+                f"{host}/api/chat",
+                json={
+                    "model": model,
+                    "messages": build_prompt(text, choices),
+                    "stream": False,
+                    "options": {"num_ctx": cfg.chat_num_ctx, "temperature": 0},
+                },
+            )
+        reply = data.get("message", {}).get("content", "")
+        choice = parse_choice(reply, choices)
+        if choice is None and reply.strip().lower() != "none":
+            log.info("classify: unparseable reply %r for %r", reply[:60], text[:60])
+        return choice
+
+    return classify

tests/llm/test_classify.py

@@ -0,0 +1,95 @@
+"""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