unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/pipeline/reformat/chunking.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os.path
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+from unstructured_ingest.interfaces import ChunkingConfig, PartitionConfig
+from unstructured_ingest.logger import logger
+from unstructured_ingest.pipeline.interfaces import ReformatNode
+from unstructured_ingest.utils.chunking import assign_and_map_hash_ids
+
+if TYPE_CHECKING:
+    from unstructured.documents.elements import Element
+
+
+@dataclass
+class Chunker(ReformatNode):
+    """Implementation for the chunking node in the ingest pipeline.
+
+    Parameters
+    ----------
+    pipeline_context: PipelineContext (inherited from parent class)
+    chunking_config: ChunkingConfig
+    partition_config: PartitionConfig
+    """
+
+    chunking_config: ChunkingConfig
+    partition_config: PartitionConfig
+
+    def initialize(self):
+        logger.info(
+            f"Running chunking node. Chunking config: {self.chunking_config.to_json()}]",
+        )
+        super().initialize()
+
+    def create_hash(self) -> str:
+        hash_dict = self.chunking_config.to_dict()
+        return hashlib.sha256(json.dumps(hash_dict, sort_keys=True).encode()).hexdigest()[:32]
+
+    def run(self, elements_json: str) -> Optional[str]:
+        try:
+            elements_json_filename = os.path.basename(elements_json)
+            filename_ext = os.path.basename(elements_json_filename)
+            filename = os.path.splitext(filename_ext)[0]
+            hashed_filename = hashlib.sha256(
+                f"{self.create_hash()}{filename}".encode(),
+            ).hexdigest()[:32]
+            json_filename = f"{hashed_filename}.json"
+            json_path = (Path(self.get_path()) / json_filename).resolve()
+            self.pipeline_context.ingest_docs_map[hashed_filename] = (
+                self.pipeline_context.ingest_docs_map[filename]
+            )
+            if (
+                not self.pipeline_context.reprocess
+                and json_path.is_file()
+                and json_path.stat().st_size
+            ):
+                logger.debug(f"file exists: {json_path}, skipping chunking")
+                return str(json_path)
+
+            chunked_elements = self.chunk(elements_json)
+
+            # -- return if chunking_strategy is None --
+            if chunked_elements is None:
+                logger.info(f"chunking_strategy is None, skipping chunking for {filename_ext}")
+                return
+
+            element_dicts = [e.to_dict() for e in chunked_elements]
+            assign_and_map_hash_ids(elements=element_dicts)
+
+            with open(json_path, "w", encoding="utf8") as output_f:
+                logger.info(f"writing chunking content to {json_path}")
+                json.dump(element_dicts, output_f, ensure_ascii=False, indent=2)
+                return str(json_path)
+
+        except Exception as e:
+            if self.pipeline_context.raise_on_error:
+                raise
+            logger.error(f"failed to run chunking on file {elements_json}, {e}", exc_info=True)
+            return None
+
+    def get_path(self) -> Path:
+        return (Path(self.pipeline_context.work_dir) / "chunked").resolve()
+
+    def chunk(self, elements_json_file: str) -> Optional[list["Element"]]:
+        """Called by Chunker.run() to properly execute the defined chunking_strategy."""
+        # -- No chunking_strategy means no chunking --
+        if self.chunking_config.chunking_strategy is None:
+            return
+        # -- Chunk locally for open-source chunking strategies, even when partitioning remotely --
+        if self.chunking_config.chunking_strategy in ("basic", "by_title"):
+            from unstructured.chunking import dispatch
+            from unstructured.staging.base import elements_from_json
+
+            return dispatch.chunk(
+                elements=elements_from_json(filename=elements_json_file),
+                chunking_strategy=self.chunking_config.chunking_strategy,
+                combine_text_under_n_chars=self.chunking_config.combine_text_under_n_chars,
+                include_orig_elements=self.chunking_config.include_orig_elements,
+                max_characters=self.chunking_config.max_characters,
+                multipage_sections=self.chunking_config.multipage_sections,
+                new_after_n_chars=self.chunking_config.new_after_n_chars,
+                overlap=self.chunking_config.overlap,
+                overlap_all=self.chunking_config.overlap_all,
+            )
+        # -- Chunk remotely --
+        if self.partition_config.partition_by_api:
+            from unstructured.partition.api import partition_via_api
+
+            return partition_via_api(
+                filename=elements_json_file,
+                # -- NOTE(jennings): If api_key or api_url are None, partition_via_api will raise an
+                # -- error, which will be caught and logged by Chunker.run()
+                api_key=self.partition_config.api_key,  # type: ignore
+                api_url=self.partition_config.partition_endpoint,  # type: ignore
+                chunking_strategy=self.chunking_config.chunking_strategy,
+                combine_under_n_chars=self.chunking_config.combine_text_under_n_chars,
+                include_orig_elements=self.chunking_config.include_orig_elements,
+                max_characters=self.chunking_config.max_characters,
+                multipage_sections=self.chunking_config.multipage_sections,
+                new_after_n_chars=self.chunking_config.new_after_n_chars,
+                overlap=self.chunking_config.overlap,
+                overlap_all=self.chunking_config.overlap_all,
+            )
+        # -- Warn that the defined chunking_strategy is not locally available --
+        logger.warning(
+            f"There is no locally available chunking_strategy:"
+            f" {self.chunking_config.chunking_strategy}."
+            f" If trying to partition remotely, check that `partition_by_api`, `api_url`,"
+            f" and `api_key` are correctly defined."
+        )

