homelab-codex-ws/packages/kb-retrieval/tests/test_embed.py

"""Unit tests for the Ollama embedding client -- no real HTTP, no real Ollama."""
from __future__ import annotations

import aiohttp
import pytest

from kb_retrieval.embed import (
    TRANSIENT_EMBED_ERRORS,
    EmbeddingDimensionError,
    _vector_literal,
    check_ollama_health,
    embed_batch,
    embed_batch_resilient,
    embed_chunk,
)


class TestVectorLiteral:
    def test_formats_as_bracketed_csv(self):
        assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]"


class _FakeEmbedResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self._status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def json(self):
        return self._payload

    def raise_for_status(self):
        if self._status >= 400:
            raise RuntimeError(f"HTTP {self._status}")


class _FakeOllamaSession:
    """Serves a fixed embedding vector for every /api/embeddings POST, or errors by prompt."""

    def __init__(self, dim=1024, fail_for=None):
        self._dim = dim
        self._fail_for = fail_for or set()
        self.requests: list[dict] = []

    def post(self, url, json):
        self.requests.append({"url": url, "json": json})
        if json["prompt"] in self._fail_for:
            return _FakeEmbedResponse({}, status=500)
        return _FakeEmbedResponse({"embedding": [0.01] * self._dim})


class TestEmbedChunk:
    async def test_returns_embedding_and_elapsed(self):
        session = _FakeOllamaSession(dim=1024)
        embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world")
        assert len(embedding) == 1024
        assert elapsed >= 0
        assert session.requests == [
            {"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}}
        ]

    async def test_missing_embedding_key_raises(self):
        class _EmptyResponse(_FakeEmbedResponse):
            pass

        class _Session:
            def post(self, url, json):
                return _EmptyResponse({})

        with pytest.raises(ValueError):
            await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text")


class _FakeBatchResponse(_FakeEmbedResponse):
    pass


class _FakeBatchOllamaSession:
    """Serves a fixed set of embeddings for every /api/embed POST."""

    def __init__(self, dim=1024, n_override=None, bad_index=None):
        self._dim = dim
        self._n_override = n_override
        self._bad_index = bad_index
        self.requests: list[dict] = []

    def post(self, url, json, timeout=None):
        self.requests.append({"url": url, "json": json, "timeout": timeout})
        texts = json["input"]
        n = self._n_override if self._n_override is not None else len(texts)
        embeddings = [[0.01] * self._dim for _ in range(n)]
        if self._bad_index is not None and self._bad_index < len(embeddings):
            embeddings[self._bad_index] = [0.01] * (self._dim - 1)
        return _FakeBatchResponse({"embeddings": embeddings})


class TestEmbedBatch:
    async def test_returns_embeddings_and_elapsed(self):
        session = _FakeBatchOllamaSession(dim=1024)
        texts = ["hello", "world", "third"]
        embeddings, elapsed = await embed_batch(session, "http://fake-ollama", "bge-m3", texts)
        assert len(embeddings) == 3
        assert all(len(e) == 1024 for e in embeddings)
        assert elapsed >= 0
        assert session.requests == [
            {"url": "http://fake-ollama/api/embed",
             "json": {"model": "bge-m3", "input": texts}, "timeout": None}
        ]

    async def test_timeout_s_bounds_the_single_request(self):
        session = _FakeBatchOllamaSession(dim=1024)
        await embed_batch(session, "http://fake-ollama", "bge-m3", ["a"], timeout_s=5.0)
        assert session.requests[0]["timeout"].total == 5.0

    async def test_length_mismatch_raises_dimension_error(self):
        session = _FakeBatchOllamaSession(dim=1024, n_override=2)
        with pytest.raises(EmbeddingDimensionError):
            await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"])

    async def test_wrong_dimension_on_one_item_raises(self):
        session = _FakeBatchOllamaSession(dim=1024, bad_index=1)
        with pytest.raises(EmbeddingDimensionError):
            await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"])

    async def test_missing_embeddings_key_raises_value_error(self):
        class _Session:
            def post(self, url, json):
                return _FakeBatchResponse({})

        with pytest.raises(ValueError):
            await embed_batch(_Session(), "http://fake-ollama", "bge-m3", ["a"])


class _FakeTagsResponse:
    def __init__(self, status):
        self.status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False


class _FakeHealthSession:
    def __init__(self, status=200, raises=None):
        self._status = status
        self._raises = raises

    def get(self, url, timeout=None):
        if self._raises is not None:
            raise self._raises
        return _FakeTagsResponse(self._status)


class TestCheckOllamaHealth:
    async def test_up_on_2xx(self):
        session = _FakeHealthSession(status=200)
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is True

    async def test_down_on_error_status(self):
        session = _FakeHealthSession(status=500)
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False

    async def test_down_on_connection_error(self):
        import aiohttp

        session = _FakeHealthSession(raises=aiohttp.ClientConnectionError("refused"))
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False

    async def test_down_on_timeout(self):
        session = _FakeHealthSession(raises=TimeoutError())
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False


class _FailResponse(_FakeEmbedResponse):
    """A transport failure, i.e. something in TRANSIENT_EMBED_ERRORS -- not the plain
    RuntimeError `_FakeEmbedResponse` raises, which would (correctly) never be retried."""

    def raise_for_status(self):
        raise aiohttp.ClientConnectionError("simulated transport failure")


class _ResilientSession:
    """Drives every branch of `embed_batch_resilient`.

    * `poison` -- texts that fail whatever batch they land in (the all-or-nothing property of
      /api/embed that bisection exists to work around)
    * `fail_first` -- the first N posts fail regardless (a transient blip)
    * `post_exc` -- raised instead of returning a response (used for the TimeoutError regression)
    * `health_sequence` -- consumed one entry per /api/tags probe, then falls back to `health`
    """

    def __init__(self, *, dim=1024, poison=(), fail_first=0, health=True,
                 health_sequence=None, post_exc=None):
        self._dim = dim
        self._poison = set(poison)
        self._fail_first = fail_first
        self._health = health
        self._health_sequence = list(health_sequence or [])
        self._post_exc = post_exc
        self.posts: list[list[str]] = []
        self.health_checks = 0

    def post(self, url, json, timeout=None):
        texts = json["input"]
        self.posts.append(list(texts))
        if self._post_exc is not None:
            raise self._post_exc
        if self._fail_first > 0:
            self._fail_first -= 1
            return _FailResponse({})
        if self._poison & set(texts):
            return _FailResponse({})
        return _FakeBatchResponse({"embeddings": [[0.01] * self._dim for _ in texts]})

    def get(self, url, timeout=None):
        self.health_checks += 1
        up = self._health_sequence.pop(0) if self._health_sequence else self._health
        return _FakeTagsResponse(200 if up else 500)


def _recorder():
    sleeps: list[float] = []

    async def fake_sleep(seconds):
        sleeps.append(seconds)

    return sleeps, fake_sleep


class TestEmbedBatchResilient:
    async def test_clean_run_costs_one_request_and_no_health_probe(self):
        session = _ResilientSession()
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", ["a", "b", "c"]
        )
        assert outcome.ok_count == 3
        assert outcome.failed_indices == []
        assert outcome.gave_up is False
        assert outcome.requests == 1
        assert outcome.retries == 0
        assert session.health_checks == 0

    async def test_transient_failure_is_retried_with_exponential_backoff(self):
        session = _ResilientSession(fail_first=2)
        sleeps, fake_sleep = _recorder()
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", ["a", "b"],
            retries=2, backoff_s=1.0, sleep=fake_sleep,
        )
        assert outcome.ok_count == 2
        assert outcome.requests == 3
        assert outcome.retries == 2
        assert sleeps == [1.0, 2.0]
        assert session.health_checks == 0  # never needed -- the retry recovered it

    async def test_timeout_is_transient_and_never_escapes(self):
        """Regression: aiohttp raises a bare builtins.TimeoutError when ClientTimeout(total=...)
        expires, and that is NOT an aiohttp.ClientError. Callers catching ClientError alone
        crashed instead of degrading -- exactly Ollama@SOLARIA's hang-not-refuse failure mode."""
        assert not isinstance(TimeoutError(), aiohttp.ClientError)  # the trap this guards
        assert isinstance(TimeoutError(), TRANSIENT_EMBED_ERRORS)

        session = _ResilientSession(post_exc=TimeoutError(), health=False)
        _sleeps, fake_sleep = _recorder()
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", ["a", "b"], retries=1, sleep=fake_sleep,
        )
        assert outcome.gave_up is True
        assert outcome.failed_indices == [0, 1]

    async def test_dead_backend_gives_up_without_bisecting(self):
        texts = [f"t{i}" for i in range(8)]
        session = _ResilientSession(poison=texts, health=False)
        _sleeps, fake_sleep = _recorder()

        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", texts, retries=1, sleep=fake_sleep,
        )

        assert outcome.gave_up is True
        assert outcome.ok_count == 0
        assert outcome.failed_indices == list(range(8))
        # 1 attempt + 1 retry, then the health probe stops it. Bisecting a dead backend would
        # cost 2n-1 requests and delay the caller's circuit breaker.
        assert len(session.posts) == 2
        assert session.health_checks == 1

    async def test_poison_item_is_isolated_and_the_rest_still_embed(self):
        texts = [f"t{i}" for i in range(8)]
        session = _ResilientSession(poison={"t3"}, health=True)
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", texts, retries=0,
        )

        assert outcome.gave_up is False
        assert outcome.failed_indices == [3]
        assert outcome.ok_count == 7
        assert outcome.embeddings[3] is None
        assert all(outcome.embeddings[i] is not None for i in range(8) if i != 3)
        assert len(session.posts) < 2 * len(texts)  # bisection, not one request per item

    async def test_bisected_spans_do_not_retry(self):
        """The top-level attempt already proved this isn't a blip; re-retrying every sub-span
        would multiply a poison chunk's cost by (retries+1) at every level of the bisection."""
        texts = ["a", "b", "c", "d"]
        session = _ResilientSession(poison={"c"}, health=True)
        sleeps, fake_sleep = _recorder()

        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", texts,
            retries=2, backoff_s=1.0, sleep=fake_sleep,
        )

        assert outcome.failed_indices == [2]
        assert sleeps == [1.0, 2.0]  # top level only -- no backoff inside the bisection
        assert outcome.retries == 2

    async def test_backend_dying_mid_bisection_fails_every_unresolved_item(self):
        """The abandoned right-hand span must still be reported failed. Accumulating
        failed_indices per span silently under-reported it -- it never reaches a leaf."""
        texts = ["a", "b", "c", "d"]
        session = _ResilientSession(poison=texts, health_sequence=[True, False])
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", texts, retries=0,
        )

        assert outcome.gave_up is True
        assert outcome.ok_count == 0
        assert outcome.failed_indices == [0, 1, 2, 3]
        assert all(e is None for e in outcome.embeddings)

    async def test_dimension_error_is_never_retried_or_swallowed(self):
        session = _FakeBatchOllamaSession(dim=1024, n_override=2)
        with pytest.raises(EmbeddingDimensionError):
            await embed_batch_resilient(
                session, "http://fake-ollama", "bge-m3", ["a", "b", "c"], retries=3,
            )
        assert len(session.requests) == 1

    async def test_empty_input_makes_no_requests(self):
        session = _ResilientSession()
        outcome = await embed_batch_resilient(session, "http://fake-ollama", "bge-m3", [])
        assert outcome.embeddings == []
        assert outcome.requests == 0
        assert session.posts == []

    async def test_elapsed_and_requests_are_accumulated_across_retries(self):
        session = _ResilientSession(fail_first=1)
        _sleeps, fake_sleep = _recorder()
        outcome = await embed_batch_resilient(
            session, "http://fake-ollama", "bge-m3", ["a"], retries=1, sleep=fake_sleep,
        )
        assert outcome.requests == 2
        assert outcome.elapsed_s >= 0