stack/tests/llm/test_pool.py

"""llm.pool — least-loaded fan-out across Ollama hosts."""

from unittest.mock import MagicMock, patch

import pytest

from llm.config import LlmConfig
from llm.pool import HostPool, PoolEmbeddings, embed_texts

H1, H2 = "http://h1:11434", "http://h2:11434"


def _resp(payload, status=200):
    r = MagicMock()
    r.status_code = status
    r.json.return_value = payload
    return r


class TestCheck:
    @patch("llm.pool.httpx.Client")
    def test_drops_host_missing_model(self, MockClient):
        client = MockClient.return_value.__enter__.return_value
        client.get.side_effect = [
            _resp({"models": [{"name": "nomic-embed-text:latest"}]}),
            _resp({"models": [{"name": "other:latest"}]}),
        ]
        pool = HostPool([H1, H2])
        assert pool.check("nomic-embed-text") == [H1]

    @patch("llm.pool.httpx.Client")
    def test_drops_unreachable_host(self, MockClient):
        import httpx

        client = MockClient.return_value.__enter__.return_value
        client.get.side_effect = [
            _resp({"models": [{"name": "m:latest"}]}),
            httpx.ConnectError("down"),
        ]
        pool = HostPool([H1, H2])
        assert pool.check("m") == [H1]

    @patch("llm.pool.httpx.Client")
    def test_no_hosts_left_raises(self, MockClient):
        client = MockClient.return_value.__enter__.return_value
        client.get.return_value = _resp({"models": []})
        with pytest.raises(RuntimeError, match="no Ollama host"):
            HostPool([H1]).check("m")

    @patch("llm.pool.httpx.Client")
    def test_drops_host_with_malformed_json(self, MockClient):
        """A host that's up but returns a non-JSON/malformed body is
        dropped, not treated as a fatal error for the whole run."""
        client = MockClient.return_value.__enter__.return_value
        bad = MagicMock()
        bad.json.side_effect = ValueError("not JSON")
        client.get.side_effect = [bad, _resp({"models": [{"name": "m:latest"}]})]
        pool = HostPool([H1, H2])
        assert pool.check("m") == [H2]

    @patch("llm.pool.httpx.Client")
    def test_drops_host_missing_name_key(self, MockClient):
        """A model dict missing 'name' raises KeyError inside check();
        that host should be dropped, not abort the whole check()."""
        client = MockClient.return_value.__enter__.return_value
        client.get.side_effect = [
            _resp({"models": [{"no_name": "whatever"}]}),
            _resp({"models": [{"name": "m:latest"}]}),
        ]
        pool = HostPool([H1, H2])
        assert pool.check("m") == [H2]


class TestFromConfig:
    def test_from_config_uses_ollama_hosts(self):
        cfg = LlmConfig(
            ollama_hosts=(H1, H2),
            embed_model="m",
            instruct_model="g",
            embed_dim=768,
            build_ann_index=False,
            pg_host="x",
            pg_port=5432,
            pg_db="llm",
            pg_user="llm",
        )
        pool = HostPool.from_config(cfg)
        assert pool.hosts == list(cfg.ollama_hosts)


class TestAcquire:
    def test_prefers_least_in_flight(self):
        pool = HostPool([H1, H2])
        with pool.acquire() as first, pool.acquire() as second:
            assert {first, second} == {H1, H2}

    def test_released_host_is_reused(self):
        pool = HostPool([H1, H2])
        with pool.acquire() as first:
            pass
        with pool.acquire() as a, pool.acquire() as b:
            assert {a, b} == {H1, H2}
        assert first in {H1, H2}


class TestEmbedTexts:
    @patch("llm.pool.httpx.Client")
    def test_order_preserved_across_batches(self, MockClient):
        client = MockClient.return_value.__enter__.return_value

        def fake_post(url, json=None, **kw):
            return _resp({"embeddings": [[float(len(t))] for t in json["input"]]})

        client.post.side_effect = fake_post
        pool = HostPool([H1, H2])
        out = embed_texts(pool, "m", ["a", "bb", "ccc"], batch_size=1)
        assert out == [[1.0], [2.0], [3.0]]

    @patch("llm.pool.httpx.Client")
    def test_http_error_raises(self, MockClient):
        client = MockClient.return_value.__enter__.return_value
        resp = _resp({}, status=500)
        resp.raise_for_status.side_effect = Exception("boom")
        client.post.return_value = resp
        with pytest.raises(Exception, match="boom"):
            embed_texts(HostPool([H1]), "m", ["a"])


class TestCheckAcquireInterleave:
    @patch("llm.pool.httpx.Client")
    def test_check_drops_held_host_without_corruption(self, MockClient):
        client = MockClient.return_value.__enter__.return_value
        pool = HostPool([H1, H2])
        with pool.acquire() as held:
            survivor = H2 if held == H1 else H1

            def fake_get(url, **kw):
                if url.startswith(held):
                    return _resp({"models": []})
                return _resp({"models": [{"name": "m:latest"}]})

            client.get.side_effect = fake_get
            assert pool.check("m") == [survivor]
        # exiting acquire() above must not raise despite the dropped host
        with pool.acquire() as again:
            assert again == survivor
        assert pool._in_flight[survivor] == 0


class TestPoolEmbeddings:
    @patch("llm.pool.embed_texts")
    def test_adapter_delegates(self, mock_embed):
        mock_embed.return_value = [[0.1]]
        pool = HostPool([H1])
        emb = PoolEmbeddings(pool, "m")
        assert emb.embed_documents(["x"]) == [[0.1]]
        assert emb.embed_query("x") == [0.1]