bugfix/support async indexing (#192)

* Support async indexing in pipeling

* bump changelog

* fix asyncgenerator typing

Author: rbiseck3

CHANGELOG.md

@@ -1,4 +1,4 @@
-## 0.1.1-dev3
+## 0.1.1-dev4
 
 ### Enhancements
 
@@ -9,6 +9,9 @@
 * **Migrate Slack Source Connector to V2**
 * **Migrate Slack Source Connector to V2**
 * **Add Delta Table destination to v2**
+* **Migrate Slack Source Connector to V2**
+
+>>>>>>> 9214214 (bump changelog)
 
 ## 0.1.0
 

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.1-dev3"  # pragma: no cover
+__version__ = "0.1.1-dev4"  # pragma: no cover

unstructured_ingest/v2/interfaces/indexer.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Generator, Optional, TypeVar
+from typing import Any, AsyncGenerator, Generator, Optional, TypeVar
 
 from pydantic import BaseModel
 
@@ -25,3 +25,6 @@ def is_async(self) -> bool:
     @abstractmethod
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
         pass
+
+    async def run_async(self, **kwargs: Any) -> AsyncGenerator[FileData, None]:
+        raise NotImplementedError()

unstructured_ingest/v2/pipeline/pipeline.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import logging
 import multiprocessing as mp
 import shutil
@@ -186,6 +187,14 @@ def apply_filter(self, records: list[dict]) -> list[dict]:
         filtered_records = [r for r in records if r["file_data_path"] in filtered_file_data_paths]
         return filtered_records
 
+    def get_indices(self) -> list[dict]:
+        if self.indexer_step.process.is_async():
+            indices = asyncio.run(self.indexer_step.run_async())
+        else:
+            indices = self.indexer_step.run()
+        indices_inputs = [{"file_data_path": i} for i in indices]
+        return indices_inputs
+
     def _run(self):
         logger.info(
             f"running local pipeline: {self} with configs: " f"{self.context.model_dump_json()}"
@@ -197,8 +206,7 @@ def _run(self):
             self.context.status = {}
 
         # Index into data source
-        indices = self.indexer_step.run()
-        indices_inputs = [{"file_data_path": i} for i in indices]
+        indices_inputs = self.get_indices()
         if not indices_inputs:
             logger.info("No files to process after indexer, exiting")
             return

unstructured_ingest/v2/pipeline/steps/index.py

@@ -1,7 +1,7 @@
 import hashlib
 import json
 from dataclasses import dataclass
-from typing import Generator, Optional, TypeVar
+from typing import AsyncGenerator, Generator, Optional, TypeVar
 
 from unstructured_ingest.v2.interfaces.indexer import Indexer
 from unstructured_ingest.v2.logger import logger
@@ -52,6 +52,23 @@ def run(self) -> Generator[str, None, None]:
                     raise e
                 continue
 
+    async def run_async(self) -> AsyncGenerator[str, None]:
+        async for file_data in self.process.run_async():
+            logger.debug(f"generated file data: {file_data.to_dict()}")
+            try:
+                record_hash = self.get_hash(extras=[file_data.identifier])
+                filename = f"{record_hash}.json"
+                filepath = (self.cache_dir / filename).resolve()
+                filepath.parent.mkdir(parents=True, exist_ok=True)
+                with open(str(filepath), "w") as f:
+                    json.dump(file_data.to_dict(), f, indent=2)
+                yield str(filepath)
+            except Exception as e:
+                logger.error(f"failed to create index for file data: {file_data}", exc_info=True)
+                if self.context.raise_on_error:
+                    raise e
+                continue
+
     def get_hash(self, extras: Optional[list[str]]) -> str:
         index_config_dict = json.loads(
             serialize_base_model_json(model=self.process.index_config, sort_keys=True)