unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/pipeline/pipeline.py

@@ -0,0 +1,284 @@
+import logging
+import multiprocessing as mp
+from dataclasses import InitVar, dataclass, field
+from time import time
+from typing import Any, Optional, Union
+
+from unstructured_ingest.v2.interfaces import ProcessorConfig
+from unstructured_ingest.v2.logger import logger, make_default_logger
+from unstructured_ingest.v2.pipeline.steps.chunk import Chunker, ChunkStep
+from unstructured_ingest.v2.pipeline.steps.download import DownloaderT, DownloadStep
+from unstructured_ingest.v2.pipeline.steps.embed import Embedder, EmbedStep
+from unstructured_ingest.v2.pipeline.steps.index import IndexerT, IndexStep
+from unstructured_ingest.v2.pipeline.steps.partition import Partitioner, PartitionStep
+from unstructured_ingest.v2.pipeline.steps.stage import UploadStager, UploadStageStep
+from unstructured_ingest.v2.pipeline.steps.uncompress import Uncompressor, UncompressStep
+from unstructured_ingest.v2.pipeline.steps.upload import Uploader, UploadStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+from unstructured_ingest.v2.processes.chunker import ChunkerConfig
+from unstructured_ingest.v2.processes.connector_registry import (
+    ConnectionConfig,
+    DownloaderConfigT,
+    IndexerConfigT,
+    UploaderConfigT,
+    UploadStagerConfigT,
+    destination_registry,
+    source_registry,
+)
+from unstructured_ingest.v2.processes.connectors.local import LocalUploader
+from unstructured_ingest.v2.processes.embedder import EmbedderConfig
+from unstructured_ingest.v2.processes.partitioner import PartitionerConfig
+
+
+class PipelineError(Exception):
+    pass
+
+
+@dataclass
+class Pipeline:
+    context: ProcessorConfig
+    indexer: InitVar[IndexerT]
+    indexer_step: IndexStep = field(init=False)
+    downloader: InitVar[DownloaderT]
+    downloader_step: DownloadStep = field(init=False)
+    partitioner: InitVar[Partitioner]
+    partitioner_step: PartitionStep = field(init=False)
+    chunker: InitVar[Optional[Chunker]] = None
+    chunker_step: ChunkStep = field(init=False, default=None)
+    embedder: InitVar[Optional[Embedder]] = None
+    embedder_step: EmbedStep = field(init=False, default=None)
+    stager: InitVar[Optional[UploadStager]] = None
+    stager_step: UploadStageStep = field(init=False, default=None)
+    uploader: InitVar[Uploader] = field(default=LocalUploader())
+    uploader_step: UploadStep = field(init=False, default=None)
+    uncompress_step: UncompressStep = field(init=False, default=None)
+
+    def __post_init__(
+        self,
+        indexer: IndexerT,
+        downloader: DownloaderT,
+        partitioner: Partitioner,
+        chunker: Chunker = None,
+        embedder: Embedder = None,
+        stager: UploadStager = None,
+        uploader: Uploader = None,
+    ):
+        make_default_logger(level=logging.DEBUG if self.context.verbose else logging.INFO)
+        self.indexer_step = IndexStep(process=indexer, context=self.context)
+        self.downloader_step = DownloadStep(process=downloader, context=self.context)
+        self.partitioner_step = PartitionStep(process=partitioner, context=self.context)
+        self.chunker_step = ChunkStep(process=chunker, context=self.context) if chunker else None
+
+        self.embedder_step = EmbedStep(process=embedder, context=self.context) if embedder else None
+        # TODO: support initialize() call from each step process
+        # Potential long call to download embedder models, run before any fanout:
+        if embedder and embedder.config:
+            embedder.config.get_embedder().initialize()
+
+        self.stager_step = UploadStageStep(process=stager, context=self.context) if stager else None
+        self.uploader_step = UploadStep(process=uploader, context=self.context)
+        if self.context.uncompress:
+            process = Uncompressor()
+            self.uncompress_step = UncompressStep(process=process, context=self.context)
+
+        self.check_destination_connector()
+
+    def check_destination_connector(self):
+        # Make sure that if the set destination connector expects a stager, one is also set
+        if not self.uploader_step:
+            return
+        uploader_connector_type = self.uploader_step.process.connector_type
+        registry_entry = destination_registry[uploader_connector_type]
+        if registry_entry.upload_stager and self.stager_step is None:
+            raise ValueError(
+                f"pipeline with uploader type {self.uploader_step.process.__class__.__name__} "
+                f"expects a stager of type {registry_entry.upload_stager.__name__} "
+                f"but one was not set"
+            )
+
+    def cleanup(self):
+        pass
+
+    def log_statuses(self):
+        if status := self.context.status:
+            logger.error(f"{len(status)} failed documents:")
+            for k, v in status.items():
+                for kk, vv in v.items():
+                    logger.error(f"{k}: [{kk}] {vv}")
+
+    def run(self):
+        try:
+            start_time = time()
+            self._run()
+            logger.info(f"Finished ingest process in {time() - start_time}s")
+        finally:
+            self.log_statuses()
+            self.cleanup()
+            if self.context.status:
+                raise PipelineError("Pipeline did not run successfully")
+
+    def clean_results(self, results: Optional[list[Union[Any, list[Any]]]]) -> Optional[list[Any]]:
+        if not results:
+            return None
+        results = [r for r in results if r]
+        flat = []
+        for r in results:
+            if isinstance(r, list):
+                flat.extend(r)
+            else:
+                flat.append(r)
+        final = [f for f in flat if f]
+        return final or None
+
+    def _run(self):
+        logger.info(
+            f"Running local pipline: {self} with configs: "
+            f"{sterilize_dict(self.context.to_dict(redact_sensitive=True))}"
+        )
+        if self.context.mp_supported:
+            manager = mp.Manager()
+            self.context.status = manager.dict()
+        else:
+            self.context.status = {}
+
+        # Index into data source
+        indices = self.indexer_step.run()
+        indices_inputs = [{"file_data_path": i} for i in indices]
+        if not indices_inputs:
+            return
+
+        # Download associated content to local file system
+        downloaded_data = self.downloader_step(indices_inputs)
+        downloaded_data = self.clean_results(results=downloaded_data)
+        if not downloaded_data:
+            return
+
+        # Run uncompress if available
+        if self.uncompress_step:
+            downloaded_data = self.uncompress_step(downloaded_data)
+            # Flatten list of lists
+            downloaded_data = self.clean_results(results=downloaded_data)
+
+        if not downloaded_data:
+            return
+
+        # Partition content
+        elements = self.partitioner_step(downloaded_data)
+        elements = self.clean_results(results=elements)
+        if not elements:
+            return
+
+        # Run element specific modifiers
+        for step in [self.chunker_step, self.embedder_step, self.stager_step]:
+            elements = step(elements) if step else elements
+            elements = self.clean_results(results=elements)
+            if not elements:
+                return
+
+        # Upload the final result
+        self.uploader_step(iterable=elements)
+
+    def __str__(self):
+        s = [str(self.indexer_step), str(self.downloader_step)]
+        if uncompress_step := self.uncompress_step:
+            s.append(str(uncompress_step))
+        s.append(str(self.partitioner_step))
+        if chunker_step := self.chunker_step:
+            s.append(str(chunker_step))
+        if embedder_step := self.embedder_step:
+            s.append(str(embedder_step))
+        if stager_step := self.stager_step:
+            s.append(str(stager_step))
+        s.append(str(self.uploader_step))
+        return " -> ".join(s)
+
+    @classmethod
+    def from_configs(
+        cls,
+        context: ProcessorConfig,
+        indexer_config: IndexerConfigT,
+        downloader_config: DownloaderConfigT,
+        source_connection_config: ConnectionConfig,
+        partitioner_config: PartitionerConfig,
+        chunker_config: Optional[ChunkerConfig] = None,
+        embedder_config: Optional[EmbedderConfig] = None,
+        destination_connection_config: Optional[ConnectionConfig] = None,
+        stager_config: Optional[UploadStagerConfigT] = None,
+        uploader_config: Optional[UploaderConfigT] = None,
+    ) -> "Pipeline":
+        # Get registry key based on indexer config
+        source_entry = {
+            k: v
+            for k, v in source_registry.items()
+            if isinstance(indexer_config, v.indexer_config)
+            and isinstance(downloader_config, v.downloader_config)
+            and isinstance(source_connection_config, v.connection_config)
+        }
+        if len(source_entry) > 1:
+            raise ValueError(
+                f"multiple entries found matching provided indexer, "
+                f"downloader and connection configs: {source_entry}"
+            )
+        if len(source_entry) != 1:
+            raise ValueError(
+                "no entry found in source registry with matching indexer, "
+                "downloader and connection configs"
+            )
+        source = list(source_entry.values())[0]
+        pipeline_kwargs = {
+            "context": context,
+            "indexer": source.indexer(
+                index_config=indexer_config, connection_config=source_connection_config
+            ),
+            "downloader": source.downloader(
+                download_config=downloader_config, connection_config=source_connection_config
+            ),
+            "partitioner": Partitioner(config=partitioner_config),
+        }
+        if chunker_config:
+            pipeline_kwargs["chunker"] = Chunker(config=chunker_config)
+        if embedder_config:
+            pipeline_kwargs["embedder"] = Embedder(config=embedder_config)
+        if not uploader_config:
+            return Pipeline(**pipeline_kwargs)
+
+        destination_entry = {
+            k: v
+            for k, v in destination_registry.items()
+            if isinstance(uploader_config, v.uploader_config)
+        }
+        if destination_connection_config:
+            destination_entry = {
+                k: v
+                for k, v in destination_entry.items()
+                if isinstance(destination_connection_config, v.connection_config)
+            }
+        if stager_config:
+            destination_entry = {
+                k: v
+                for k, v in destination_entry.items()
+                if isinstance(stager_config, v.upload_stager_config)
+            }
+
+        if len(destination_entry) > 1:
+            raise ValueError(
+                f"multiple entries found matching provided uploader, "
+                f"stager and connection configs: {destination_entry}"
+            )
+        if len(destination_entry) != 1:
+            raise ValueError(
+                "no entry found in source registry with matching uploader, "
+                "stager and connection configs"
+            )
+
+        destination = list(destination_entry.values())[0]
+        if stager_config:
+            pipeline_kwargs["stager"] = destination.upload_stager(
+                upload_stager_config=stager_config
+            )
+        if uploader_config:
+            uploader_kwargs = {"upload_config": uploader_config}
+            if destination_connection_config:
+                uploader_kwargs["connection_config"] = destination_connection_config
+            pipeline_kwargs["uploader"] = destination.uploader(**uploader_kwargs)
+        return cls(**pipeline_kwargs)

