From 869d87452b9e49680bd110e9dffd19ed1047b7ca Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 27 Aug 2026 13:45:34 +0100 Subject: [PATCH] fix: minor embed compat fixes --- application/scripts/reembed.py | 7 ++- application/vectorstore/base.py | 13 +++++ application/vectorstore/embeddings_local.py | 22 ++++++++ tests/scripts/test_reembed.py | 8 +++ tests/vectorstore/test_embeddings_local.py | 51 +++++++++++++++++++ .../test_remote_embeddings_truncation.py | 22 ++++++++ 6 files changed, 121 insertions(+), 2 deletions(-) diff --git a/application/scripts/reembed.py b/application/scripts/reembed.py index d861049f..47eb916c 100644 --- a/application/scripts/reembed.py +++ b/application/scripts/reembed.py @@ -202,7 +202,7 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, ``(chunks_seen, chunks_written)``. """ store = VectorCreator.create_vectorstore( - "faiss", source_id=source_id, embeddings_key="embeddings" + "faiss", source_id=source_id, embeddings_key=settings.EMBEDDINGS_KEY ) chunks: List[Dict[str, Any]] = store.get_chunks() or [] if dry_run or not chunks: @@ -217,7 +217,10 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, # index and touches nothing on disk. The existing index is only replaced by # ``save_local`` below, so a failure while embedding leaves it intact. rebuilt = VectorCreator.create_vectorstore( - "faiss", source_id=source_id, embeddings_key="embeddings", docs_init=docs + "faiss", + source_id=source_id, + embeddings_key=settings.EMBEDDINGS_KEY, + docs_init=docs, ) rebuilt.save_local() return len(chunks), len(docs) diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index 2e30ad40..eceff89c 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -13,6 +13,11 @@ from application.vectorstore.model_registry import ( ) +def _embeddings_name_is_explicit() -> bool: + """True when ``EMBEDDINGS_NAME`` was configured rather than defaulted.""" + return "EMBEDDINGS_NAME" in getattr(settings, "model_fields_set", set()) + + class RemoteEmbeddings: """ Wrapper for remote embeddings API (OpenAI-compatible). @@ -47,10 +52,18 @@ class RemoteEmbeddings: model contributes its own context window, so a request that the server would reject -- or silently truncate -- is clipped here instead of being sent and paid for. + + That fallback needs the name to mean something. For a remote server it + is only a label forwarded as the ``model`` field, so the settings + default must not lend the server mpnet's 384-token window: a name + nobody chose describes nothing, and clipping on it would silently + discard most of every chunk. """ configured = settings.EMBEDDINGS_MAX_INPUT_TOKENS if configured and configured > 0: return configured + if not _embeddings_name_is_explicit(): + return None model_limit = max_input_tokens_for(self.model_name) return model_limit if model_limit and model_limit > 0 else None diff --git a/application/vectorstore/embeddings_local.py b/application/vectorstore/embeddings_local.py index 675ceb97..b9a55865 100644 --- a/application/vectorstore/embeddings_local.py +++ b/application/vectorstore/embeddings_local.py @@ -94,6 +94,27 @@ def _spec_for(model_name: str) -> EmbeddingModel: ) +def _pad_to_longest_in_batch(model: Any) -> None: + """Undo a fixed padding width baked into a model's ``tokenizer.json``.""" + # FastEmbed enables padding only when the tokenizer declares none, so a + # fixed ``length`` survives loading. Shorter inputs are then padded to that + # width while longer ones keep their own, the batch is ragged, and the ONNX + # tensor build fails. mpnet ships ``length: 128``; granite does not. + # Mean pooling masks pad tokens, so the vectors are unaffected. + tokenizer = getattr(getattr(model, "model", None), "tokenizer", None) + padding = getattr(tokenizer, "padding", None) + if not isinstance(padding, dict) or padding.get("length") is None: + return + tokenizer.enable_padding( + direction=padding.get("direction", "right"), + pad_id=padding.get("pad_id", 0), + pad_type_id=padding.get("pad_type_id", 0), + pad_token=padding.get("pad_token", ""), + length=None, + pad_to_multiple_of=padding.get("pad_to_multiple_of"), + ) + + class EmbeddingsWrapper: """Runs an embedding model locally through FastEmbed. @@ -131,6 +152,7 @@ class EmbeddingsWrapper: f"{exc}. Known models: {', '.join(known_names())}." ) from exc + _pad_to_longest_in_batch(self.model) self.dimension = self.spec.dimension or self._probe_dimension() logger.info("Embeddings model ready (dimension=%d)", self.dimension) diff --git a/tests/scripts/test_reembed.py b/tests/scripts/test_reembed.py index e384662b..a77fc4ca 100644 --- a/tests/scripts/test_reembed.py +++ b/tests/scripts/test_reembed.py @@ -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 diff --git a/tests/vectorstore/test_embeddings_local.py b/tests/vectorstore/test_embeddings_local.py index 257ede11..68e84dca 100644 --- a/tests/vectorstore/test_embeddings_local.py +++ b/tests/vectorstore/test_embeddings_local.py @@ -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_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"] == "" + + def test_dynamic_padding_is_left_alone(self, fake_fastembed): + _, instance = fake_fastembed + tokenizer = self._tokenizer({"length": None, "pad_id": 0, "pad_token": ""}) + 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) diff --git a/tests/vectorstore/test_remote_embeddings_truncation.py b/tests/vectorstore/test_remote_embeddings_truncation.py index dbc88c7b..8b7e2d64 100644 --- a/tests/vectorstore/test_remote_embeddings_truncation.py +++ b/tests/vectorstore/test_remote_embeddings_truncation.py @@ -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):