unstructured_ingest/pipeline/reformat/embedding.py

@@ -0,0 +1,64 @@
+import hashlib
+import json
+import os.path
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from unstructured_ingest.interfaces import (
+    EmbeddingConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.pipeline.interfaces import ReformatNode
+
+
+@dataclass
+class Embedder(ReformatNode):
+    embedder_config: EmbeddingConfig
+
+    def initialize(self):
+        logger.info(
+            f"Running embedding node. Embedding config: {self.embedder_config.to_json()}]",
+        )
+        super().initialize()
+
+    def create_hash(self) -> str:
+        hash_dict = self.embedder_config.to_dict()
+        return hashlib.sha256(json.dumps(hash_dict, sort_keys=True).encode()).hexdigest()[:32]
+
+    def run(self, elements_json: str) -> Optional[str]:
+        try:
+            elements_json_filename = os.path.basename(elements_json)
+            filename_ext = os.path.basename(elements_json_filename)
+            filename = os.path.splitext(filename_ext)[0]
+            hashed_filename = hashlib.sha256(
+                f"{self.create_hash()}{filename}".encode(),
+            ).hexdigest()[:32]
+            json_filename = f"{hashed_filename}.json"
+            json_path = (Path(self.get_path()) / json_filename).resolve()
+            self.pipeline_context.ingest_docs_map[hashed_filename] = (
+                self.pipeline_context.ingest_docs_map[filename]
+            )
+            if (
+                not self.pipeline_context.reprocess
+                and json_path.is_file()
+                and json_path.stat().st_size
+            ):
+                logger.debug(f"file exists: {json_path}, skipping embedding")
+                return str(json_path)
+            with open(elements_json) as f:
+                elements = json.load(f)
+            embedder = self.embedder_config.get_embedder()
+            element_dicts = embedder.embed_documents(elements=elements)
+            with open(json_path, "w", encoding="utf8") as output_f:
+                logger.info(f"writing embeddings content to {json_path}")
+                json.dump(element_dicts, output_f, ensure_ascii=False, indent=2)
+            return str(json_path)
+        except Exception as e:
+            if self.pipeline_context.raise_on_error:
+                raise
+            logger.error(f"failed to embed content from file {elements_json}, {e}", exc_info=True)
+            return None
+
+    def get_path(self) -> Path:
+        return (Path(self.pipeline_context.work_dir) / "embedded.py").resolve()

unstructured_ingest/pipeline/source.py

@@ -0,0 +1,77 @@
+import os
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.connector.registry import create_ingest_doc_from_dict
+from unstructured_ingest.interfaces import (
+    BaseIngestDocBatch,
+    BaseSessionHandle,
+    BaseSingleIngestDoc,
+    IngestDocSessionHandleMixin,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.pipeline.interfaces import SourceNode
+
+# module-level variable to store session handle
+session_handle: t.Optional[BaseSessionHandle] = None
+
+
+@dataclass
+class Reader(SourceNode):
+    def get_single(self, doc: BaseSingleIngestDoc, ingest_doc_dict: dict) -> str:
+        if (
+            not self.read_config.re_download
+            and doc.filename.is_file()
+            and doc.filename.stat().st_size
+        ):
+            logger.info(f"file exists: {doc.filename}, skipping download")
+            # Still need to fetch metadata if file exists locally
+            doc.update_source_metadata()
+        else:
+            serialized_doc = doc.to_json(redact_sensitive=True)
+            logger.debug(f"fetching {serialized_doc} - PID: {os.getpid()}")
+            if self.retry_strategy:
+                self.retry_strategy(doc.get_file)
+            else:
+                doc.get_file()
+        for k, v in doc.to_dict().items():
+            ingest_doc_dict[k] = v
+        return doc.filename
+
+    def get_batch(self, doc_batch: BaseIngestDocBatch, ingest_doc_dict: dict) -> t.List[str]:
+        if self.retry_strategy:
+            self.retry_strategy(doc_batch.get_files)
+        else:
+            doc_batch.get_files()
+        for k, v in doc_batch.to_dict().items():
+            ingest_doc_dict[k] = v
+        return [doc.filename for doc in doc_batch.ingest_docs]
+
+    def run(self, ingest_doc_dict: dict) -> t.Optional[t.Union[str, t.List[str]]]:
+        try:
+            global session_handle
+            doc = create_ingest_doc_from_dict(ingest_doc_dict)
+            if isinstance(doc, IngestDocSessionHandleMixin):
+                if session_handle is None:
+                    # create via doc.session_handle, which is a property that creates a
+                    # session handle if one is not already defined
+                    session_handle = doc.session_handle
+                else:
+                    doc._session_handle = session_handle
+            if isinstance(doc, BaseSingleIngestDoc):
+                return self.get_single(doc=doc, ingest_doc_dict=ingest_doc_dict)
+            elif isinstance(doc, BaseIngestDocBatch):
+                return self.get_batch(doc_batch=doc, ingest_doc_dict=ingest_doc_dict)
+            else:
+                raise ValueError(
+                    f"type of doc ({type(doc)}) is not a recognized type: "
+                    f"BaseSingleIngestDoc or BaseSingleIngestDoc"
+                )
+        except Exception as e:
+            if self.pipeline_context.raise_on_error:
+                raise
+            logger.error(
+                f"failed to get data associated with source doc: {ingest_doc_dict}, {e}",
+                exc_info=True,
+            )
+            return None

unstructured_ingest/pipeline/utils.py

@@ -0,0 +1,6 @@
+import hashlib
+
+
+def get_ingest_doc_hash(json_as_dict: dict) -> str:
+    hashed = hashlib.sha256(json_as_dict["unique_id"].encode()).hexdigest()[:32]
+    return hashed

unstructured_ingest/pipeline/write.py

@@ -0,0 +1,18 @@
+import os.path
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.connector.registry import create_ingest_doc_from_dict
+from unstructured_ingest.pipeline.interfaces import WriteNode
+
+
+@dataclass
+class Writer(WriteNode):
+    def run(self, json_paths: t.List[str]):
+        ingest_docs = []
+        for json_path in json_paths:
+            filename = os.path.basename(json_path)
+            doc_hash = os.path.splitext(filename)[0]
+            ingest_doc_dict = self.pipeline_context.ingest_docs_map[doc_hash]
+            ingest_docs.append(create_ingest_doc_from_dict(ingest_doc_dict))
+        self.dest_doc_connector.write(docs=ingest_docs)

unstructured_ingest/processor.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import multiprocessing as mp
+from contextlib import suppress
+from typing import Optional
+
+from unstructured_ingest.interfaces import (
+    BaseDestinationConnector,
+    BaseSourceConnector,
+    ChunkingConfig,
+    EmbeddingConfig,
+    PartitionConfig,
+    PermissionsConfig,
+    ProcessorConfig,
+    RetryStrategyConfig,
+)
+from unstructured_ingest.pipeline import (
+    Chunker,
+    DocFactory,
+    Embedder,
+    Partitioner,
+    PermissionsDataCleaner,
+    Pipeline,
+    PipelineContext,
+    Reader,
+    ReformatNode,
+    Writer,
+)
+
+with suppress(RuntimeError):
+    mp.set_start_method("spawn")
+
+
+def process_documents(
+    processor_config: ProcessorConfig,
+    source_doc_connector: BaseSourceConnector,
+    partition_config: PartitionConfig,
+    dest_doc_connector: Optional[BaseDestinationConnector] = None,
+    chunking_config: Optional[ChunkingConfig] = None,
+    embedder_config: Optional[EmbeddingConfig] = None,
+    permissions_config: Optional[PermissionsConfig] = None,
+    retry_strategy_config: Optional[RetryStrategyConfig] = None,
+) -> None:
+    pipeline_config = PipelineContext.from_dict(processor_config.to_dict())
+    doc_factory = DocFactory(
+        pipeline_context=pipeline_config,
+        source_doc_connector=source_doc_connector,
+    )
+    reader = Reader(
+        pipeline_context=pipeline_config,
+        retry_strategy_config=retry_strategy_config,
+        read_config=source_doc_connector.read_config,
+    )
+    partitioner = Partitioner(pipeline_context=pipeline_config, partition_config=partition_config)
+    reformat_nodes: list[ReformatNode] = []
+    if chunking_config:
+        reformat_nodes.append(
+            Chunker(
+                pipeline_context=pipeline_config,
+                chunking_config=chunking_config,
+                partition_config=partition_config,
+            ),
+        )
+    if embedder_config:
+        reformat_nodes.append(
+            Embedder(
+                pipeline_context=pipeline_config,
+                embedder_config=embedder_config,
+            ),
+        )
+    writer = (
+        Writer(
+            pipeline_context=pipeline_config,
+            dest_doc_connector=dest_doc_connector,
+        )
+        if dest_doc_connector
+        else None
+    )
+    permissions_data_cleaner = (
+        PermissionsDataCleaner(pipeline_context=pipeline_config, processor_config=processor_config)
+        if permissions_config
+        else None
+    )
+    pipeline = Pipeline(
+        pipeline_context=pipeline_config,
+        doc_factory_node=doc_factory,
+        source_node=reader,
+        partition_node=partitioner,
+        reformat_nodes=reformat_nodes,
+        write_node=writer,
+        permissions_node=permissions_data_cleaner,
+    )
+    pipeline.run()

unstructured_ingest/runner/__init__.py

@@ -0,0 +1,104 @@
+import typing as t
+from typing import Type
+
+from .airtable import AirtableRunner
+from .astradb import AstraDBRunner
+from .base_runner import Runner
+from .biomed import BiomedRunner
+from .confluence import ConfluenceRunner
+from .delta_table import DeltaTableRunner
+from .discord import DiscordRunner
+from .elasticsearch import ElasticSearchRunner
+from .fsspec.azure import AzureRunner
+from .fsspec.box import BoxRunner
+from .fsspec.dropbox import DropboxRunner
+from .fsspec.fsspec import FsspecRunner
+from .fsspec.gcs import GCSRunner
+from .fsspec.s3 import S3Runner
+from .fsspec.sftp import SftpRunner
+from .github import GithubRunner
+from .gitlab import GitlabRunner
+from .google_drive import GoogleDriveRunner
+from .hubspot import HubSpotRunner
+from .jira import JiraRunner
+from .kafka import KafkaRunner
+from .local import LocalRunner
+from .mongodb import MongoDBRunner
+from .notion import NotionRunner
+from .onedrive import OneDriveRunner
+from .opensearch import OpenSearchRunner
+from .outlook import OutlookRunner
+from .reddit import RedditRunner
+from .salesforce import SalesforceRunner
+from .sharepoint import SharePointRunner
+from .slack import SlackRunner
+from .wikipedia import WikipediaRunner
+
+runner_map: t.Dict[str, Type[Runner]] = {
+    "airtable": AirtableRunner,
+    "astradb": AstraDBRunner,
+    "azure": AzureRunner,
+    "biomed": BiomedRunner,
+    "box": BoxRunner,
+    "confluence": ConfluenceRunner,
+    "delta_table": DeltaTableRunner,
+    "discord": DiscordRunner,
+    "dropbox": DropboxRunner,
+    "elasticsearch": ElasticSearchRunner,
+    "fsspec": FsspecRunner,
+    "gcs": GCSRunner,
+    "github": GithubRunner,
+    "gitlab": GitlabRunner,
+    "gdrive": GoogleDriveRunner,
+    "google_drive": GoogleDriveRunner,
+    "hubspot": HubSpotRunner,
+    "jira": JiraRunner,
+    "kafka": KafkaRunner,
+    "local": LocalRunner,
+    "mongodb": MongoDBRunner,
+    "notion": NotionRunner,
+    "onedrive": OneDriveRunner,
+    "opensearch": OpenSearchRunner,
+    "outlook": OutlookRunner,
+    "reddit": RedditRunner,
+    "s3": S3Runner,
+    "salesforce": SalesforceRunner,
+    "sftp": SftpRunner,
+    "sharepoint": SharePointRunner,
+    "slack": SlackRunner,
+    "wikipedia": WikipediaRunner,
+}
+
+__all__ = [
+    "AirtableRunner",
+    "AstraRunner",
+    "AzureRunner",
+    "BiomedRunner",
+    "BoxRunner",
+    "ConfluenceRunner",
+    "DeltaTableRunner",
+    "DiscordRunner",
+    "DropboxRunner",
+    "ElasticSearchRunner",
+    "FsspecRunner",
+    "GCSRunner",
+    "GoogleDriveRunner",
+    "GithubRunner",
+    "GitlabRunner",
+    "JiraRunner",
+    "KafkaRunner",
+    "LocalRunner",
+    "MongoDBRunner",
+    "NotionRunner",
+    "OneDriveRunner",
+    "OpenSearchRunner",
+    "OutlookRunner",
+    "RedditRunner",
+    "S3Runner",
+    "SalesforceRunner",
+    "SharePointRunner",
+    "SlackRunner",
+    "WikipediaRunner",
+    "runner_map",
+    "Runner",
+]

unstructured_ingest/runner/airtable.py

@@ -0,0 +1,35 @@
+import hashlib
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseSourceConnector
+from unstructured_ingest.logger import logger
+from unstructured_ingest.runner.base_runner import Runner
+from unstructured_ingest.runner.utils import update_download_dir_hash
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.airtable import SimpleAirtableConfig
+
+
+@dataclass
+class AirtableRunner(Runner):
+    connector_config: "SimpleAirtableConfig"
+
+    def update_read_config(self):
+        hashed_dir_name = hashlib.sha256(
+            self.connector_config.access_config.personal_access_token.encode("utf-8"),
+        )
+
+        self.read_config.download_dir = update_download_dir_hash(
+            connector_name="airtable",
+            read_config=self.read_config,
+            hashed_dir_name=hashed_dir_name,
+            logger=logger,
+        )
+
+    def get_source_connector_cls(self) -> t.Type[BaseSourceConnector]:
+        from unstructured_ingest.connector.airtable import (
+            AirtableSourceConnector,
+        )
+
+        return AirtableSourceConnector

unstructured_ingest/runner/astradb.py

@@ -0,0 +1,34 @@
+import hashlib
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseSourceConnector
+from unstructured_ingest.logger import logger
+from unstructured_ingest.runner.base_runner import Runner
+from unstructured_ingest.runner.utils import update_download_dir_hash
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.astradb import SimpleAstraDBConfig
+
+
+@dataclass
+class AstraDBRunner(Runner):
+    connector_config: "SimpleAstraDBConfig"
+
+    def update_read_config(self):
+        hashed_dir_name = hashlib.sha256(
+            str(self.connector_config.access_config.api_endpoint).encode("utf-8"),
+        )
+        self.read_config.download_dir = update_download_dir_hash(
+            connector_name="astradb",
+            read_config=self.read_config,
+            hashed_dir_name=hashed_dir_name,
+            logger=logger,
+        )
+
+    def get_source_connector_cls(self) -> t.Type[BaseSourceConnector]:
+        from unstructured_ingest.connector.astradb import (
+            AstraDBSourceConnector,
+        )
+
+        return AstraDBSourceConnector

unstructured_ingest/runner/base_runner.py

@@ -0,0 +1,89 @@
+import logging
+import typing as t
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseSourceConnector,
+    ChunkingConfig,
+    EmbeddingConfig,
+    PartitionConfig,
+    PermissionsConfig,
+    ProcessorConfig,
+    ReadConfig,
+    RetryStrategyConfig,
+)
+from unstructured_ingest.logger import ingest_log_streaming_init
+from unstructured_ingest.processor import process_documents
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+
+@dataclass
+class Runner(EnhancedDataClassJsonMixin, ABC):
+    connector_config: BaseConnectorConfig
+    processor_config: ProcessorConfig
+    read_config: ReadConfig
+    partition_config: PartitionConfig
+    writer: t.Optional[Writer] = None
+    writer_kwargs: t.Optional[dict] = None
+    embedding_config: t.Optional[EmbeddingConfig] = None
+    chunking_config: t.Optional[ChunkingConfig] = None
+    permissions_config: t.Optional[PermissionsConfig] = None
+    retry_strategy_config: t.Optional[RetryStrategyConfig] = None
+
+    def run(self, *args, **kwargs):
+        ingest_log_streaming_init(logging.DEBUG if self.processor_config.verbose else logging.INFO)
+        self.update_read_config()
+        source_connector = self.get_source_connector()
+        self.process_documents(
+            source_doc_connector=source_connector,
+        )
+
+    @abstractmethod
+    def update_read_config(self):
+        pass
+
+    @abstractmethod
+    def get_source_connector_cls(self) -> t.Type[BaseSourceConnector]:
+        pass
+
+    def get_source_connector(self) -> BaseSourceConnector:
+        source_connector_cls = self.get_source_connector_cls()
+        return source_connector_cls(
+            processor_config=self.processor_config,
+            connector_config=self.connector_config,
+            read_config=self.read_config,
+        )
+
+    def get_dest_doc_connector(self) -> t.Optional[BaseDestinationConnector]:
+        writer_kwargs = self.writer_kwargs if self.writer_kwargs else {}
+        if self.writer:
+            return self.writer.get_connector(**writer_kwargs)
+        return None
+
+    def get_permissions_config(self) -> t.Optional[PermissionsConfig]:
+        if self.permissions_config is None:
+            return None
+
+        permissions_config_filled = bool(
+            self.permissions_config.application_id
+            and self.permissions_config.client_cred
+            and self.permissions_config.tenant,
+        )
+
+        return self.permissions_config if permissions_config_filled else None
+
+    def process_documents(self, source_doc_connector: BaseSourceConnector):
+        process_documents(
+            processor_config=self.processor_config,
+            source_doc_connector=source_doc_connector,
+            partition_config=self.partition_config,
+            dest_doc_connector=self.get_dest_doc_connector(),
+            embedder_config=self.embedding_config,
+            chunking_config=self.chunking_config,
+            permissions_config=self.get_permissions_config(),
+            retry_strategy_config=self.retry_strategy_config,
+        )