unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/v2/pipeline/steps/partition.py

@@ -0,0 +1,78 @@
+import asyncio
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, TypedDict
+
+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.partitioner import Partitioner
+
+STEP_ID = "partition"
+
+
+class PartitionStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+@dataclass
+class PartitionStep(PipelineStep):
+    process: Partitioner
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.config.strategy})"
+
+    def __post_init__(self):
+        config = sterilize_dict(self.process.config.to_dict(redact_sensitive=True))
+        logger.info(f"Created {self.identifier} with configs: {config}")
+
+    def should_partition(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, partitioned_content: list[dict]):
+        with open(str(output_filepath), "w") as f:
+            logger.debug(f"Writing partitioned output to: {output_filepath}")
+            json.dump(partitioned_content, f, indent=2)
+
+    async def _run_async(
+        self, fn: Callable, path: str, file_data_path: str
+    ) -> Optional[PartitionStepResponse]:
+        path = Path(path)
+        file_data = FileData.from_file(path=file_data_path)
+        output_filepath = self.get_output_filepath(filename=Path(file_data_path))
+        if not self.should_partition(filepath=output_filepath, file_data=file_data):
+            logger.debug(f"Skipping partitioning, output already exists: {output_filepath}")
+            return PartitionStepResponse(file_data_path=file_data_path, path=str(output_filepath))
+        fn_kwargs = {"filename": path, "metadata": file_data.metadata}
+        if not asyncio.iscoroutinefunction(fn):
+            partitioned_content = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                partitioned_content = await fn(**fn_kwargs)
+        else:
+            partitioned_content = await fn(**fn_kwargs)
+        self._save_output(
+            output_filepath=str(output_filepath), partitioned_content=partitioned_content
+        )
+        return PartitionStepResponse(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/stage.py

@@ -0,0 +1,64 @@
+import asyncio
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, TypedDict
+
+from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.upload_stager import UploadStager
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+
+STEP_ID = "upload_stage"
+
+
+class UploadStageStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+@dataclass
+class UploadStageStep(PipelineStep):
+    process: UploadStager
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.__class__.__name__})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.upload_stager_config.to_dict(redact_sensitive=True))
+            if self.process.upload_stager_config
+            else None
+        )
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Created {self.identifier} with configs: {config}")
+
+    async def _run_async(
+        self, fn: Callable, path: str, file_data_path: str
+    ) -> UploadStageStepResponse:
+        path = Path(path)
+        fn_kwargs = {
+            "elements_filepath": path,
+            "file_data": FileData.from_file(path=file_data_path),
+            "output_dir": self.cache_dir,
+            "output_filename": self.get_hash(extras=[path.name]),
+        }
+        if not asyncio.iscoroutinefunction(fn):
+            staged_output_path = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                staged_output_path = await fn(**fn_kwargs)
+        else:
+            staged_output_path = await fn(**fn_kwargs)
+        return UploadStageStepResponse(file_data_path=file_data_path, path=str(staged_output_path))
+
+    def get_hash(self, extras: Optional[list[str]]) -> str:
+        hashable_string = json.dumps(
+            self.process.upload_stager_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/uncompress.py

@@ -0,0 +1,68 @@
+import asyncio
+from pathlib import Path
+from typing import Callable, TypedDict
+
+from unstructured_ingest.v2.interfaces.file_data 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.uncompress import Uncompressor
+
+STEP_ID = "uncompress"
+
+
+class UncompressStepResponse(TypedDict):
+    file_data_path: str
+    path: str
+
+
+class UncompressStep(PipelineStep):
+    process: Uncompressor
+    identifier: str = STEP_ID
+
+    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 _run(self, path: str, file_data_path: str) -> list[UncompressStepResponse]:
+        file_data = FileData.from_file(path=file_data_path)
+        new_file_data = self.process.run(file_data=file_data)
+        responses = []
+        for new_file in new_file_data:
+            new_file_data_path = Path(file_data_path).parent / f"{new_file.identifier}.json"
+            new_file.to_file(path=str(new_file_data_path.resolve()))
+            responses.append(
+                UncompressStepResponse(
+                    path=new_file.source_identifiers.fullpath,
+                    file_data_path=str(new_file_data_path),
+                )
+            )
+        return responses
+
+    async def _run_async(
+        self, fn: Callable, path: str, file_data_path: str
+    ) -> list[UncompressStepResponse]:
+        file_data = FileData.from_file(path=file_data_path)
+        fn_kwargs = {"file_data": file_data}
+        if not asyncio.iscoroutinefunction(fn):
+            new_file_data = fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                new_file_data = await fn(**fn_kwargs)
+        else:
+            new_file_data = await fn(**fn_kwargs)
+        responses = []
+        for new_file in new_file_data:
+            new_file_data_path = Path(file_data_path).parent / f"{new_file.identifier}.json"
+            new_file.to_file(path=str(new_file_data_path.resolve()))
+            responses.append(
+                UncompressStepResponse(
+                    path=new_file.source_identifiers.fullpath,
+                    file_data_path=str(new_file_data_path),
+                )
+            )
+        return responses

unstructured_ingest/v2/pipeline/steps/upload.py

@@ -0,0 +1,73 @@
+import asyncio
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, TypedDict
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.uploader import UploadContent, Uploader
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.pipeline.interfaces import PipelineStep, iterable_input, timed
+from unstructured_ingest.v2.pipeline.utils import sterilize_dict
+
+STEP_ID = "upload"
+
+
+class UploadStepContent(TypedDict):
+    path: str
+    file_data_path: str
+
+
+@dataclass
+class UploadStep(PipelineStep):
+    process: Uploader
+    identifier: str = STEP_ID
+
+    def __str__(self):
+        return f"{self.identifier} ({self.process.__class__.__name__})"
+
+    def __post_init__(self):
+        config = (
+            sterilize_dict(self.process.upload_config.to_dict(redact_sensitive=True))
+            if self.process.upload_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 process_whole(self, iterable: iterable_input):
+        self.run(contents=iterable)
+
+    @timed
+    def __call__(self, iterable: iterable_input):
+        logger.info(
+            f"Calling {self.__class__.__name__} " f"with {len(iterable)} docs",  # type: ignore
+        )
+        if self.process.is_async():
+            self.process_async(iterable=iterable)
+        else:
+            self.process_whole(iterable=iterable)
+
+    def _run(self, fn: Callable, contents: list[UploadStepContent]):
+        upload_contents = [
+            UploadContent(path=Path(c["path"]), file_data=FileData.from_file(c["file_data_path"]))
+            for c in contents
+        ]
+        fn(contents=upload_contents)
+
+    async def _run_async(self, path: str, file_data_path: str, fn: Optional[Callable] = None):
+        fn = fn or self.process.run_async
+        fn_kwargs = {"path": Path(path), "file_data": FileData.from_file(path=file_data_path)}
+        if not asyncio.iscoroutinefunction(fn):
+            fn(**fn_kwargs)
+        elif semaphore := self.context.semaphore:
+            async with semaphore:
+                await fn(**fn_kwargs)
+        else:
+            await fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/utils.py

@@ -0,0 +1,15 @@
+import json
+from datetime import datetime
+from pathlib import Path
+
+
+def sterilize_dict(data: dict) -> dict:
+    def json_serial(obj):
+        if isinstance(obj, Path):
+            return obj.as_posix()
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+        raise TypeError("Type %s not serializable" % type(obj))
+
+    data_s = json.dumps(data, default=json_serial)
+    return json.loads(data_s)

unstructured_ingest/v2/processes/chunker.py

@@ -0,0 +1,97 @@
+from abc import ABC
+from dataclasses import dataclass, fields
+from pathlib import Path
+from typing import Any, Optional
+
+from unstructured.chunking import dispatch
+from unstructured.documents.elements import Element, assign_and_map_hash_ids
+from unstructured.staging.base import dict_to_elements, elements_from_json
+
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
+from unstructured_ingest.v2.interfaces.process import BaseProcess
+from unstructured_ingest.v2.logger import logger
+
+
+@dataclass
+class ChunkerConfig(EnhancedDataClassJsonMixin):
+    chunking_strategy: Optional[str] = None
+    chunking_endpoint: Optional[str] = "https://api.unstructured.io/general/v0/general"
+    chunk_by_api: bool = False
+    chunk_api_key: Optional[str] = enhanced_field(default=None, sensitive=True)
+
+    chunk_combine_text_under_n_chars: Optional[int] = None
+    chunk_include_orig_elements: Optional[bool] = None
+    chunk_max_characters: Optional[int] = None
+    chunk_multipage_sections: Optional[bool] = None
+    chunk_new_after_n_chars: Optional[int] = None
+    chunk_overlap: Optional[int] = None
+    chunk_overlap_all: Optional[bool] = None
+
+    def to_chunking_kwargs(self) -> dict[str, Any]:
+        return {
+            "chunking_strategy": self.chunking_strategy,
+            "combine_under_n_chars": self.chunk_combine_text_under_n_chars,
+            "max_characters": self.chunk_max_characters,
+            "include_orig_elements": self.chunk_include_orig_elements,
+            "multipage_sections": self.chunk_multipage_sections,
+            "new_after_n_chars": self.chunk_new_after_n_chars,
+            "overlap": self.chunk_overlap,
+            "overlap_all": self.chunk_overlap_all,
+        }
+
+
+@dataclass
+class Chunker(BaseProcess, ABC):
+    config: ChunkerConfig
+
+    def is_async(self) -> bool:
+        return self.config.chunk_by_api
+
+    def run(self, elements_filepath: Path, **kwargs: Any) -> list[Element]:
+        elements = elements_from_json(filename=str(elements_filepath))
+        if not elements:
+            return elements
+        local_chunking_strategies = ("basic", "by_title")
+        if self.config.chunking_strategy not in local_chunking_strategies:
+            logger.warning(
+                "chunking strategy not supported for local chunking: {}, must be one of: {}".format(
+                    self.config.chunking_strategy, ", ".join(local_chunking_strategies)
+                )
+            )
+            return elements
+        chunked_elements = dispatch.chunk(elements=elements, **self.config.to_chunking_kwargs())
+        assign_and_map_hash_ids(chunked_elements)
+        return chunked_elements
+
+    async def run_async(self, elements_filepath: Path, **kwargs: Any) -> list[Element]:
+        from unstructured_client import UnstructuredClient
+        from unstructured_client.models.shared import Files, PartitionParameters
+
+        client = UnstructuredClient(
+            api_key_auth=self.config.chunk_api_key,
+            server_url=self.config.chunking_endpoint,
+        )
+        partition_request = self.config.to_chunking_kwargs()
+        possible_fields = [f.name for f in fields(PartitionParameters)]
+        filtered_partition_request = {
+            k: v for k, v in partition_request.items() if k in possible_fields
+        }
+        if len(filtered_partition_request) != len(partition_request):
+            logger.debug(
+                "Following fields were omitted due to not being "
+                "supported by the currently used unstructured client: {}".format(
+                    ", ".join([v for v in partition_request if v not in filtered_partition_request])
+                )
+            )
+        with open(elements_filepath, "rb") as f:
+            files = Files(
+                content=f.read(),
+                file_name=str(elements_filepath.resolve()),
+            )
+            filtered_partition_request["files"] = files
+        partition_params = PartitionParameters(**filtered_partition_request)
+        resp = client.general.partition(partition_params)
+        elements_raw = resp.elements or []
+        elements = dict_to_elements(elements_raw)
+        assign_and_map_hash_ids(elements)
+        return elements

unstructured_ingest/v2/processes/connector_registry.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+from typing import Optional, Type, TypeVar
+
+from unstructured_ingest.v2.interfaces import (
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    Indexer,
+    IndexerConfig,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+
+IndexerT = TypeVar("IndexerT", bound=Indexer)
+IndexerConfigT = TypeVar("IndexerConfigT", bound=IndexerConfig)
+DownloaderT = TypeVar("DownloaderT", bound=Downloader)
+DownloaderConfigT = TypeVar("DownloaderConfigT", bound=DownloaderConfig)
+ConnectionConfigT = TypeVar("ConnectionConfigT", bound=ConnectionConfig)
+UploadStagerConfigT = TypeVar("UploadStagerConfigT", bound=UploadStagerConfig)
+UploadStagerT = TypeVar("UploadStagerT", bound=UploadStager)
+UploaderConfigT = TypeVar("UploaderConfigT", bound=UploaderConfig)
+UploaderT = TypeVar("UploaderT", bound=Uploader)
+
+
+@dataclass
+class SourceRegistryEntry:
+    indexer: Type[IndexerT]
+    downloader: Type[DownloaderT]
+
+    downloader_config: Optional[Type[DownloaderConfigT]] = None
+    indexer_config: Optional[Type[IndexerConfigT]] = None
+    connection_config: Optional[Type[ConnectionConfigT]] = None
+
+
+source_registry: dict[str, SourceRegistryEntry] = {}
+
+
+def add_source_entry(source_type: str, entry: SourceRegistryEntry):
+    if source_type in source_registry:
+        raise ValueError(f"source {source_type} has already been registered")
+    source_registry[source_type] = entry
+
+
+@dataclass
+class DestinationRegistryEntry:
+    uploader: Type[UploaderT]
+    upload_stager: Optional[Type[UploadStagerT]] = None
+
+    upload_stager_config: Optional[Type[UploadStagerConfigT]] = None
+    uploader_config: Optional[Type[UploaderConfigT]] = None
+
+    connection_config: Optional[Type[ConnectionConfigT]] = None
+
+
+destination_registry: dict[str, DestinationRegistryEntry] = {}
+
+
+def add_destination_entry(destination_type: str, entry: DestinationRegistryEntry):
+    if destination_type in destination_registry:
+        raise ValueError(f"destination {destination_type} has already been registered")
+    destination_registry[destination_type] = entry

unstructured_ingest/v2/processes/connectors/__init__.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import unstructured.ingest.v2.processes.connectors.fsspec  # noqa: F401
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+    add_source_entry,
+)
+
+from .astra import CONNECTOR_TYPE as ASTRA_CONNECTOR_TYPE
+from .astra import astra_destination_entry
+from .chroma import CONNECTOR_TYPE as CHROMA_CONNECTOR_TYPE
+from .chroma import chroma_destination_entry
+from .databricks_volumes import CONNECTOR_TYPE as DATABRICKS_VOLUMES_CONNECTOR_TYPE
+from .databricks_volumes import databricks_volumes_destination_entry
+from .elasticsearch import CONNECTOR_TYPE as ELASTICSEARCH_CONNECTOR_TYPE
+from .elasticsearch import elasticsearch_destination_entry, elasticsearch_source_entry
+from .google_drive import CONNECTOR_TYPE as GOOGLE_DRIVE_CONNECTOR_TYPE
+from .google_drive import google_drive_source_entry
+from .local import CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE
+from .local import local_destination_entry, local_source_entry
+from .mongodb import CONNECTOR_TYPE as MONGODB_CONNECTOR_TYPE
+from .mongodb import mongodb_destination_entry
+from .onedrive import CONNECTOR_TYPE as ONEDRIVE_CONNECTOR_TYPE
+from .onedrive import onedrive_source_entry
+from .opensearch import CONNECTOR_TYPE as OPENSEARCH_CONNECTOR_TYPE
+from .opensearch import opensearch_destination_entry, opensearch_source_entry
+from .pinecone import CONNECTOR_TYPE as PINECONE_CONNECTOR_TYPE
+from .pinecone import pinecone_destination_entry
+from .salesforce import CONNECTOR_TYPE as SALESFORCE_CONNECTOR_TYPE
+from .salesforce import salesforce_source_entry
+from .sharepoint import CONNECTOR_TYPE as SHAREPOINT_CONNECTOR_TYPE
+from .sharepoint import sharepoint_source_entry
+from .singlestore import CONNECTOR_TYPE as SINGLESTORE_CONNECTOR_TYPE
+from .singlestore import singlestore_destination_entry
+from .sql import CONNECTOR_TYPE as SQL_CONNECTOR_TYPE
+from .sql import sql_destination_entry
+from .weaviate import CONNECTOR_TYPE as WEAVIATE_CONNECTOR_TYPE
+from .weaviate import weaviate_destination_entry
+
+add_destination_entry(destination_type=ASTRA_CONNECTOR_TYPE, entry=astra_destination_entry)
+
+add_destination_entry(destination_type=CHROMA_CONNECTOR_TYPE, entry=chroma_destination_entry)
+
+add_source_entry(source_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_source_entry)
+add_destination_entry(
+    destination_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_destination_entry
+)
+
+add_source_entry(source_type=GOOGLE_DRIVE_CONNECTOR_TYPE, entry=google_drive_source_entry)
+
+add_source_entry(source_type=LOCAL_CONNECTOR_TYPE, entry=local_source_entry)
+add_destination_entry(destination_type=LOCAL_CONNECTOR_TYPE, entry=local_destination_entry)
+
+add_source_entry(source_type=ONEDRIVE_CONNECTOR_TYPE, entry=onedrive_source_entry)
+
+add_source_entry(source_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_source_entry)
+add_destination_entry(
+    destination_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_destination_entry
+)
+
+add_source_entry(source_type=SALESFORCE_CONNECTOR_TYPE, entry=salesforce_source_entry)
+
+add_destination_entry(destination_type=WEAVIATE_CONNECTOR_TYPE, entry=weaviate_destination_entry)
+
+add_destination_entry(
+    destination_type=DATABRICKS_VOLUMES_CONNECTOR_TYPE, entry=databricks_volumes_destination_entry
+)
+
+add_destination_entry(destination_type=SQL_CONNECTOR_TYPE, entry=sql_destination_entry)
+
+add_destination_entry(destination_type=MONGODB_CONNECTOR_TYPE, entry=mongodb_destination_entry)
+add_destination_entry(destination_type=PINECONE_CONNECTOR_TYPE, entry=pinecone_destination_entry)
+add_source_entry(source_type=SHAREPOINT_CONNECTOR_TYPE, entry=sharepoint_source_entry)
+add_destination_entry(
+    destination_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_destination_entry
+)