diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 59cfb6d9d..30d722a47 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -1,3 +1,4 @@ +import enum import logging from collections.abc import Iterable from contextlib import contextmanager @@ -192,18 +193,49 @@ def write_store(embed_model_name: str | None = None): yield store -def _check_and_run_migrations(store: "PaperlessSqliteVecVectorStore") -> bool: - """Run any pending structural migrations, returning True if a pending - re-embed migration needs the caller to force a rebuild -- never - triggered automatically here. Safe to call before any write, including +class MigrationCheckResult(enum.Enum): + """Outcome of _check_and_run_migrations(). + + CURRENT: no migration was pending, or a pending structural migration + was applied successfully -- safe to write. + + REEMBED_REQUIRED: a pending migration needs fresh embeddings, which is + never triggered automatically -- the caller must force a rebuild. + + DEFERRED: a migration was pending but could not run because active + index readers did not drain within LLM_INDEX_COMPACTION_LOCK_TIMEOUT -- + the store is still on its old schema. Callers must NOT proceed to + write: collapsing this into the same falsy value as CURRENT (as a + plain bool return once did) would let a write proceed against an + unmigrated schema. + """ + + CURRENT = "current" + REEMBED_REQUIRED = "reembed_required" + DEFERRED = "deferred" + + +def _check_and_run_migrations( + store: "PaperlessSqliteVecVectorStore", +) -> MigrationCheckResult: + """Run any pending structural migrations, reporting the outcome as a + tri-state result. Safe to call before any write, including delete()/upsert_document(): has_pending_migration() (see its docstring) keeps this a no-op, with no exclusive access taken, once the store is current. """ if not store.has_pending_migration(): - return False - return bool( - _with_exclusive_access("migration check", store.check_and_run_migrations), + return MigrationCheckResult.CURRENT + result = _with_exclusive_access( + "migration check", + store.check_and_run_migrations, + ) + if result is None: + return MigrationCheckResult.DEFERRED + return ( + MigrationCheckResult.REEMBED_REQUIRED + if result + else MigrationCheckResult.CURRENT ) @@ -378,12 +410,21 @@ def update_llm_index( happens, since a rebuild always covers the whole library regardless. """ with write_store() as store: - needs_reembed = _check_and_run_migrations(store) - if needs_reembed: + migration_result = _check_and_run_migrations(store) + if migration_result is MigrationCheckResult.REEMBED_REQUIRED: logger.warning( "LLM index migration requires re-embedding; forcing rebuild.", ) rebuild = True + elif migration_result is MigrationCheckResult.DEFERRED: + logger.info( + "Skipping LLM index update: migration check deferred while " + "index readers are active; will retry next run.", + ) + return ( + "Skipping LLM index update: migration check deferred; " + "will retry next run." + ) documents = Document.objects.select_related( "correspondent", "document_type", @@ -461,8 +502,8 @@ def llm_index_add_or_update_document(document: Document): _embed_nodes(new_nodes, get_embedding_model(config)) with write_store(embed_model_name=get_configured_model_name(config)) as store: - needs_reembed = _check_and_run_migrations(store) - if needs_reembed: + migration_result = _check_and_run_migrations(store) + if migration_result is MigrationCheckResult.REEMBED_REQUIRED: logger.warning( "Skipping incremental LLM index update for document %s: the " "index requires re-embedding first. Run 'document_llmindex " @@ -470,6 +511,14 @@ def llm_index_add_or_update_document(document: Document): document.id, ) return + if migration_result is MigrationCheckResult.DEFERRED: + logger.info( + "Skipping incremental LLM index update for document %s: " + "migration check deferred while index readers are active; " + "will retry on the next write.", + document.id, + ) + return store.upsert_document(str(document.id), new_nodes) @@ -488,14 +537,19 @@ def llm_index_migrate() -> None: if not AIConfig().llm_index_enabled: return with write_store() as store: - needs_reembed = _check_and_run_migrations(store) - if needs_reembed: + migration_result = _check_and_run_migrations(store) + if migration_result is MigrationCheckResult.REEMBED_REQUIRED: logger.warning( "LLM index requires re-embedding, which this automatic migration " "check will not do on its own -- it can be slow and, for a " "metered embedding backend, cost money. Run " "'document_llmindex rebuild' manually when ready.", ) + elif migration_result is MigrationCheckResult.DEFERRED: + logger.info( + "LLM index migration check deferred while index readers are " + "active; will retry next run.", + ) def llm_index_compact() -> None: @@ -507,7 +561,8 @@ def llm_index_compact() -> None: def llm_index_remove_document(document: Document): """Remove a document's chunks from the LLM index.""" with write_store() as store: - if _check_and_run_migrations(store): + migration_result = _check_and_run_migrations(store) + if migration_result is MigrationCheckResult.REEMBED_REQUIRED: logger.warning( "Skipping removal of document %s from the LLM index: the " "index requires re-embedding first. Run 'document_llmindex " @@ -515,6 +570,14 @@ def llm_index_remove_document(document: Document): document.id, ) return + if migration_result is MigrationCheckResult.DEFERRED: + logger.info( + "Skipping removal of document %s from the LLM index: " + "migration check deferred while index readers are active; " + "will retry on the next write.", + document.id, + ) + return store.delete(str(document.id)) diff --git a/src/paperless_ai/migrations/m0001_v1_to_v2.py b/src/paperless_ai/migrations/m0001_v1_to_v2.py index c7edf7f41..18409d5fc 100644 --- a/src/paperless_ai/migrations/m0001_v1_to_v2.py +++ b/src/paperless_ai/migrations/m0001_v1_to_v2.py @@ -7,7 +7,7 @@ from paperless_ai.tables import DocumentChunksTable from paperless_ai.tables import DocumentMetaRow from paperless_ai.tables import DocumentMetaTable from paperless_ai.tables import IndexMetaTable -from paperless_ai.vector_store import COMPACT_BATCH_SIZE +from paperless_ai.vector_store import BATCH_SIZE from paperless_ai.vector_store import DEFAULT_TABLE_NAME # v1's vec0 shape has never changed since it first shipped and is the ONLY @@ -75,7 +75,7 @@ def _migrate_v1_to_v2( dst_conn.execute("BEGIN IMMEDIATE") src_cursor = src_conn.execute(_V1_SELECT) live = 0 - while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): + while batch := src_cursor.fetchmany(BATCH_SIZE): vec0_rows = [] chunk_rows = [] meta_by_document: dict[int, str] = {} diff --git a/src/paperless_ai/tables.py b/src/paperless_ai/tables.py index f3ac5710f..c8d6d838c 100644 --- a/src/paperless_ai/tables.py +++ b/src/paperless_ai/tables.py @@ -131,7 +131,7 @@ class DocumentMetaTable: connections (compact()/migrations) -- an unbounded fetchall here would defeat the same OOM-avoidance the vec0 row copy already relies on. batch_size has no default: forces the call site to think about - it (pass COMPACT_BATCH_SIZE).""" + it (pass BATCH_SIZE).""" cursor = src_conn.execute( "SELECT document_id, modified FROM document_meta", ) diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 25d0a7204..2da225053 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -9,6 +9,7 @@ from django.db import connection from django.test import override_settings from django.test.utils import CaptureQueriesContext from django.utils import timezone +from filelock import Timeout from llama_index.core.schema import MetadataMode from documents.models import Correspondent @@ -830,6 +831,44 @@ class TestLlmIndexLocking: mock_store.upsert_document.assert_not_called() + def test_add_or_update_document_skips_write_when_migration_check_deferred( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + mocker: pytest_mock.MockerFixture, + ) -> None: + """A migration check that times out waiting for readers to drain + must be treated the same as a pending migration -- proceeding to + write would target a store still on its old schema. Regression + test for the tri-state fix: a bare bool collapsed this outcome + into the same falsy value as "already current". + """ + mock_store = MagicMock() + mock_store.has_pending_migration.return_value = True + mocker.patch( + "paperless_ai.indexing.write_store", + return_value=mocker.MagicMock( + __enter__=mocker.MagicMock(return_value=mock_store), + __exit__=mocker.MagicMock(return_value=False), + ), + ) + mocker.patch( + "paperless_ai.indexing._exclude_readers", + side_effect=Timeout("test"), + ) + mock_node = MagicMock() + mock_node.get_content.return_value = "fake node text" + mocker.patch( + "paperless_ai.indexing.build_document_node", + return_value=[mock_node], + ) + + doc = MagicMock(spec=Document) + doc.id = 1 + indexing.llm_index_add_or_update_document(doc) + + mock_store.upsert_document.assert_not_called() + def test_remove_document_uses_write_store( self, temp_llm_index_dir: Path, @@ -876,6 +915,34 @@ class TestLlmIndexLocking: mock_store.delete.assert_not_called() + def test_remove_document_skips_write_when_migration_check_deferred( + self, + temp_llm_index_dir: Path, + mocker: pytest_mock.MockerFixture, + ) -> None: + """A migration check deferred by a reader-lock timeout must block + the delete too, for the same reason as the incremental-update path. + """ + mock_store = MagicMock() + mock_store.has_pending_migration.return_value = True + mocker.patch( + "paperless_ai.indexing.write_store", + return_value=mocker.MagicMock( + __enter__=mocker.MagicMock(return_value=mock_store), + __exit__=mocker.MagicMock(return_value=False), + ), + ) + mocker.patch( + "paperless_ai.indexing._exclude_readers", + side_effect=Timeout("test"), + ) + + doc = MagicMock(spec=Document) + doc.id = 1 + indexing.llm_index_remove_document(doc) + + mock_store.delete.assert_not_called() + def test_update_llm_index_rebuild_uses_write_store( self, temp_llm_index_dir: Path, @@ -899,6 +966,35 @@ class TestLlmIndexLocking: mock_store.drop_table.assert_called_once() + def test_update_llm_index_skips_when_migration_check_deferred( + self, + temp_llm_index_dir: Path, + mocker: pytest_mock.MockerFixture, + ) -> None: + """A migration check deferred by a reader-lock timeout must short- + circuit before the second write_store() block (document scanning, + add/upsert, compaction) ever runs -- that block would otherwise + write against a store still on its old schema. + """ + mock_store = MagicMock() + mock_store.has_pending_migration.return_value = True + write_store_mock = mocker.patch( + "paperless_ai.indexing.write_store", + return_value=mocker.MagicMock( + __enter__=mocker.MagicMock(return_value=mock_store), + __exit__=mocker.MagicMock(return_value=False), + ), + ) + mocker.patch( + "paperless_ai.indexing._exclude_readers", + side_effect=Timeout("test"), + ) + + result = indexing.update_llm_index(rebuild=False) + + assert "deferred" in result + write_store_mock.assert_called_once() + @pytest.mark.django_db @pytest.mark.django_db @@ -1017,6 +1113,37 @@ class TestLlmIndexMigrate: indexing.llm_index_migrate() assert "requires re-embedding" in caplog.text + def test_logs_info_when_migration_check_deferred( + self, + mocker: pytest_mock.MockerFixture, + caplog: pytest.LogCaptureFixture, + ) -> None: + """ + GIVEN: + - AI/LLM index support is enabled + - A pending migration cannot run because readers are active + WHEN: + - llm_index_migrate() is called + THEN: + - An info line notes the deferral, not the re-embed warning + """ + mocker.patch( + "paperless_ai.indexing.AIConfig", + return_value=mocker.Mock(llm_index_enabled=True), + ) + store_mock = mocker.MagicMock() + store_mock.has_pending_migration.return_value = True + write_store_cm = mocker.patch("paperless_ai.indexing.write_store") + write_store_cm.return_value.__enter__.return_value = store_mock + mocker.patch( + "paperless_ai.indexing._exclude_readers", + side_effect=Timeout("test"), + ) + with caplog.at_level(logging.INFO, logger="paperless_ai.indexing"): + indexing.llm_index_migrate() + assert "deferred" in caplog.text + assert "requires re-embedding" not in caplog.text + @pytest.mark.django_db class TestQuerySimilarDocuments: diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index e4cf883bb..3d2dc13f5 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -508,7 +508,7 @@ class TestCompact: A tiny batch size forces several fetchmany()/executemany() cycles so a regression in the streaming loop (dropped tail, off-by-one) surfaces. """ - monkeypatch.setattr("paperless_ai.vector_store.COMPACT_BATCH_SIZE", 3) + monkeypatch.setattr("paperless_ai.vector_store.BATCH_SIZE", 3) store.add([make_node(f"n{i}", 1, seed=float(i)) for i in range(10)]) store.compact(force=True) ids = {n.node_id for n in store.get_nodes(filters=_in_filter([1]))} diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index bb3007323..975b6c57b 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -47,10 +47,12 @@ SCHEMA_VERSION = 2 # a rebuild copies the live rows into a fresh table. COMPACT_BLOAT_RATIO = 2.0 -# compact(): number of rows copied per executemany() when rebuilding the file. -# Rows are streamed from the source cursor in batches of this size rather than -# materialized all at once, keeping memory bounded regardless of index size. -COMPACT_BATCH_SIZE = 500 +# Number of rows fetched/copied per batch whenever this module streams rows +# instead of materializing them all at once, keeping memory bounded regardless +# of index size -- used by compact()'s rebuild, m0001_v1_to_v2's migration +# copy, and DocumentMetaTable.copy_all(). No longer compact()-specific, hence +# the plain name. +BATCH_SIZE = 500 # Filterable vec0 metadata columns. _build_where() only ever receives filter # keys we construct ourselves, but allowlisting keeps SQL identifiers safe by @@ -605,7 +607,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): + DEFAULT_TABLE_NAME, ) copied = 0 - while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): + while batch := src_cursor.fetchmany(BATCH_SIZE): dst_conn.executemany( _INSERT, [ @@ -623,7 +625,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): (ChunkRow(r["id"], r["document_id"]) for r in batch), ) copied += len(batch) - DocumentMetaTable.copy_all(src_conn, dst_conn, COMPACT_BATCH_SIZE) + DocumentMetaTable.copy_all(src_conn, dst_conn, BATCH_SIZE) # Reset the cumulative counter: after a rebuild, total_inserts == live. IndexMetaTable.reset_total_inserts(dst_conn, copied) dst_conn.execute("COMMIT")