codex-py/tests/embed/test_embedder.py

"""Tests for codex.embed.

The real :class:`FlagEmbedding.BGEM3FlagModel` would download 2 GB of
weights and load them onto the device.  We mock it out so the suite
runs offline in milliseconds.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import pytest

import codex.embed as embed_module
from codex.embed import Embedder, get_embedder

# ---------------------------------------------------------------------------
# Fakes / fixtures
# ---------------------------------------------------------------------------


class _FakeBGEModel:
    """Stand-in for :class:`BGEM3FlagModel` that records calls.

    ``encode()`` returns deterministic fake vectors so the L2-norm
    assertions are exact.  The ``call_count`` attribute lets tests
    verify the single-forward-pass invariant of :meth:`Embedder.encode`.
    """

    DIM = 1024

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.init_args = args
        self.init_kwargs = kwargs
        self.call_count = 0
        self.last_kwargs: dict[str, Any] | None = None

    def encode(
        self,
        sentences: list[str],
        **kwargs: Any,
    ) -> dict[str, Any]:
        self.call_count += 1
        self.last_kwargs = kwargs
        n = len(sentences)
        rng = np.random.default_rng(seed=42)
        return {
            "dense_vecs": rng.random((n, self.DIM)).astype(np.float32),
            "lexical_weights": [{0: 0.5, 7: 0.25} for _ in range(n)],
        }


@pytest.fixture
def fake_model(monkeypatch: pytest.MonkeyPatch) -> type[_FakeBGEModel]:
    """Replace ``BGEM3FlagModel`` in ``codex.embed`` with the fake class."""
    monkeypatch.setattr(embed_module, "BGEM3FlagModel", _FakeBGEModel)
    return _FakeBGEModel


@pytest.fixture(autouse=True)
def reset_singleton() -> None:
    """Reset the module-level singleton between tests."""
    embed_module._embedder = None


# ---------------------------------------------------------------------------
# encode_dense
# ---------------------------------------------------------------------------


def test_encode_dense_shape_dtype_and_norm(fake_model: type[_FakeBGEModel]) -> None:
    """Dense output is (N, dim) float32 with unit L2 rows."""
    e = Embedder()
    out = e.encode_dense(["a", "b"])

    assert out.shape == (2, 1024)
    assert out.dtype == np.float32
    norms = np.linalg.norm(out, axis=1)
    np.testing.assert_allclose(norms, [1.0, 1.0], atol=1e-5)


def test_encode_dense_empty_input_returns_empty_matrix(
    fake_model: type[_FakeBGEModel],
) -> None:
    """Empty input -> (0, dim) without invoking the model."""
    e = Embedder()
    # Reach into the wrapped model to verify it is not called.
    fake = e._model
    assert isinstance(fake, _FakeBGEModel)

    out = e.encode_dense([])

    assert out.shape == (0, 1024)
    assert out.dtype == np.float32
    assert fake.call_count == 0


# ---------------------------------------------------------------------------
# encode_sparse
# ---------------------------------------------------------------------------


def test_encode_sparse_returns_list_of_dicts(fake_model: type[_FakeBGEModel]) -> None:
    """Sparse output has one dict per input string."""
    e = Embedder()
    out = e.encode_sparse(["a", "b"])

    assert len(out) == 2
    for row in out:
        assert isinstance(row, dict)
        for key, value in row.items():
            assert isinstance(key, int)
            assert isinstance(value, float)


def test_encode_sparse_empty_input_returns_empty_list(
    fake_model: type[_FakeBGEModel],
) -> None:
    """Empty input -> [] without invoking the model."""
    e = Embedder()
    fake = e._model
    assert isinstance(fake, _FakeBGEModel)

    out = e.encode_sparse([])

    assert out == []
    assert fake.call_count == 0


# ---------------------------------------------------------------------------
# encode (combined)
# ---------------------------------------------------------------------------


def test_encode_returns_tuple_and_uses_single_forward_pass(
    fake_model: type[_FakeBGEModel],
) -> None:
    """encode() must issue exactly one model.encode() call for efficiency."""
    e = Embedder()
    fake = e._model
    assert isinstance(fake, _FakeBGEModel)

    dense, sparse = e.encode(["a"])

    assert isinstance(dense, np.ndarray)
    assert dense.shape == (1, 1024)
    assert dense.dtype == np.float32
    np.testing.assert_allclose(np.linalg.norm(dense, axis=1), [1.0], atol=1e-5)

    assert isinstance(sparse, list)
    assert len(sparse) == 1
    assert isinstance(sparse[0], dict)

    assert fake.call_count == 1
    assert fake.last_kwargs is not None
    assert fake.last_kwargs.get("return_dense") is True
    assert fake.last_kwargs.get("return_sparse") is True


def test_encode_empty_input(fake_model: type[_FakeBGEModel]) -> None:
    """Empty input -> ((0, dim), [])."""
    e = Embedder()
    fake = e._model
    assert isinstance(fake, _FakeBGEModel)

    dense, sparse = e.encode([])

    assert dense.shape == (0, 1024)
    assert dense.dtype == np.float32
    assert sparse == []
    assert fake.call_count == 0


# ---------------------------------------------------------------------------
# Singleton
# ---------------------------------------------------------------------------


def test_get_embedder_returns_singleton(fake_model: type[_FakeBGEModel]) -> None:
    """Two calls return the same object."""
    first = get_embedder()
    second = get_embedder()

    assert first is second
    assert isinstance(first, Embedder)