Feature: Switch all indexing to use rich (#12193)

src/documents/management/commands/document_index.py

@@ -1,22 +1,25 @@
-from django.core.management import BaseCommand
 from django.db import transaction
 
-from documents.management.commands.mixins import ProgressBarMixin
+from documents.management.commands.base import PaperlessCommand
 from documents.tasks import index_optimize
 from documents.tasks import index_reindex
 
 
-class Command(ProgressBarMixin, BaseCommand):
+class Command(PaperlessCommand):
     help = "Manages the document index."
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument("command", choices=["reindex", "optimize"])
-        self.add_argument_progress_bar_mixin(parser)
 
     def handle(self, *args, **options):
-        self.handle_progress_bar_mixin(**options)
         with transaction.atomic():
             if options["command"] == "reindex":
-                index_reindex(progress_bar_disable=self.no_progress_bar)
+                index_reindex(
+                    iter_wrapper=lambda docs: self.track(
+                        docs,
+                        description="Indexing documents...",
+                    ),
+                )
             elif options["command"] == "optimize":
                 index_optimize()

src/documents/management/commands/document_llmindex.py

@@ -1,22 +1,22 @@
-from django.core.management import BaseCommand
-from django.db import transaction
+from typing import Any
 
-from documents.management.commands.mixins import ProgressBarMixin
+from documents.management.commands.base import PaperlessCommand
 from documents.tasks import llmindex_index
 
 
-class Command(ProgressBarMixin, BaseCommand):
+class Command(PaperlessCommand):
     help = "Manages the LLM-based vector index for Paperless."
 
-    def add_arguments(self, parser):
+    def add_arguments(self, parser: Any) -> None:
+        super().add_arguments(parser)
         parser.add_argument("command", choices=["rebuild", "update"])
-        self.add_argument_progress_bar_mixin(parser)
 
-    def handle(self, *args, **options):
-        self.handle_progress_bar_mixin(**options)
-        with transaction.atomic():
-            llmindex_index(
-                progress_bar_disable=self.no_progress_bar,
-                rebuild=options["command"] == "rebuild",
-                scheduled=False,
-            )
+    def handle(self, *args: Any, **options: Any) -> None:
+        llmindex_index(
+            rebuild=options["command"] == "rebuild",
+            scheduled=False,
+            iter_wrapper=lambda docs: self.track(
+                docs,
+                description="Indexing documents...",
+            ),
+        )

src/documents/tasks.py

@@ -4,11 +4,13 @@ import logging
 import shutil
 import uuid
 import zipfile
+from collections.abc import Callable
+from collections.abc import Iterable
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from tempfile import mkstemp
+from typing import TypeVar
 
-import tqdm
 from celery import Task
 from celery import shared_task
 from celery import states
@@ -66,11 +68,19 @@ from paperless_ai.indexing import llm_index_add_or_update_document
 from paperless_ai.indexing import llm_index_remove_document
 from paperless_ai.indexing import update_llm_index
 
+_T = TypeVar("_T")
+IterWrapper = Callable[[Iterable[_T]], Iterable[_T]]
+
+
 if settings.AUDIT_LOG_ENABLED:
     from auditlog.models import LogEntry
 logger = logging.getLogger("paperless.tasks")
 
 
+def _identity(iterable: Iterable[_T]) -> Iterable[_T]:
+    return iterable
+
+
 @shared_task
 def index_optimize() -> None:
     ix = index.open_index()
@@ -78,13 +88,13 @@ def index_optimize() -> None:
     writer.commit(optimize=True)
 
 
-def index_reindex(*, progress_bar_disable=False) -> None:
+def index_reindex(*, iter_wrapper: IterWrapper[Document] = _identity) -> None:
     documents = Document.objects.all()
 
     ix = index.open_index(recreate=True)
 
     with AsyncWriter(ix) as writer:
-        for document in tqdm.tqdm(documents, disable=progress_bar_disable):
+        for document in iter_wrapper(documents):
             index.update_document(writer, document)
 
 
@@ -265,7 +275,6 @@ def bulk_update_documents(document_ids) -> None:
     ai_config = AIConfig()
     if ai_config.llm_index_enabled:
         update_llm_index(
-            progress_bar_disable=True,
             rebuild=False,
         )
 
@@ -606,7 +615,7 @@ def update_document_parent_tags(tag: Tag, new_parent: Tag) -> None:
 @shared_task
 def llmindex_index(
     *,
-    progress_bar_disable=True,
+    iter_wrapper: IterWrapper[Document] = _identity,
     rebuild=False,
     scheduled=True,
     auto=False,
@@ -629,7 +638,7 @@ def llmindex_index(
 
         try:
             result = update_llm_index(
-                progress_bar_disable=progress_bar_disable,
+                iter_wrapper=iter_wrapper,
                 rebuild=rebuild,
             )
             task.status = states.SUCCESS