diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 9678f8fc0..d4d381a40 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 @@ -186,18 +187,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 ) @@ -372,12 +404,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", @@ -451,8 +492,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 " @@ -460,6 +501,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) @@ -478,14 +527,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: @@ -497,7 +551,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 " @@ -505,6 +560,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/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 25d0a7204..296554b3f 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, @@ -1017,6 +1084,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: