stack/tests/llm/test_structured.py

"""llm.structured — schema-constrained extraction on the local pool."""

from __future__ import annotations

import httpx
import pytest

from llm.structured import JsonExtractor, json_extractor, validate

SCHEMA = {
    "type": "object",
    "properties": {"ok": {"type": "boolean"}},
    "required": ["ok"],
}


class Cfg:
    instruct_model = "small"
    instruct_model_large = "big"
    large_min_vram_gb = 20.0
    chat_num_ctx = 4096


class Pool:
    def __init__(self, hosts):
        self.hosts = list(hosts)
        self.checked = []
        self.dropped = []

    def check(self, model):
        self.checked.append(model)
        return self.hosts

    def acquire_generation(self):
        host = self.hosts[0]

        class _cm:
            def __enter__(s):
                return host

            def __exit__(s, *a):
                return False

        return _cm()

    def drop(self, host):
        self.dropped.append(host)
        self.hosts.remove(host)


def _reply(text):
    return {"message": {"content": text}}


def test_extractor_posts_schema_and_fixed_model():
    sent = []

    def post(url, json=None):
        sent.append((url, json))
        return _reply('{"ok": true}')

    pool = Pool(["http://rig"])
    ask = json_extractor(Cfg(), pool, post=post)
    assert isinstance(ask, JsonExtractor) and ask.model == "small"
    assert ask("sys", "usr", SCHEMA) == {"ok": True}
    assert pool.checked == ["small"]
    url, body = sent[0]
    assert url == "http://rig/api/chat"
    assert body["model"] == "small" and body["format"] == SCHEMA
    assert body["options"] == {"num_ctx": 4096, "temperature": 0}
    assert body["stream"] is False
    assert [m["role"] for m in body["messages"]] == ["system", "user"]


def test_extractor_default_post_is_classifys():
    from llm.classify import _default_post

    assert json_extractor(Cfg(), Pool(["h"]))._send is _default_post


def test_extractor_drops_dead_host_and_retries():
    seen = []

    def post(url, json=None):
        seen.append(url)
        if "dead" in url:
            raise httpx.ConnectError("down")
        return _reply('{"ok": false}')

    pool = Pool(["http://dead", "http://live"])
    assert json_extractor(Cfg(), pool, post=post)("s", "u", SCHEMA) == {"ok": False}
    assert pool.dropped == ["http://dead"] and len(seen) == 2


def test_extractor_raises_when_every_host_fails():
    def post(url, json=None):
        raise httpx.ConnectError("down")

    with pytest.raises(RuntimeError, match="every Ollama host failed"):
        json_extractor(Cfg(), Pool(["http://a"]), post=post)("s", "u", SCHEMA)


@pytest.mark.parametrize(
    "reply", ['{"ok": "yes"}', "not json", "[1]", "{}", '{"message": 1}']
)
def test_extractor_returns_none_on_schema_violation(reply):
    ask = json_extractor(Cfg(), Pool(["h"]), post=lambda url, json=None: _reply(reply))
    assert ask("s", "u", SCHEMA) is None


def test_extractor_empty_message_is_none():
    ask = json_extractor(Cfg(), Pool(["h"]), post=lambda url, json=None: {})
    assert ask("s", "u", SCHEMA) is None


def test_extractor_non_object_with_permissive_schema_is_none():
    ask = json_extractor(Cfg(), Pool(["h"]), post=lambda url, json=None: _reply("[1]"))
    assert ask("s", "u", {}) is None


def test_validate_enum_required_maxlength():
    schema = {
        "type": "object",
        "required": ["kind", "note"],
        "properties": {
            "kind": {"type": "string", "enum": ["a", "b"]},
            "note": {"type": "string", "maxLength": 3},
        },
    }
    assert validate({"kind": "a", "note": "abc"}, schema) == []
    errs = validate({"kind": "c", "note": "abcd"}, schema)
    assert len(errs) == 2 and "enum" in errs[0] and "longer" in errs[1]
    assert validate({}, schema) == ["$.kind: required", "$.note: required"]


def test_validate_types_and_arrays():
    assert validate(1, {"type": "integer"}) == []
    assert validate(True, {"type": "integer"})
    assert validate(1.5, {"type": "number"}) == []
    assert validate(True, {"type": "number"})
    assert validate(None, {"type": ["string", "null"]}) == []
    assert validate(3, {"type": ["string", "null"]}) == [
        "$: expected string|null, got int"
    ]
    arr = {"type": "array", "maxItems": 2, "items": {"type": "string"}}
    assert validate(["a", "b"], arr) == []
    errs = validate(["a", 1, "c"], arr)
    assert len(errs) == 2 and "more than 2" in errs[0] and "[1]" in errs[1]
    assert validate([1, 2], {"items": {"type": "string"}})
    assert validate("x", {"type": "array"})


def test_extractor_non_string_content_is_none():
    ask = json_extractor(
        Cfg(), Pool(["h"]), post=lambda url, json=None: {"message": {"content": None}}
    )
    assert ask("s", "u", SCHEMA) is None