diff --git a/docker/rootfs/etc/s6-overlay/s6-rc.d/init-complete/dependencies.d/init-llmindex-migrate b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-complete/dependencies.d/init-llmindex-migrate new file mode 100644 index 000000000..e69de29bb diff --git a/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/dependencies.d/init-migrations b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/dependencies.d/init-migrations new file mode 100644 index 000000000..e69de29bb diff --git a/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/run b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/run new file mode 100644 index 000000000..50fa691c2 --- /dev/null +++ b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/run @@ -0,0 +1,12 @@ +#!/command/with-contenv /usr/bin/bash +# shellcheck shell=bash + +declare -r log_prefix="[init-llmindex-migrate]" + +echo "${log_prefix} Checking for pending LLM index migrations..." +cd "${PAPERLESS_SRC_DIR}" +if [[ -n "${USER_IS_NON_ROOT}" ]]; then + python3 manage.py document_llmindex migrate +else + s6-setuidgid paperless python3 manage.py document_llmindex migrate +fi diff --git a/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/type b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/type new file mode 100644 index 000000000..bdd22a185 --- /dev/null +++ b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/type @@ -0,0 +1 @@ +oneshot diff --git a/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/up b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/up new file mode 100644 index 000000000..c2016d47a --- /dev/null +++ b/docker/rootfs/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/up @@ -0,0 +1 @@ +/etc/s6-overlay/s6-rc.d/init-llmindex-migrate/run diff --git a/src/documents/management/commands/document_llmindex.py b/src/documents/management/commands/document_llmindex.py index 7b34ca9a8..216ab631d 100644 --- a/src/documents/management/commands/document_llmindex.py +++ b/src/documents/management/commands/document_llmindex.py @@ -3,6 +3,7 @@ from typing import Any from documents.management.commands.base import PaperlessCommand from documents.tasks import llmindex_index from paperless_ai.indexing import llm_index_compact +from paperless_ai.indexing import llm_index_migrate class Command(PaperlessCommand): @@ -13,12 +14,18 @@ class Command(PaperlessCommand): def add_arguments(self, parser: Any) -> None: super().add_arguments(parser) - parser.add_argument("command", choices=["rebuild", "update", "compact"]) + parser.add_argument( + "command", + choices=["rebuild", "update", "compact", "migrate"], + ) def handle(self, *args: Any, **options: Any) -> None: if options["command"] == "compact": llm_index_compact() return + if options["command"] == "migrate": + llm_index_migrate() + return llmindex_index( rebuild=options["command"] == "rebuild", iter_wrapper=lambda docs: self.track( diff --git a/src/documents/tests/management/test_management_document_llmindex.py b/src/documents/tests/management/test_management_document_llmindex.py index b8a05dd85..3b75338d2 100644 --- a/src/documents/tests/management/test_management_document_llmindex.py +++ b/src/documents/tests/management/test_management_document_llmindex.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: _COMPACT = "documents.management.commands.document_llmindex.llm_index_compact" _INDEX = "documents.management.commands.document_llmindex.llmindex_index" +_MIGRATE = "documents.management.commands.document_llmindex.llm_index_migrate" class TestDocumentLlmindexCommand: @@ -17,6 +18,11 @@ class TestDocumentLlmindexCommand: call_command("document_llmindex", "compact") mock_compact.assert_called_once_with() + def test_migrate_calls_llm_index_migrate(self, mocker: MockerFixture) -> None: + mock_migrate = mocker.patch(_MIGRATE) + call_command("document_llmindex", "migrate") + mock_migrate.assert_called_once_with() + def test_rebuild_calls_llmindex_index_with_rebuild_true( self, mocker: MockerFixture, diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 93b850cbf..88e29bace 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -144,6 +144,24 @@ def _exclude_readers(): lock.close() +def _with_exclusive_access(operation: str, fn): + """Run ``fn()`` with exclusive index access (see ``_exclude_readers()``), + for compaction/migration file swaps that must not run while readers are + active. Returns ``fn()``'s result, or None (after logging) if active + readers do not drain within ``LLM_INDEX_COMPACTION_LOCK_TIMEOUT`` -- + callers skip the operation this run; it retries next time. + """ + try: + with _exclude_readers(): + return fn() + except Timeout: + logger.info( + "Skipping LLM index %s: index readers are active; will retry next run.", + operation, + ) + return None + + @contextmanager def write_store(embed_model_name: str | None = None): """Acquire the write lock and yield the vector store. @@ -168,6 +186,21 @@ 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 + 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), + ) + + def _safe_related_name(document: Document, field: str) -> str | None: """ Returns the ``name`` of a related object (correspondent, document_type, @@ -339,15 +372,7 @@ def update_llm_index( happens, since a rebuild always covers the whole library regardless. """ with write_store() as store: - try: - with _exclude_readers(): - needs_reembed = store.check_and_run_migrations() - except Timeout: - logger.info( - "Skipping LLM index migration check: index readers are active; " - "will retry next run.", - ) - needs_reembed = False + needs_reembed = _check_and_run_migrations(store) if needs_reembed: logger.warning( "LLM index migration requires re-embedding; forcing rebuild.", @@ -412,14 +437,7 @@ def update_llm_index( else "No changes detected in LLM index." ) - try: - with _exclude_readers(): - store.compact() - except Timeout: - logger.info( - "Skipping LLM index compaction: index readers are active; " - "will retry next run.", - ) + _with_exclusive_access("compaction", store.compact) return msg @@ -434,25 +452,45 @@ 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: + _check_and_run_migrations(store) store.upsert_document(str(document.id), new_nodes) +def llm_index_migrate() -> None: + """Apply any pending LLM index schema migrations, with no reindex. + + Intended to run unconditionally on every startup (see the + init-llmindex-migrate container step and the bare-metal upgrade docs): + has_pending_migration() short-circuits to a metadata-only read once the + store is current, so a healthy install pays almost nothing here. Only + ever applies structural migrations -- a pending re-embed migration is + left for the explicit, deliberate rebuild path (``document_llmindex + update``/``rebuild``) to resolve, since re-embedding can be slow and, + for a metered embedding backend, cost money. + """ + if not AIConfig().llm_index_enabled: + return + with write_store() as store: + needs_reembed = _check_and_run_migrations(store) + if needs_reembed: + 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.", + ) + + def llm_index_compact() -> None: """Compact the index immediately, rebuilding the table to reclaim space.""" with write_store() as store: - try: - with _exclude_readers(): - store.compact(force=True) - except Timeout: - logger.info( - "Skipping LLM index compaction: index readers are active; " - "will retry next run.", - ) + _with_exclusive_access("compaction", lambda: store.compact(force=True)) def llm_index_remove_document(document: Document): """Remove a document's chunks from the LLM index.""" with write_store() as store: + _check_and_run_migrations(store) 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 667b43d06..406224ba7 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -849,6 +849,48 @@ class TestVectorStoreIndexing: assert rows >= 1 +class TestLlmIndexMigrate: + def test_noop_when_ai_disabled(self, mocker: pytest_mock.MockerFixture) -> None: + """ + GIVEN: + - AI/LLM index support is disabled in configuration + WHEN: + - llm_index_migrate() is called + THEN: + - No store is opened and no migration check runs + """ + mocker.patch( + "paperless_ai.indexing.AIConfig", + return_value=mocker.Mock(llm_index_enabled=False), + ) + write_store_mock = mocker.patch("paperless_ai.indexing.write_store") + indexing.llm_index_migrate() + write_store_mock.assert_not_called() + + def test_runs_pending_migration_when_enabled( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - AI/LLM index support is enabled + WHEN: + - llm_index_migrate() is called + THEN: + - The store is opened for write and a migration check runs + """ + 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 = False + write_store_cm = mocker.patch("paperless_ai.indexing.write_store") + write_store_cm.return_value.__enter__.return_value = store_mock + indexing.llm_index_migrate() + store_mock.has_pending_migration.assert_called_once() + + @pytest.mark.django_db class TestQuerySimilarDocuments: def test_query_similar_documents_respects_allowed_ids(