unstructured_ingest/v2/pipeline/steps/chunk.py

@@ -0,0 +1,85 @@
+import asyncio
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, TypedDict
+
+from unstructured.staging.base import elements_to_dicts
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+from unstructured_ingest.v2.processes.chunker import Chunker
+
+STEP_ID = "chunk"
+
+
+class ChunkStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+@dataclass
+class ChunkStep(PipelineStep):
+    process: Chunker
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.config.chunking_strategy})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.config.to_dict(redact_sensitive=True))
+            if self.process.config
+            else None
+        )
+        logger.info(f"Created {self.identifier} with configs: {config}")
+
+    def should_chunk(self, filepath: Path, file_data: FileData) -> bool:
+        if self.context.reprocess or file_data.reprocess:
+            return True
+        return not filepath.exists()
+
+    def get_output_filepath(self, filename: Path) -> Path:
+        hashed_output_file = f"{self.get_hash(extras=[filename.name])}.json"
+        filepath = (self.cache_dir / hashed_output_file).resolve()
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        return filepath
+
+    def _save_output(self, output_filepath: str, chunked_content: list[dict]):
+        with open(str(output_filepath), "w") as f:
+            logger.debug(f"Writing chunker output to: {output_filepath}")
+            json.dump(chunked_content, f, indent=2)
+
+    async def _run_async(
+        self, fn: Callable, path: str, file_data_path: str, **kwargs
+    ) -> ChunkStepResponse:
+        path = Path(path)
+        file_data = FileData.from_file(path=file_data_path)
+        output_filepath = self.get_output_filepath(filename=path)
+        if not self.should_chunk(filepath=output_filepath, file_data=file_data):
+            logger.debug(f"Skipping chunking, output already exists: {output_filepath}")
+            return ChunkStepResponse(file_data_path=file_data_path, path=str(output_filepath))
+        fn_kwargs = {"elements_filepath": path}
+        if not asyncio.iscoroutinefunction(fn):
+            chunked_content_raw = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                chunked_content_raw = await fn(**fn_kwargs)
+        else:
+            chunked_content_raw = await fn(**fn_kwargs)
+        self._save_output(
+            output_filepath=str(output_filepath),
+            chunked_content=elements_to_dicts(chunked_content_raw),
+        )
+        return ChunkStepResponse(file_data_path=file_data_path, path=str(output_filepath))
+
+    def get_hash(self, extras: Optional[list[str]]) -> str:
+        hashable_string = json.dumps(
+            self.process.config.to_dict(), sort_keys=True, ensure_ascii=True
+        )
+        if extras:
+            hashable_string += "".join(extras)
+        return hashlib.sha256(hashable_string.encode()).hexdigest()[:12]

