Ensures an embed migration stops the migration chain, just in case

src/paperless_ai/tests/test_ai_indexing.py

@@ -6,6 +6,7 @@ import pytest
 import pytest_mock
 from django.test import override_settings
 from django.utils import timezone
+from llama_index.core.schema import MetadataMode
 
 from documents.models import Document
 from documents.models import PaperlessTask
@@ -16,6 +17,7 @@ from documents.tests.factories import PaperlessTaskFactory
 from paperless.models import ApplicationConfiguration
 from paperless_ai import indexing
 from paperless_ai.tests.conftest import FakeEmbedding
+from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
 
 
 @pytest.fixture
@@ -67,8 +69,6 @@ def test_build_document_node_excludes_metadata_from_embedding(
     double the token count and exceed embedding models with small context
     windows (e.g. nomic-embed-text via Ollama defaults to num_ctx=2048).
     """
-    from llama_index.core.schema import MetadataMode
-
     nodes = indexing.build_document_node(real_document)
     for node in nodes:
         embed_text = node.get_content(metadata_mode=MetadataMode.EMBED)
@@ -100,8 +100,6 @@ def test_build_document_node_excludes_document_id_from_llm_context(
     real_document: Document,
 ) -> None:
     """document_id is an internal key and must not appear in LLM context text."""
-    from llama_index.core.schema import MetadataMode
-
     nodes = indexing.build_document_node(real_document)
     assert len(nodes) > 0
     for node in nodes:
@@ -670,8 +668,6 @@ class TestVectorStoreIndexing:
         temp_llm_index_dir: Path,
         mock_embed_model: FakeEmbedding,
     ) -> None:
-        from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
-
         store = indexing.get_vector_store()
         assert isinstance(store, PaperlessSqliteVecVectorStore)
 

src/paperless_ai/tests/test_chat.py

@@ -3,9 +3,13 @@ from unittest.mock import MagicMock
 from unittest.mock import patch
 
 import pytest
+from llama_index.core import settings as llama_settings
+from llama_index.core.embeddings.mock_embed_model import MockEmbedding
 from llama_index.core.schema import TextNode
 
+from documents.tests.factories import DocumentFactory
 from paperless_ai import chat
+from paperless_ai import indexing
 from paperless_ai.chat import CHAT_ERROR_MESSAGE
 from paperless_ai.chat import CHAT_METADATA_DELIMITER
 from paperless_ai.chat import stream_chat_with_documents
@@ -13,9 +17,6 @@ from paperless_ai.chat import stream_chat_with_documents
 
 @pytest.fixture(autouse=True)
 def patch_embed_model():
-    from llama_index.core import settings as llama_settings
-    from llama_index.core.embeddings.mock_embed_model import MockEmbedding
-
     # Use a real BaseEmbedding subclass to satisfy llama-index 0.14 validation
     llama_settings.Settings.embed_model = MockEmbedding(embed_dim=1536)
     yield
@@ -241,8 +242,6 @@ class TestStreamChatRetrieval:
         temp_llm_index_dir,
         mock_embed_model,
     ) -> None:
-        from documents.tests.factories import DocumentFactory
-
         doc = DocumentFactory.create(content="hello world")
         # Nothing indexed for this document yet.
         out = list(chat.stream_chat_with_documents("question?", [doc]))
@@ -258,9 +257,6 @@ class TestStreamChatRetrieval:
         requested documents only — content from other indexed documents must
         not be surfaced.
         """
-        from documents.tests.factories import DocumentFactory
-        from paperless_ai import indexing
-
         included = DocumentFactory.create(content="included document content")
         excluded = DocumentFactory.create(content="excluded document content")
         indexing.llm_index_add_or_update_document(included)

src/paperless_ai/tests/test_vector_store.py

@@ -3,6 +3,10 @@ from pathlib import Path
 
 import pytest
 from llama_index.core.schema import TextNode
+from llama_index.core.vector_stores.types import FilterOperator
+from llama_index.core.vector_stores.types import MetadataFilter
+from llama_index.core.vector_stores.types import MetadataFilters
+from llama_index.core.vector_stores.types import VectorStoreQuery
 
 from paperless_ai.vector_store import DB_FILENAME
 from paperless_ai.vector_store import DEFAULT_TABLE_NAME
@@ -43,8 +47,6 @@ def _query(
     top_k: int = 5,
     filters=None,
 ):
-    from llama_index.core.vector_stores.types import VectorStoreQuery
-
     return store.query(
         VectorStoreQuery(
             query_embedding=embedding,
@@ -55,20 +57,12 @@ def _query(
 
 
 def _eq_filter(key: str, value: str):
-    from llama_index.core.vector_stores.types import FilterOperator
-    from llama_index.core.vector_stores.types import MetadataFilter
-    from llama_index.core.vector_stores.types import MetadataFilters
-
     return MetadataFilters(
         filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)],
     )
 
 
 def _in_filter(document_ids: list[str]):
-    from llama_index.core.vector_stores.types import FilterOperator
-    from llama_index.core.vector_stores.types import MetadataFilter
-    from llama_index.core.vector_stores.types import MetadataFilters
-
     return MetadataFilters(
         filters=[
             MetadataFilter(
@@ -454,3 +448,84 @@ class TestMigrations:
         assert self._schema_version(store) == SCHEMA_VERSION
         store.compact(force=True)
         assert self._schema_version(store) == SCHEMA_VERSION
+
+    def test_stop_at_reembed_boundary(self, store) -> None:
+        # Registry: structural v2, re-embed v3, structural v4.
+        # Only v2 should apply; the re-embed boundary must stop execution
+        # before v4 runs, and the stored version must stay at 2.
+        store.add([make_node("a1", "1"), make_node("b1", "2")])
+
+        def copy_apply(
+            src: sqlite3.Connection,
+            dst: sqlite3.Connection,
+            dim: int,
+        ) -> None:
+            dst.execute(  # nosemgrep
+                f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0("
+                "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT,"
+                f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine"
+                ")",
+            )
+            dst.execute(
+                "INSERT INTO index_meta (key, value) VALUES ('dim', ?) "
+                "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+                (str(dim),),
+            )
+            rows = src.execute(
+                "SELECT id, document_id, modified, node_content, embedding "
+                f"FROM {DEFAULT_TABLE_NAME}",
+            ).fetchall()
+            dst.execute("BEGIN IMMEDIATE")
+            dst.executemany(
+                f"INSERT INTO {DEFAULT_TABLE_NAME} "
+                "(id, document_id, modified, node_content, embedding) "
+                "VALUES (?, ?, ?, ?, ?)",
+                [
+                    (
+                        r["id"],
+                        r["document_id"],
+                        r["modified"],
+                        r["node_content"],
+                        bytes(r["embedding"]),
+                    )
+                    for r in rows
+                ],
+            )
+            dst.execute("COMMIT")
+
+        migrations = [
+            Migration(
+                from_version=1,
+                to_version=2,
+                kind="structural",
+                description="v2 structural",
+                apply=copy_apply,
+            ),
+            Migration(
+                from_version=2,
+                to_version=3,
+                kind="re-embed",
+                description="v3 re-embed boundary",
+            ),
+            Migration(
+                from_version=3,
+                to_version=4,
+                kind="structural",
+                description="v4 structural - must not run",
+                apply=copy_apply,
+            ),
+        ]
+        MIGRATIONS.extend(migrations)
+        try:
+            from paperless_ai import vector_store as vs_mod
+
+            original = vs_mod.SCHEMA_VERSION
+            vs_mod.SCHEMA_VERSION = 4
+            result = store.check_and_run_migrations()
+        finally:
+            for m in migrations:
+                MIGRATIONS.remove(m)
+            vs_mod.SCHEMA_VERSION = original
+
+        assert result is True
+        assert self._schema_version(store) == 2