Make the file strategy parallel with a default of 10 files

Author: tonybaloney

app/backend/prepdocs.py

@@ -311,6 +311,12 @@ async def main(strategy: Strategy, setup_index: bool = True):
         required=False,
         help="Search service system assigned Identity (Managed identity) (used for integrated vectorization)",
     )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=10,
+        help="Max. number of concurrent tasks to run for processing files (file strategy only) (default: 10)",
+    )
 
     parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
     args = parser.parse_args()
@@ -467,6 +473,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
             category=args.category,
             use_content_understanding=use_content_understanding,
             content_understanding_endpoint=os.getenv("AZURE_CONTENTUNDERSTANDING_ENDPOINT"),
+            concurrency=args.concurrency,
         )
 
     loop.run_until_complete(main(ingestion_strategy, setup_index=not args.remove and not args.removeall))

app/backend/prepdocslib/filestrategy.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import Optional
 
@@ -56,6 +57,7 @@ def __init__(
         category: Optional[str] = None,
         use_content_understanding: bool = False,
         content_understanding_endpoint: Optional[str] = None,
+        concurrency: int = 10,
     ):
         self.list_file_strategy = list_file_strategy
         self.blob_manager = blob_manager
@@ -70,6 +72,7 @@ def __init__(
         self.category = category
         self.use_content_understanding = use_content_understanding
         self.content_understanding_endpoint = content_understanding_endpoint
+        self.concurrency = concurrency
 
     def setup_search_manager(self):
         self.search_manager = SearchManager(
@@ -98,9 +101,9 @@ async def setup(self):
 
     async def run(self):
         self.setup_search_manager()
-        if self.document_action == DocumentAction.Add:
-            files = self.list_file_strategy.list()
-            async for file in files:
+
+        async def process_file_worker(semaphore: asyncio.Semaphore, file: File):
+            async with semaphore:
                 try:
                     sections = await parse_file(file, self.file_processors, self.category, self.image_embeddings)
                     if sections:
@@ -112,6 +115,12 @@ async def run(self):
                 finally:
                     if file:
                         file.close()
+
+        if self.document_action == DocumentAction.Add:
+            files = self.list_file_strategy.list()
+            semaphore = asyncio.Semaphore(self.concurrency)
+            tasks = [process_file_worker(semaphore, file) async for file in files]
+            await asyncio.gather(*tasks)
         elif self.document_action == DocumentAction.Remove:
             paths = self.list_file_strategy.list_paths()
             async for path in paths: