Coverage for documents/management/commands/document_llmindex.py: 0%

1from typing import Any 

2 

3from documents.management.commands.base import PaperlessCommand 

4from documents.tasks import llmindex_index 

5from paperless_ai.indexing import llm_index_compact 

6from paperless_ai.indexing import llm_index_migrate 

7 

8 

9class Command(PaperlessCommand): 

10 help = "Manages the LLM-based vector index for Paperless." 

11 

12 supports_progress_bar = True 

13 supports_multiprocessing = False 

14 

15 def add_arguments(self, parser: Any) -> None: 

16 super().add_arguments(parser) 

17 parser.add_argument( 

18 "command", 

19 choices=["rebuild", "update", "compact", "migrate"], 

20 ) 

21 

22 def handle(self, *args: Any, **options: Any) -> None: 

23 if options["command"] == "compact": 

24 llm_index_compact() 

25 return 

26 if options["command"] == "migrate": 

27 llm_index_migrate() 

28 return 

29 llmindex_index( 

30 rebuild=options["command"] == "rebuild", 

31 iter_wrapper=lambda docs: self.track( 

32 docs, 

33 description="Indexing documents...", 

34 ), 

35 )