From dd3876fdcb757ec74de1ad52e8e090ca6c622869 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 26 Aug 2026 16:37:03 +0100 Subject: [PATCH] fix: mini fixes --- application/scripts/reembed.py | 23 ++++- tests/parser/test_tokenization.py | 101 +++++++++++++++++- tests/scripts/test_prefetch_models.py | 96 +++++++++++++++++ tests/scripts/test_reembed.py | 109 ++++++++++++++++++++ tests/scripts/test_reembed_pgvector_live.py | 4 + 5 files changed, 327 insertions(+), 6 deletions(-) create mode 100644 tests/scripts/test_prefetch_models.py diff --git a/application/scripts/reembed.py b/application/scripts/reembed.py index d8eb16b2..d861049f 100644 --- a/application/scripts/reembed.py +++ b/application/scripts/reembed.py @@ -26,6 +26,8 @@ import time from dataclasses import dataclass from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple +from psycopg import sql + from application.core.settings import settings from application.vectorstore.model_registry import resolve from application.vectorstore.vector_creator import VectorCreator @@ -84,9 +86,15 @@ def _pgvector_source_ids() -> List[str]: conn = store._get_connection() cursor = conn.cursor() try: + # Table and column names cannot be bound as parameters, so they go + # through ``sql.Identifier``, which quotes them. They are internal + # constants rather than user input, but composing SQL by f-string is + # the habit worth not having. cursor.execute( - f"SELECT DISTINCT source_id FROM {store._table_name} " - f"WHERE source_id IS NOT NULL ORDER BY source_id" + sql.SQL( + "SELECT DISTINCT source_id FROM {table} " + "WHERE source_id IS NOT NULL ORDER BY source_id" + ).format(table=sql.Identifier(store._table_name)) ) return [row[0] for row in cursor.fetchall()] finally: @@ -135,7 +143,9 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in cursor = conn.cursor() try: cursor.execute( - f"SELECT id, text FROM {table} WHERE source_id = %s ORDER BY id", + sql.SQL("SELECT id, text FROM {table} WHERE source_id = %s ORDER BY id").format( + table=sql.Identifier(table) + ), (source_id,), ) rows = cursor.fetchall() @@ -153,7 +163,12 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in cursor = conn.cursor() try: cursor.executemany( - f"UPDATE {table} SET {vector_column} = %s::vector WHERE id = %s", + sql.SQL( + "UPDATE {table} SET {column} = %s::vector WHERE id = %s" + ).format( + table=sql.Identifier(table), + column=sql.Identifier(vector_column), + ), [ (str(list(vector)), row[0]) for vector, row in zip(vectors, batch) diff --git a/tests/parser/test_tokenization.py b/tests/parser/test_tokenization.py index 6e9cc16a..8f391b81 100644 --- a/tests/parser/test_tokenization.py +++ b/tests/parser/test_tokenization.py @@ -42,7 +42,7 @@ class TestSplittingPreservesText: @pytest.mark.parametrize("text", SAMPLES) def test_tiktoken_split_reassembles_exactly(self, text, monkeypatch): - monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding()) + monkeypatch.setattr(tokenization, "get_encoding", _StubEncoding) counter = TiktokenCounter() pieces = counter.split(text, 7) assert "".join(pieces) == text @@ -80,7 +80,7 @@ class TestSplittingPreservesText: class TestCounting: def test_counts_differ_between_tokenizers(self, hf_counter, monkeypatch): """The whole point: mpnet and cl100k disagree, so units matter.""" - monkeypatch.setattr(tokenization, "get_encoding", lambda: _StubEncoding()) + monkeypatch.setattr(tokenization, "get_encoding", _StubEncoding) text = "internationalisation tokenization" assert hf_counter.count(text) != TiktokenCounter().count(text) @@ -105,3 +105,100 @@ class TestSelection: def test_counter_is_cached_per_model(self): first = get_token_counter("granite-311m") assert get_token_counter("granite-311m") is first + + +class TestTiktokenCounterEdges: + """The cl100k path is the fallback, so its edges matter as much.""" + + @pytest.fixture + def counter(self, monkeypatch): + monkeypatch.setattr(tokenization, "get_encoding", _StubEncoding) + return TiktokenCounter() + + def test_empty_text_counts_zero_and_splits_to_nothing(self, counter): + assert counter.count("") == 0 + assert counter.split("", 10) == [] + + def test_text_within_budget_is_returned_whole(self, counter): + assert counter.split("abc", 10) == ["abc"] + + def test_first_window_can_be_smaller_than_the_rest(self, counter): + """A header eats into the first chunk's budget only.""" + pieces = counter.split("abcdefghij", 4, first_max_tokens=2) + assert pieces[0] == "ab" + assert "".join(pieces) == "abcdefghij" + + +class TestCounterContract: + def test_base_class_requires_an_implementation(self): + base = tokenization.TokenCounter() + with pytest.raises(NotImplementedError): + base.count("x") + with pytest.raises(NotImplementedError): + base.split("x", 1) + + +class TestFallbackWhenTokenizerUnavailable: + def test_load_failure_returns_none_rather_than_raising(self, monkeypatch, caplog): + """Chunking must survive an offline host or a bad repo name.""" + import builtins + + real_import = builtins.__import__ + + def boom(name, *args, **kwargs): + if name == "tokenizers": + raise ImportError("no tokenizers here") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", boom) + assert tokenization._load_hf_counter("some/repo") is None + + def test_selection_falls_back_to_cl100k_on_failure(self, monkeypatch): + monkeypatch.setattr(tokenization, "_load_hf_counter", lambda repo: None) + assert isinstance(get_token_counter("granite-97m"), TiktokenCounter) + + def test_reset_cache_forces_reselection(self, monkeypatch): + first = get_token_counter("granite-311m") + tokenization.reset_cache() + monkeypatch.setattr(tokenization, "_load_hf_counter", lambda repo: None) + assert get_token_counter("granite-311m") is not first + + +class TestOffsetsWithoutSpans: + """Some tokenizers emit ``(0, 0)`` for specials or normalised-away chars. + + Those tokens consume budget but point at no text, so the splitter has to + skip the window rather than emit an empty piece or lose the tail. + """ + + class _Encoded: + def __init__(self, offsets): + self.offsets = offsets + self.ids = list(range(len(offsets))) + + class _Tokenizer: + def __init__(self, offsets): + self._offsets = offsets + + def encode(self, text, add_special_tokens=False): + return TestOffsetsWithoutSpans._Encoded(self._offsets) + + def _counter(self, offsets): + return HuggingFaceCounter(self._Tokenizer(offsets), "stub") + + def test_span_less_windows_are_skipped_not_emitted_empty(self): + # Two real tokens, then a window of pure (0, 0) padding-like entries. + counter = self._counter([(0, 2), (2, 4), (0, 0), (0, 0)]) + pieces = counter.split("abcd", 2) + assert "" not in pieces + assert "".join(pieces) == "abcd" + + def test_trailing_text_is_never_dropped(self): + """Offsets that stop short of the string must not lose the remainder.""" + counter = self._counter([(0, 1), (1, 2), (2, 3)]) + pieces = counter.split("abcdef", 2) + assert "".join(pieces) == "abcdef" + + def test_all_span_less_offsets_still_return_the_text(self): + counter = self._counter([(0, 0), (0, 0), (0, 0)]) + assert "".join(counter.split("abc", 1)) == "abc" diff --git a/tests/scripts/test_prefetch_models.py b/tests/scripts/test_prefetch_models.py new file mode 100644 index 00000000..ea8f5708 --- /dev/null +++ b/tests/scripts/test_prefetch_models.py @@ -0,0 +1,96 @@ +"""Model pre-fetching, which the Docker build runs to bake artifacts in.""" + +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest + +from application.scripts import prefetch_models +from application.vectorstore.model_registry import GRANITE_311M, MPNET, OPENAI_ADA_002 + + +@pytest.fixture +def fake_fastembed(): + """Stand in for FastEmbed so nothing is downloaded.""" + text_embedding = MagicMock() + pooling = types.SimpleNamespace(CLS="CLS", MEAN="MEAN") + module = types.ModuleType("fastembed") + module.TextEmbedding = text_embedding + desc = types.ModuleType("fastembed.common.model_description") + desc.PoolingType = pooling + desc.ModelSource = lambda hf=None: {"hf": hf} + with patch.dict( + sys.modules, + {"fastembed": module, "fastembed.common.model_description": desc}, + ): + yield text_embedding + + +class TestPrefetch: + def test_defaults_cover_legacy_and_new_install(self): + """An upgraded image must still serve mpnet; a new one needs granite.""" + assert MPNET.name in prefetch_models.DEFAULT_MODELS + assert GRANITE_311M.name in prefetch_models.DEFAULT_MODELS + + def test_fetches_named_models(self, fake_fastembed): + fetched = prefetch_models.prefetch([GRANITE_311M.name]) + assert fetched == [GRANITE_311M.repo] + assert fake_fastembed.call_args.kwargs["model_name"] == GRANITE_311M.repo + + def test_registers_with_registry_pooling_and_dim(self, fake_fastembed): + prefetch_models.prefetch([GRANITE_311M.name]) + kwargs = fake_fastembed.add_custom_model.call_args.kwargs + assert kwargs["model"] == GRANITE_311M.repo + assert kwargs["pooling"] == "CLS" + assert kwargs["dim"] == GRANITE_311M.dimension + assert kwargs["model_file"] == GRANITE_311M.onnx_file + + def test_mean_pooled_model_registers_as_mean(self, fake_fastembed): + prefetch_models.prefetch([MPNET.name]) + assert fake_fastembed.add_custom_model.call_args.kwargs["pooling"] == "MEAN" + + def test_aliases_resolve(self, fake_fastembed): + assert prefetch_models.prefetch(["granite-311m"]) == [GRANITE_311M.repo] + + def test_cache_dir_is_forwarded(self, fake_fastembed): + prefetch_models.prefetch([MPNET.name], cache_dir="/models") + assert fake_fastembed.call_args.kwargs["cache_dir"] == "/models" + + def test_cache_dir_omitted_when_absent(self, fake_fastembed): + prefetch_models.prefetch([MPNET.name]) + assert "cache_dir" not in fake_fastembed.call_args.kwargs + + def test_remote_only_model_is_skipped(self, fake_fastembed): + """OpenAI embeddings have no local artifacts to cache.""" + assert prefetch_models.prefetch([OPENAI_ADA_002.name]) == [] + fake_fastembed.assert_not_called() + + def test_unknown_model_fails_loudly(self, fake_fastembed): + """A silent skip at build time is a download at run time, offline.""" + with pytest.raises(SystemExit) as excinfo: + prefetch_models.prefetch(["nope/nope"]) + assert "nope/nope" in str(excinfo.value) + assert MPNET.name in str(excinfo.value) + + def test_several_models_in_one_run(self, fake_fastembed): + fetched = prefetch_models.prefetch([MPNET.name, GRANITE_311M.name]) + assert fetched == [MPNET.repo, GRANITE_311M.repo] + + +class TestMain: + def test_no_args_fetches_the_defaults(self, fake_fastembed): + with patch.object(prefetch_models, "prefetch", return_value=[]) as spy: + assert prefetch_models.main([]) == 0 + assert spy.call_args.args[0] == list(prefetch_models.DEFAULT_MODELS) + + def test_explicit_args_override_the_defaults(self, fake_fastembed): + with patch.object(prefetch_models, "prefetch", return_value=[]) as spy: + prefetch_models.main(["granite-97m"]) + assert spy.call_args.args[0] == ["granite-97m"] + + def test_cache_dir_read_from_environment(self, fake_fastembed, monkeypatch): + monkeypatch.setenv("EMBEDDINGS_CACHE_DIR", "/app/models") + with patch.object(prefetch_models, "prefetch", return_value=[]) as spy: + prefetch_models.main(["granite-97m"]) + assert spy.call_args.args[1] == "/app/models" diff --git a/tests/scripts/test_reembed.py b/tests/scripts/test_reembed.py index 1717bbcd..e384662b 100644 --- a/tests/scripts/test_reembed.py +++ b/tests/scripts/test_reembed.py @@ -122,3 +122,112 @@ class TestFaissRebuild: reembed.VectorCreator, "create_vectorstore", return_value=existing ): assert reembed.reembed_faiss("s1", batch_size=8, dry_run=False) == (0, 0) + + +class TestPgvectorWithoutTheExtension: + """Mocked pgvector paths. + + The live tests in ``test_reembed_pgvector_live`` skip wherever the cluster + has no pgvector build -- which includes CI -- so the SQL shape and the + batching contract are pinned here too. + """ + + @pytest.fixture + def store(self): + store = MagicMock() + store._table_name = "documents" + store._vector_column = "embedding" + cursor = MagicMock() + cursor.fetchall.return_value = [(1, "alpha"), (2, "beta"), (3, "gamma")] + conn = MagicMock() + conn.cursor.return_value = cursor + store._get_connection.return_value = conn + store._embedding.embed_documents.side_effect = lambda texts: [ + [0.5] * 4 for _ in texts + ] + with patch.object( + reembed.VectorCreator, "create_vectorstore", return_value=store + ): + yield store, conn, cursor + + def test_reads_and_rewrites_every_chunk(self, store): + _, conn, cursor = store + seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) + assert (seen, written) == (3, 3) + cursor.executemany.assert_called_once() + conn.commit.assert_called() + + def test_dry_run_neither_embeds_nor_writes(self, store): + fake_store, conn, cursor = store + seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=True) + assert (seen, written) == (3, 0) + fake_store._embedding.embed_documents.assert_not_called() + cursor.executemany.assert_not_called() + + def test_batches_commit_separately(self, store): + _, conn, cursor = store + reembed.reembed_pgvector("s1", batch_size=2, dry_run=False) + # 3 rows at batch 2 is two write transactions. + assert cursor.executemany.call_count == 2 + assert conn.commit.call_count == 2 + + def test_failed_batch_rolls_back_and_raises(self, store): + fake_store, conn, cursor = store + cursor.executemany.side_effect = RuntimeError("write failed") + with pytest.raises(RuntimeError): + reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) + conn.rollback.assert_called_once() + + def test_connection_is_returned_even_on_failure(self, store): + fake_store, _, cursor = store + cursor.executemany.side_effect = RuntimeError("boom") + with pytest.raises(RuntimeError): + reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) + fake_store.close.assert_called_once() + + def test_null_text_does_not_crash_the_embed_call(self, store): + fake_store, _, cursor = store + cursor.fetchall.return_value = [(1, None), (2, "beta")] + seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) + assert (seen, written) == (2, 2) + assert fake_store._embedding.embed_documents.call_args.args[0] == ["", "beta"] + + def test_empty_source_is_a_no_op(self, store): + fake_store, _, cursor = store + cursor.fetchall.return_value = [] + assert reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) == (0, 0) + fake_store._embedding.embed_documents.assert_not_called() + + def test_source_discovery_returns_sorted_ids(self, store): + _, _, cursor = store + cursor.fetchall.return_value = [("b",), ("a",)] + assert reembed.list_source_ids("pgvector") == ["b", "a"] + + +class TestFaissSourceDiscovery: + def test_source_ids_come_from_index_directories(self): + storage = MagicMock() + storage.list_files.return_value = [ + "indexes/src-a/index.faiss", + "indexes/src-a/index.pkl", + "indexes/src-b/index.faiss", + ] + with patch( + "application.storage.storage_creator.StorageCreator.get_storage", + return_value=storage, + ): + assert reembed.list_source_ids("faiss") == ["src-a", "src-b"] + + def test_storage_failure_is_reported_as_a_usable_error(self): + storage = MagicMock() + storage.list_files.side_effect = OSError("permission denied") + with patch( + "application.storage.storage_creator.StorageCreator.get_storage", + return_value=storage, + ): + with pytest.raises(reembed.ReembedError, match="permission denied"): + reembed.list_source_ids("faiss") + + def test_unsupported_store_error_names_the_alternatives(self): + with patch.object(reembed.settings, "VECTOR_STORE", "milvus", create=True): + assert reembed.main([]) == 2 diff --git a/tests/scripts/test_reembed_pgvector_live.py b/tests/scripts/test_reembed_pgvector_live.py index c7a0ac6c..165e311d 100644 --- a/tests/scripts/test_reembed_pgvector_live.py +++ b/tests/scripts/test_reembed_pgvector_live.py @@ -44,11 +44,15 @@ def _dsn(info) -> str: @pytest.fixture(autouse=True) def _close_pools(): + """Never leak a pool into another test; the DSN dies with the test DB.""" yield for dsn, pool in list(pgvector_module._POOLS.items()): try: pool.close() except Exception: + # Teardown only: the ephemeral cluster may already be gone, and a + # failure to close a pool for a dead DSN must not fail the test + # that just passed. Dropping the entry below is what matters. pass pgvector_module._POOLS.pop(dsn, None)