diff --git a/src/documents/management/commands/document_index.py b/src/documents/management/commands/document_index.py index 45dd98824..a5e8fd090 100644 --- a/src/documents/management/commands/document_index.py +++ b/src/documents/management/commands/document_index.py @@ -67,18 +67,22 @@ class Command(PaperlessCommand): if options.get("recreate"): wipe_index(settings.INDEX_DIR) - documents = Document.objects.select_related( - "correspondent", - "document_type", - "storage_path", - "owner", - ).prefetch_related( - "tags", - "notes__user", - "custom_fields__field", - "versions", - "barcodes", - "versions__barcodes", + documents = ( + Document.objects.filter(root_document__isnull=True) + .select_related( + "correspondent", + "document_type", + "storage_path", + "owner", + ) + .prefetch_related( + "tags", + "notes__user", + "custom_fields__field", + "versions", + "barcodes", + "versions__barcodes", + ) ) total = documents.count() rebuild_kwargs = {} diff --git a/src/documents/search/_backend.py b/src/documents/search/_backend.py index de11f215b..930a7f48c 100644 --- a/src/documents/search/_backend.py +++ b/src/documents/search/_backend.py @@ -284,9 +284,14 @@ class WriteBatch: and adding the new version. This ensures stale document data (e.g., after permission changes) doesn't persist in the index. + Only root documents are indexed, with their effective content, so a + version is indexed as its root document. + Args: document: Django Document instance to index """ + if document.root_document_id is not None: + document = document.root_document self.remove(document.pk) doc = self._backend._build_tantivy_doc(document) self._writer.add_document(doc) @@ -311,20 +316,22 @@ class WriteBatch: An id with no matching document (e.g. deleted between the caller collecting ids and the batch running) is silently skipped, matching ``add_or_update()``'s existing single-document deferred-task behavior - rather than erroring or leaving a stale index entry. + rather than erroring or leaving a stale index entry. The id of a + version stands for its root document. Args: ids: Primary keys of Document instances to index """ from documents.models import Document from documents.versioning import annotate_effective_content + from documents.versioning import root_document_ids ids = list(ids) if not ids: return queryset = annotate_effective_content( - Document.objects.filter(pk__in=ids) + Document.objects.filter(pk__in=root_document_ids(ids)) .select_related("correspondent", "document_type", "storage_path", "owner") .prefetch_related( "tags", diff --git a/src/documents/tasks.py b/src/documents/tasks.py index 0a7409f17..778ee9673 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -490,23 +490,20 @@ def update_document_content_maybe_archive_file( shutil.move(thumbnail, document.thumbnail_path) document.refresh_from_db() - root_document = ( - document.root_document if document.root_document_id else document - ) logger.info( - f"Updating index for document {root_document.pk} ({document.archive_checksum})", + f"Updating index for document {document_id} ({document.archive_checksum})", ) from documents.search import get_backend - get_backend().add_or_update(root_document) + get_backend().add_or_update(document) ai_config = AIConfig() if ai_config.llm_index_enabled: - llm_index_add_or_update_document(root_document) + llm_index_add_or_update_document(document) clear_document_caches(document.pk) - if root_document.pk != document.pk: - clear_document_caches(root_document.pk) + if document.root_document_id is not None: + clear_document_caches(document.root_document_id) except Exception: logger.exception( diff --git a/src/documents/tests/search/test_backend.py b/src/documents/tests/search/test_backend.py index c97558fb4..4350280f7 100644 --- a/src/documents/tests/search/test_backend.py +++ b/src/documents/tests/search/test_backend.py @@ -293,6 +293,80 @@ class TestAddOrUpdateIds: assert backend.search_ids("updated", user=None) == [doc.pk] +class TestVersionsAreIndexedAsTheirRoot: + """Only root documents are indexed, with their effective content, so + every write path that is handed a version indexes its root instead.""" + + @staticmethod + def _root_with_version() -> tuple[Document, Document]: + root = DocumentFactory(title="Statement", content="stale text") + version = DocumentFactory( + title="Statement", + content="latest text", + root_document=root, + version_index=1, + ) + return root, version + + def test_add_or_update_indexes_the_root_of_a_version( + self, + backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - The version is passed to add_or_update + THEN: + - The root is indexed with the version's text, and the version is not + """ + root, version = self._root_with_version() + + backend.add_or_update(version) + + assert backend.search_ids("latest", user=None) == [root.pk] + assert backend.search_ids("stale", user=None) == [] + + def test_add_or_update_ids_indexes_each_root_once( + self, + backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - Both ids are passed to add_or_update_ids + THEN: + - The root is indexed once and the version is not indexed + """ + root, version = self._root_with_version() + + with backend.batch_update() as batch: + batch.add_or_update_ids([version.pk, root.pk]) + + assert backend.search_ids("Statement", user=None) == [root.pk] + assert backend.search_ids("latest", user=None) == [root.pk] + + def test_add_or_update_ids_resolves_a_lone_version( + self, + backend: TantivyBackend, + ) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - Only the version's id is passed to add_or_update_ids + THEN: + - The root is indexed + """ + root, version = self._root_with_version() + + with backend.batch_update() as batch: + batch.add_or_update_ids([version.pk]) + + assert backend.search_ids("latest", user=None) == [root.pk] + + class TestSearch: """Test search query parsing and matching via search_ids.""" diff --git a/src/documents/tests/test_api_search.py b/src/documents/tests/test_api_search.py index 102736139..20319e71e 100644 --- a/src/documents/tests/test_api_search.py +++ b/src/documents/tests/test_api_search.py @@ -1196,6 +1196,59 @@ class TestDocumentSearchApi(DirectoriesMixin, APITestCase): self.assertIn(d3.id, result_ids) self.assertNotIn(d4.id, result_ids) + def test_search_more_like_version_uses_its_root(self) -> None: + """ + GIVEN: + - A document similar in content to a root document, and one that is not + - A version of the root document, which is never indexed + WHEN: + - API request for more like the version + THEN: + - The documents similar to the root are returned, not the version + """ + indexed = {} + for name, title, content, day in ( + ("root", "bank statement 1", "things i paid for in august", (2019, 3, 4)), + ( + "similar", + "bank statement 3", + "things i paid for in september", + (2020, 7, 9), + ), + ( + "other", + "Quarterly Report", + "quarterly revenue profit margin", + (2021, 11, 30), + ), + ): + with time_machine.travel( + timezone.make_aware(datetime.datetime(*day)), + tick=False, + ): + indexed[name] = DocumentFactory( + title=title, + content=content, + created=datetime.date(*day), + added=timezone.make_aware(datetime.datetime(*day)), + ) + version = DocumentFactory( + root_document=indexed["root"], + version_index=1, + content="things i paid for in august", + ) + backend = get_backend() + for document in indexed.values(): + backend.add_or_update(document) + + response = self.client.get(f"/api/documents/?more_like_id={version.id}") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + result_ids = [r["id"] for r in response.data["results"]] + self.assertIn(indexed["similar"].id, result_ids) + self.assertNotIn(indexed["other"].id, result_ids) + self.assertNotIn(version.id, result_ids) + def test_more_like_requires_id_of_existing_document(self) -> None: """ GIVEN: diff --git a/src/documents/tests/test_management.py b/src/documents/tests/test_management.py index a5a2e6233..19ab4e269 100644 --- a/src/documents/tests/test_management.py +++ b/src/documents/tests/test_management.py @@ -22,6 +22,7 @@ from documents.models import Document from documents.tasks import update_document_content_maybe_archive_file from paperless_testing.assertions import FileSystemAssertsMixin from paperless_testing.dirs import DirectoriesMixin +from paperless_testing.factories import DocumentFactory sample_file: Path = Path(__file__).parent / "samples" / "simple.pdf" @@ -116,6 +117,27 @@ class TestMakeIndex: call_command("document_index", "reindex", skip_checks=True) mock_get_backend.return_value.rebuild.assert_called_once() + def test_reindex_skips_versions(self, mocker: MockerFixture) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - The reindex command runs + THEN: + - Only the root document is handed to the rebuild, since a version + is indexed as its root + """ + root = DocumentFactory() + DocumentFactory(root_document=root, version_index=1) + mock_get_backend = mocker.patch( + "documents.management.commands.document_index.get_backend", + ) + + call_command("document_index", "reindex", skip_checks=True) + + documents = mock_get_backend.return_value.rebuild.call_args.args[0] + assert list(documents.values_list("pk", flat=True)) == [root.pk] + def test_optimize(self) -> None: """Optimize command must execute without error (Tantivy handles optimization automatically).""" call_command("document_index", "optimize", skip_checks=True) diff --git a/src/documents/tests/test_tasks.py b/src/documents/tests/test_tasks.py index c1bea1b43..50d16233b 100644 --- a/src/documents/tests/test_tasks.py +++ b/src/documents/tests/test_tasks.py @@ -310,7 +310,7 @@ class TestUpdateContent(DirectoriesMixin, TestCase): @mock.patch("documents.tasks.clear_document_caches") @mock.patch("documents.search.get_backend") - def test_update_content_version_indexes_root( + def test_update_content_version_clears_caches_for_root( self, mock_get_backend: mock.Mock, mock_clear_caches: mock.Mock, @@ -321,8 +321,8 @@ class TestUpdateContent(DirectoriesMixin, TestCase): WHEN: - Update content task is called for the version THEN: - - The version's content is updated - - The root document is indexed rather than the version + - The version's content is updated, not the root's + - The document is indexed - Caches are cleared for both """ root, version = self._create_root_with_version() @@ -334,8 +334,7 @@ class TestUpdateContent(DirectoriesMixin, TestCase): "my document", ) self.assertEqual(Document.objects.get(pk=root.pk).content, "root content") - indexed = mock_get_backend.return_value.add_or_update.call_args.args[0] - self.assertEqual(indexed.pk, root.pk) + mock_get_backend.return_value.add_or_update.assert_called_once() mock_clear_caches.assert_has_calls( [mock.call(version.pk), mock.call(root.pk)], ) @@ -343,7 +342,7 @@ class TestUpdateContent(DirectoriesMixin, TestCase): @override_settings(AI_ENABLED=True, LLM_EMBEDDING_BACKEND="huggingface") @mock.patch("documents.tasks.llm_index_add_or_update_document") @mock.patch("documents.search.get_backend") - def test_update_content_version_updates_llm_index_for_root( + def test_update_content_version_updates_llm_index( self, mock_get_backend: mock.Mock, mock_llm_index: mock.Mock, @@ -355,14 +354,13 @@ class TestUpdateContent(DirectoriesMixin, TestCase): WHEN: - Update content task is called for the version THEN: - - The LLM index is updated for the root document, not the version + - The LLM index is updated """ - root, version = self._create_root_with_version() + _, version = self._create_root_with_version() tasks.update_document_content_maybe_archive_file(version.pk) mock_llm_index.assert_called_once() - self.assertEqual(mock_llm_index.call_args.args[0].pk, root.pk) class TestUpdateContentRemoteOCR(DirectoriesMixin, TestCase): diff --git a/src/documents/versioning.py b/src/documents/versioning.py index a3c005587..512b0c6c6 100644 --- a/src/documents/versioning.py +++ b/src/documents/versioning.py @@ -17,9 +17,26 @@ from django.db.models.functions import RowNumber from documents.models import Document if TYPE_CHECKING: + from collections.abc import Iterable + from rest_framework.request import Request +def root_document_ids(ids: Iterable[int]) -> QuerySet[int]: + """ + The ids of the root documents of the given documents: a root stands for + itself and a version for its root. Only the indexes' bookkeeping needs + this, since they hold root documents only. + """ + return ( + Document.objects.filter(pk__in=ids) + .annotate(root_id=Coalesce("root_document_id", "id")) + .order_by() + .values_list("root_id", flat=True) + .distinct() + ) + + def versions_newest_first(documents: QuerySet[Document]) -> QuerySet[Document]: """ Sorts versions so the newest one comes first using version_index and not on id, diff --git a/src/documents/views.py b/src/documents/views.py index 74ba694bd..7e5ab806f 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -329,9 +329,10 @@ def _get_tantivy_query_and_mode(params): def _get_more_like_id(query_params: dict[str, Any], user: User | None) -> int: try: more_like_doc_id = int(query_params["more_like_id"]) - more_like_doc = Document.objects.select_related("owner").get( - pk=more_like_doc_id, - ) + more_like_doc = Document.objects.select_related( + "owner", + "root_document__owner", + ).get(pk=more_like_doc_id) except (TypeError, ValueError, Document.DoesNotExist): raise PermissionDenied(_("Invalid more_like_id")) @@ -342,7 +343,8 @@ def _get_more_like_id(query_params: dict[str, Any], user: User | None) -> int: ): raise PermissionDenied(_("Insufficient permissions.")) - return more_like_doc_id + # Only root documents are indexed, a version stands for its root + return more_like_doc.root_document_id or more_like_doc.pk class SearchParams(NamedTuple): diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index 07c20b526..681ef7ad5 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -17,6 +17,7 @@ from documents.models import PaperlessTask from documents.utils import IterWrapper from documents.utils import QuerySetStream from documents.utils import identity +from documents.versioning import root_document_ids from paperless.config import AIConfig from paperless_ai.db import db_connection_released from paperless_ai.embedding import build_llm_index_text @@ -443,11 +444,11 @@ def update_llm_index( "Skipping LLM index update: migration check deferred; " "will retry next run." ) - documents = Document.objects.select_related( - "correspondent", - "document_type", - "storage_path", - ).prefetch_related("tags", "notes", "custom_fields__field") + documents = ( + Document.objects.filter(root_document__isnull=True) + .select_related("correspondent", "document_type", "storage_path") + .prefetch_related("tags", "notes", "custom_fields__field") + ) no_documents = not documents.exists() # Fast exit before touching config: nothing to index and no existing index. @@ -483,7 +484,7 @@ def update_llm_index( msg = "LLM index rebuilt successfully." else: scoped_documents = ( - documents.filter(id__in=document_ids) + documents.filter(id__in=root_document_ids(document_ids)) if document_ids is not None else documents ) @@ -510,7 +511,12 @@ def update_llm_index( def llm_index_add_or_update_document(document: Document): - """Add or atomically replace a document's chunks in the index.""" + """ + Add or atomically replace a document's chunks in the index. Only root + documents are indexed, so a version is indexed as its root document. + """ + if document.root_document_id is not None: + document = document.root_document config = AIConfig() new_nodes = build_document_node( document, diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 99a11f35f..8c91dcf3a 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -388,6 +388,83 @@ def test_update_llm_index_partial_update( assert after[str(doc2.pk)] == before[str(doc2.pk)] +@pytest.mark.django_db +class TestLlmIndexVersions: + """The LLM index holds root documents only: a version is indexed as its root.""" + + def test_add_or_update_document_indexes_a_version_as_its_root( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - The version is passed to llm_index_add_or_update_document + THEN: + - Only the root document is in the index + """ + root = DocumentFactory(content="root content") + version = DocumentFactory(root_document=root, version_index=1) + + indexing.llm_index_add_or_update_document(version) + + with indexing.get_vector_store() as store: + indexed = store.get_modified_times() + + assert set(indexed) == {str(root.pk)} + + def test_rebuild_skips_versions( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + """ + GIVEN: + - A root document with a version + WHEN: + - The LLM index is rebuilt + THEN: + - Only the root document is in the index + """ + root = DocumentFactory() + DocumentFactory(root_document=root, version_index=1) + + indexing.update_llm_index(rebuild=True) + + with indexing.get_vector_store() as store: + indexed = store.get_modified_times() + + assert set(indexed) == {str(root.pk)} + + def test_incremental_update_by_version_id_refreshes_the_root( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + """ + GIVEN: + - An indexed root document with a version whose root was modified since + WHEN: + - An incremental update is scoped to the version's id + THEN: + - The root's entry is refreshed and no entry exists for the version + """ + root = DocumentFactory() + version = DocumentFactory(root_document=root, version_index=1) + indexing.update_llm_index(rebuild=True) + Document.objects.filter(pk=root.pk).update(modified=timezone.now()) + root.refresh_from_db() + + indexing.update_llm_index(document_ids=[version.pk]) + + with indexing.get_vector_store() as store: + indexed = store.get_modified_times() + + assert indexed == {str(root.pk): root.modified.isoformat()} + + @pytest.mark.django_db def test_add_or_update_document_updates_existing_entry( temp_llm_index_dir: Path, @@ -637,6 +714,7 @@ class TestLlmIndexAddOrUpdateDocumentEmptyContent: doc = MagicMock(spec=Document) doc.id = 42 + doc.root_document_id = None # Must not raise indexing.llm_index_add_or_update_document(doc)