feat(embed): BGE-M3 dense+sparse via FlagEmbedding (ADR-0002)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/embed/test_embedder.py

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