fix: minor embed compat fixes

This commit is contained in:
Alex committed 2026-08-27 13:45:34 +01:00
1 parent dd3876fdcb
commit 869d87452b
6 files changed
+121 -2

No files matched your search

+8
View File
@@ -99,6 +99,14 @@ class TestFaissRebuild:
assert [d.page_content for d in docs] == ["alpha", "beta"]
assert [d.metadata for d in docs] == [{"i": 0}, {"i": 1}]
def test_embeddings_key_comes_from_settings(self, stores):
"""A placeholder here is sent as the server's bearer token."""
_, _, factory = stores
with patch.object(reembed.settings, "EMBEDDINGS_KEY", "sk-real", create=True):
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
keys = [call.kwargs.get("embeddings_key") for call in factory.call_args_list]
assert keys == ["sk-real", "sk-real"]
def test_existing_index_is_not_deleted(self, stores):
"""The rebuild must not destroy the old index before the new one exists."""
existing, rebuilt, _ = stores
@@ -165,3 +165,54 @@ class TestLengthSortedBatching:
wrapper, _ = self._wrapper(fake_fastembed, 2)
out = wrapper.embed_documents(["aa", "b", "aa", "ccc"])
assert out == [[2.0], [1.0], [2.0], [3.0]]
class TestTokenizerPadding:
"""A fixed padding width in ``tokenizer.json`` makes mixed batches ragged.
FastEmbed calls ``enable_padding`` only when the tokenizer declares none,
so mpnet's fixed ``length: 128`` survives loading. Any batch mixing an
input longer than 128 tokens with a shorter one then produces rows of
different widths and ONNX rejects the tensor.
"""
def _tokenizer(self, padding):
tokenizer = MagicMock()
tokenizer.padding = padding
return tokenizer
def test_fixed_width_padding_is_reset_to_batch_longest(self, fake_fastembed):
_, instance = fake_fastembed
tokenizer = self._tokenizer(
{
"length": 128,
"pad_id": 1,
"pad_token": "<pad>",
"pad_type_id": 0,
"direction": "right",
"pad_to_multiple_of": None,
}
)
instance.model.tokenizer = tokenizer
EmbeddingsWrapper(MPNET.name)
kwargs = tokenizer.enable_padding.call_args.kwargs
assert kwargs["length"] is None, "padding must follow the longest input"
# The model's own pad token must survive the reset.
assert kwargs["pad_id"] == 1
assert kwargs["pad_token"] == "<pad>"
def test_dynamic_padding_is_left_alone(self, fake_fastembed):
_, instance = fake_fastembed
tokenizer = self._tokenizer({"length": None, "pad_id": 0, "pad_token": "<pad>"})
instance.model.tokenizer = tokenizer
EmbeddingsWrapper(GRANITE_97M.name)
tokenizer.enable_padding.assert_not_called()
def test_tokenizer_that_cannot_be_reached_is_not_fatal(self, fake_fastembed):
_, instance = fake_fastembed
instance.model = None
EmbeddingsWrapper(GRANITE_97M.name)
@@ -105,10 +105,31 @@ class TestInputLimitResolution:
from application.vectorstore import base
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: True)
assert self._remote("granite-311m")._resolve_input_limit() == 32768
# mpnet genuinely stops at 384; sending more is paid for and discarded.
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() == 384
def test_default_model_name_lends_the_server_no_ceiling(self, monkeypatch):
"""An unset EMBEDDINGS_NAME describes nothing about the remote server.
The name is only forwarded as the ``model`` field. Letting the
settings default contribute mpnet's 384-token window would clip every
chunk on a server that may well serve a 32k-context model.
"""
from application.vectorstore import base
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", None)
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: False)
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() is None
def test_explicit_setting_still_wins_over_an_unset_name(self, monkeypatch):
from application.vectorstore import base
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 512)
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: False)
assert self._remote("all-mpnet-base-v2")._resolve_input_limit() == 512
def test_unknown_model_stays_unlimited(self, monkeypatch):
from application.vectorstore import base
@@ -119,6 +140,7 @@ class TestInputLimitResolution:
from application.vectorstore import base
monkeypatch.setattr(base.settings, "EMBEDDINGS_MAX_INPUT_TOKENS", 0)
monkeypatch.setattr(base, "_embeddings_name_is_explicit", lambda: True)
assert self._remote("granite-97m")._resolve_input_limit() == 32768
def test_dimension_is_taken_from_the_registry(self):