diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index f245d9100..cab2dda8a 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/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) diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index af34914bb..b55f9a03c 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/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) diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 4921243de..2e48b20a1 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/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