fix: mini fixes

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)

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"

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"

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

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)