unstructured_ingest/v2/pipeline/steps/download.py

@@ -0,0 +1,124 @@
+import asyncio
+import hashlib
+import json
+from dataclasses import dataclass
+from typing import Callable, Optional, TypedDict, TypeVar
+
+from unstructured_ingest.v2.interfaces import FileData, download_responses
+from unstructured_ingest.v2.interfaces.downloader import Downloader
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+
+DownloaderT = TypeVar("DownloaderT", bound=Downloader)
+
+STEP_ID = "download"
+
+
+class DownloadStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+@dataclass
+class DownloadStep(PipelineStep):
+    process: DownloaderT
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.__class__.__name__})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.download_config.to_dict(redact_sensitive=True))
+            if self.process.download_config
+            else None
+        )
+        connection_config = (
+            sterilize_dict(self.process.connection_config.to_dict(redact_sensitive=True))
+            if self.process.connection_config
+            else None
+        )
+        logger.info(
+            f"Created {self.identifier} with configs: {config}, "
+            f"connection configs: {connection_config}"
+        )
+
+    @staticmethod
+    def is_float(value: str):
+        try:
+            float(value)
+            return True
+        except ValueError:
+            return False
+
+    def should_download(self, file_data: FileData, file_data_path: str) -> bool:
+        if self.context.re_download:
+            return True
+        download_path = self.process.get_download_path(file_data=file_data)
+        if not download_path or not download_path.exists():
+            return True
+        if (
+            download_path.is_file()
+            and file_data.metadata.date_modified
+            and self.is_float(file_data.metadata.date_modified)
+            and download_path.stat().st_mtime > float(file_data.metadata.date_modified)
+        ):
+            # Also update file data to mark this to reprocess since this won't change the filename
+            file_data.reprocess = True
+            file_data.to_file(path=file_data_path)
+            return True
+        return False
+
+    async def _run_async(self, fn: Callable, file_data_path: str) -> list[DownloadStepResponse]:
+        file_data = FileData.from_file(path=file_data_path)
+        download_path = self.process.get_download_path(file_data=file_data)
+        if not self.should_download(file_data=file_data, file_data_path=file_data_path):
+            logger.debug(f"Skipping download, file already exists locally: {download_path}")
+            return [DownloadStepResponse(file_data_path=file_data_path, path=str(download_path))]
+        fn_kwargs = {"file_data": file_data}
+        if not asyncio.iscoroutinefunction(fn):
+            download_results = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                download_results = await fn(**fn_kwargs)
+        else:
+            download_results = await fn(**fn_kwargs)
+        return self.create_step_results(
+            current_file_data_path=file_data_path, download_results=download_results
+        )
+
+    def create_step_results(
+        self, current_file_data_path: str, download_results: download_responses
+    ) -> list[DownloadStepResponse]:
+        if not isinstance(download_results, list):
+            return [
+                DownloadStepResponse(
+                    file_data_path=current_file_data_path, path=str(download_results["path"])
+                )
+            ]
+            # Supplemental results generated as part of the download process
+        download_step_results = []
+        for res in download_results:
+            file_data_path = self.persist_new_file_data(file_data=res["file_data"])
+            download_step_results.append(
+                DownloadStepResponse(file_data_path=file_data_path, path=res["path"])
+            )
+        return download_step_results
+
+    def persist_new_file_data(self, file_data: FileData) -> str:
+        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)
+        return str(filepath)
+
+    def get_hash(self, extras: Optional[list[str]]) -> str:
+        hashable_string = json.dumps(
+            sterilize_dict(self.process.download_config.to_dict()), sort_keys=True
+        )
+        if extras:
+            hashable_string += "".join(extras)
+        return hashlib.sha256(hashable_string.encode()).hexdigest()[:12]

unstructured_ingest/v2/pipeline/steps/embed.py

@@ -0,0 +1,84 @@
+import asyncio
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, TypedDict
+
+from unstructured.staging.base import elements_to_dicts
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+from unstructured_ingest.v2.processes.embedder import Embedder
+
+STEP_ID = "embed"
+
+
+class EmbedStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+@dataclass
+class EmbedStep(PipelineStep):
+    process: Embedder
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.config.embedding_provider})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.config.to_dict(redact_sensitive=True))
+            if self.process.config
+            else None
+        )
+        logger.info(f"Created {self.identifier} with configs: {config}")
+
+    def should_embed(self, filepath: Path, file_data: FileData) -> bool:
+        if self.context.reprocess or file_data.reprocess:
+            return True
+        return not filepath.exists()
+
+    def get_output_filepath(self, filename: Path) -> Path:
+        hashed_output_file = f"{self.get_hash(extras=[filename.name])}.json"
+        filepath = (self.cache_dir / hashed_output_file).resolve()
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        return filepath
+
+    def _save_output(self, output_filepath: str, embedded_content: list[dict]):
+        with open(str(output_filepath), "w") as f:
+            logger.debug(f"Writing embedded output to: {output_filepath}")
+            json.dump(embedded_content, f, indent=2)
+
+    async def _run_async(self, fn: Callable, path: str, file_data_path: str) -> EmbedStepResponse:
+        path = Path(path)
+        file_data = FileData.from_file(path=file_data_path)
+        output_filepath = self.get_output_filepath(filename=path)
+        if not self.should_embed(filepath=output_filepath, file_data=file_data):
+            logger.debug(f"Skipping embedding, output already exists: {output_filepath}")
+            return EmbedStepResponse(file_data_path=file_data_path, path=str(output_filepath))
+        fn_kwargs = {"elements_filepath": path}
+        if not asyncio.iscoroutinefunction(fn):
+            embed_content_raw = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                embed_content_raw = await fn(**fn_kwargs)
+        else:
+            embed_content_raw = await fn(**fn_kwargs)
+
+        self._save_output(
+            output_filepath=str(output_filepath),
+            embedded_content=elements_to_dicts(embed_content_raw),
+        )
+        return EmbedStepResponse(file_data_path=file_data_path, path=str(output_filepath))
+
+    def get_hash(self, extras: Optional[list[str]]) -> str:
+        hashable_string = json.dumps(
+            self.process.config.to_dict(), sort_keys=True, ensure_ascii=True
+        )
+        if extras:
+            hashable_string += "".join(extras)
+        return hashlib.sha256(hashable_string.encode()).hexdigest()[:12]

unstructured_ingest/v2/pipeline/steps/index.py

@@ -0,0 +1,61 @@
+import hashlib
+import json
+from dataclasses import dataclass
+from typing import Generator, Optional, TypeVar
+
+from unstructured_ingest.v2.interfaces.indexer import Indexer
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+
+IndexerT = TypeVar("IndexerT", bound=Indexer)
+
+STEP_ID = "index"
+
+
+@dataclass
+class IndexStep(PipelineStep):
+    process: IndexerT
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.__class__.__name__})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.index_config.to_dict(redact_sensitive=True))
+            if self.process.index_config
+            else None
+        )
+        connection_config = (
+            sterilize_dict(self.process.connection_config.to_dict(redact_sensitive=True))
+            if self.process.connection_config
+            else None
+        )
+        logger.info(
+            f"Created {self.identifier} with configs: {config}, "
+            f"connection configs: {connection_config}"
+        )
+
+    def run(self) -> Generator[str, None, None]:
+        for file_data in self.process.run():
+            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:
+        hashable_string = json.dumps(self.process.index_config.to_dict())
+        if extras:
+            hashable_string += "".join(extras)
+        return hashlib.sha256(hashable_string.encode()).hexdigest()[:12]