unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/logger.py

@@ -0,0 +1,130 @@
+import ast
+import json
+import logging
+import typing as t
+
+logger = logging.getLogger("unstructured_ingest")
+
+
+def default_is_data_sensitive(k: str, v: t.Any) -> bool:
+    sensitive_fields = [
+        "account_name",
+        "client_id",
+    ]
+    sensitive_triggers = ["key", "cred", "token", "password", "oauth", "secret"]
+    return (
+        v
+        and any([s in k.lower() for s in sensitive_triggers])  # noqa: C419
+        or k.lower() in sensitive_fields
+    )
+
+
+def hide_sensitive_fields(
+    data: dict, is_sensitive_fn: t.Callable[[str, t.Any], bool] = default_is_data_sensitive
+) -> dict:
+    """
+    Will recursively look through every k, v pair in this dict and any nested ones and run
+    is_sensitive_fn to dynamically redact the value of the k, v pair. Will also check if
+    any string value can be parsed as valid json and process that dict as well and replace
+    the original string with the json.dumps() version of the redacted dict.
+    """
+    new_data = data.copy()
+    for k, v in new_data.items():
+        if is_sensitive_fn(k, v):
+            new_data[k] = "*******"
+        if isinstance(v, dict):
+            new_data[k] = hide_sensitive_fields(v)
+        if isinstance(v, str):
+            # Need to take into account strings generated via json.dumps() or simply printing a dict
+            try:
+                json_data = json.loads(v)
+                if isinstance(json_data, dict):
+                    updated_data = hide_sensitive_fields(json_data)
+                    new_data[k] = json.dumps(updated_data)
+            except json.JSONDecodeError:
+                pass
+
+    return new_data
+
+
+def redact_jsons(s: str) -> str:
+    """
+    Takes in a generic string and pulls out all valid json content. Leverages
+    hide_sensitive_fields() to redact any sensitive information and replaces the
+    original json with the new redacted format. There can be any number of valid
+    jsons in a generic string and this will work. Having extra '{' without a
+    closing '}' will cause this to break though. i.e '{ text, {"a": 3}'.
+
+    """
+    chars = list(s)
+    if "{" not in chars:
+        return s
+    i = 0
+    jsons = []
+    i = 0
+    while i < len(chars):
+        char = chars[i]
+        if char == "{":
+            stack = [char]
+            current = [char]
+            while len(stack) != 0 and i < len(chars):
+                i += 1
+                char = chars[i]
+                current.append(char)
+                if char == "{":
+                    stack.append(char)
+                if char == "}":
+                    stack.pop(-1)
+            jsons.append("".join(current))
+            continue
+        i += 1
+    for j in jsons:
+        try:
+            formatted_j = json.dumps(json.loads(j))
+        except json.JSONDecodeError:
+            formatted_j = json.dumps(ast.literal_eval(j))
+        hidden_j = json.dumps(hide_sensitive_fields(json.loads(formatted_j)))
+        s = s.replace(j, hidden_j)
+    return s
+
+
+class SensitiveFormatter(logging.Formatter):
+    def format(self, record):
+        s = super().format(record=record)
+        return redact_jsons(s)
+
+
+def remove_root_handlers(logger: logging.Logger) -> None:
+    # NOTE(robinson): in some environments such as Google Colab, there is a root handler
+    # that doesn't not mask secrets, meaning sensitive info such as api keys appear in logs.
+    # Removing these when they exist prevents this behavior
+    if logger.root.hasHandlers():
+        for handler in logger.root.handlers:
+            logger.root.removeHandler(handler)
+
+
+def ingest_log_streaming_init(level: int) -> None:
+    handler = logging.StreamHandler()
+    handler.name = "ingest_log_handler"
+    formatter = SensitiveFormatter("%(asctime)s %(processName)-10s %(levelname)-8s %(message)s")
+    handler.setFormatter(formatter)
+
+    # Only want to add the handler once
+    if "ingest_log_handler" not in [h.name for h in logger.handlers]:
+        logger.addHandler(handler)
+
+    remove_root_handlers(logger)
+    logger.setLevel(level)
+
+
+def make_default_logger(level: int) -> logging.Logger:
+    """Return a custom logger."""
+    logger = logging.getLogger("unstructured_ingest")
+    handler = logging.StreamHandler()
+    handler.name = "ingest_log_handler"
+    formatter = SensitiveFormatter("%(asctime)s %(processName)-10s %(levelname)-8s %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(level)
+    remove_root_handlers(logger)
+    return logger

unstructured_ingest/main.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+from unstructured_ingest.cli.cli import get_cmd
+
+
+def main():
+    ingest_cmd = get_cmd()
+    ingest_cmd()
+
+
+if __name__ == "__main__":
+    main()

unstructured_ingest/pipeline/__init__.py

@@ -0,0 +1,22 @@
+from .doc_factory import DocFactory
+from .interfaces import PipelineContext, ReformatNode
+from .partition import Partitioner
+from .permissions import PermissionsDataCleaner
+from .pipeline import Pipeline
+from .reformat.chunking import Chunker
+from .reformat.embedding import Embedder
+from .source import Reader
+from .write import Writer
+
+__all__ = [
+    "DocFactory",
+    "Partitioner",
+    "Reader",
+    "Embedder",
+    "PipelineContext",
+    "Pipeline",
+    "Writer",
+    "Chunker",
+    "ReformatNode",
+    "PermissionsDataCleaner",
+]

unstructured_ingest/pipeline/copy.py

@@ -0,0 +1,19 @@
+import os
+import shutil
+from pathlib import Path
+
+from unstructured_ingest.connector.registry import create_ingest_doc_from_dict
+from unstructured_ingest.logger import logger
+from unstructured_ingest.pipeline.interfaces import CopyNode
+
+
+class Copier(CopyNode):
+    def run(self, json_path: str):
+        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_doc = create_ingest_doc_from_dict(ingest_doc_dict)
+        desired_output = ingest_doc._output_filename
+        Path(desired_output).parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"copying {json_path} -> {desired_output}")
+        shutil.copy(json_path, desired_output)

unstructured_ingest/pipeline/doc_factory.py

@@ -0,0 +1,12 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.pipeline.interfaces import DocFactoryNode
+
+
+@dataclass
+class DocFactory(DocFactoryNode):
+    def run(self, *args, **kwargs) -> t.Iterable[dict]:
+        docs = self.source_doc_connector.get_ingest_docs()
+        json_docs = [doc.to_dict() for doc in docs]
+        return json_docs

unstructured_ingest/pipeline/interfaces.py

@@ -0,0 +1,270 @@
+import hashlib
+import json
+import logging
+import multiprocessing as mp
+import typing as t
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from multiprocessing.managers import DictProxy
+from pathlib import Path
+
+from dataclasses_json import DataClassJsonMixin
+
+from unstructured_ingest.error import SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    BaseDestinationConnector,
+    BaseSourceConnector,
+    PartitionConfig,
+    ProcessorConfig,
+    ReadConfig,
+    RetryStrategyConfig,
+)
+from unstructured_ingest.logger import ingest_log_streaming_init, logger
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.ingest_backoff import RetryHandler
+
+
+@dataclass
+class PipelineContext(ProcessorConfig):
+    """
+    Data that gets shared across each pipeline node
+    """
+
+    def __post_init__(self):
+        self._ingest_docs_map: t.Optional[DictProxy] = None
+
+    @property
+    def ingest_docs_map(self) -> DictProxy:
+        if self._ingest_docs_map is None:
+            raise ValueError("ingest_docs_map never initialized")
+        return self._ingest_docs_map
+
+    @ingest_docs_map.setter
+    def ingest_docs_map(self, value: DictProxy):
+        self._ingest_docs_map = value
+
+
+@dataclass
+class PipelineNode(DataClassJsonMixin, ABC):
+    """
+    Class that encapsulates logic to run during a single pipeline step
+    """
+
+    pipeline_context: PipelineContext
+
+    def __call__(self, iterable: t.Optional[t.Iterable[t.Any]] = None) -> t.Any:
+        iterable = iterable if iterable else []
+        if iterable:
+            logger.info(
+                f"calling {self.__class__.__name__} " f"with {len(iterable)} docs",  # type: ignore
+            )
+
+        self.initialize()
+        if not self.supported_multiprocessing():
+            if iterable:
+                self.result = self.run(iterable)
+            else:
+                self.result = self.run()
+        elif self.pipeline_context.num_processes == 1:
+            if iterable:
+                self.result = [self.run(it) for it in iterable]
+            else:
+                self.result = self.run()
+        else:
+            with mp.Pool(
+                processes=self.pipeline_context.num_processes,
+                initializer=ingest_log_streaming_init,
+                initargs=(logging.DEBUG if self.pipeline_context.verbose else logging.INFO,),
+            ) as pool:
+                self.result = pool.map(self.run, iterable)
+        # Remove None which may be caused by failed docs that didn't raise an error
+        if isinstance(self.result, t.Iterable):
+            self.result = [r for r in self.result if r is not None]
+        return self.result
+
+    def supported_multiprocessing(self) -> bool:
+        return True
+
+    @abstractmethod
+    def run(self, *args, **kwargs) -> t.Optional[t.Any]:
+        pass
+
+    def initialize(self):
+        if path := self.get_path():
+            logger.info(f"creating {path}")
+            path.mkdir(parents=True, exist_ok=True)
+        ingest_log_streaming_init(logging.DEBUG if self.pipeline_context.verbose else logging.INFO)
+
+    def get_path(self) -> t.Optional[Path]:
+        return None
+
+
+@dataclass
+class DocFactoryNode(PipelineNode):
+    """
+    Encapsulated logic to generate a list of ingest docs
+    """
+
+    source_doc_connector: BaseSourceConnector
+
+    def initialize(self):
+        logger.info(
+            f"Running doc factory to generate ingest docs. "
+            f"Source connector: {self.source_doc_connector.to_json()}",
+        )
+        super().initialize()
+        self.source_doc_connector.initialize()
+
+    @abstractmethod
+    def run(self, *args, **kwargs) -> t.Iterable[dict]:
+        pass
+
+    def supported_multiprocessing(self) -> bool:
+        return False
+
+
+@dataclass
+class SourceNode(PipelineNode):
+    """A pipeline node representing logic to pull data from a source using base ingest documents.
+
+    This class encapsulates the logic for pulling data from a specified source using base ingest
+    documents. The output of this logic is expected to be in JSON format representing the data
+    itself.
+
+    Attributes:
+        read_config: A configuration object specifying how to read data from the source.
+        retry_strategy_config: Optional configuration specifying the strategy for network errors.
+
+    Properties:
+        retry_strategy: A retry handler configured based on the retry strategy configuration.
+
+    Methods:
+        initialize: Initializes the source node and logs the process.
+        run: Abstract method for downloading data associated with ingest documents.
+    """
+
+    read_config: ReadConfig
+    retry_strategy_config: t.Optional[RetryStrategyConfig] = None
+
+    @property
+    def retry_strategy(self) -> t.Optional["RetryHandler"]:
+        if retry_strategy_config := self.retry_strategy_config:
+            import backoff
+
+            from unstructured_ingest.ingest_backoff import RetryHandler
+
+            return RetryHandler(
+                backoff.expo,
+                SourceConnectionNetworkError,
+                max_time=retry_strategy_config.max_retry_time,
+                max_tries=retry_strategy_config.max_retries,
+                logger=logger,
+                start_log_level=logger.level,
+                backoff_log_level=logger.level,
+            )
+        return None
+
+    def initialize(self):
+        logger.info("Running source node to download data associated with ingest docs")
+        super().initialize()
+
+    @abstractmethod
+    def run(self, ingest_doc_json: str) -> t.Optional[str]:
+        pass
+
+
+@dataclass
+class PartitionNode(PipelineNode):
+    """
+    Encapsulates logic to run partition on the json files as the output of the source node
+    """
+
+    partition_config: PartitionConfig
+    partition_kwargs: dict = field(default_factory=dict)
+
+    def initialize(self):
+        logger.info(
+            f"Running partition node to extract content from json files. "
+            f"Config: {self.partition_config.to_json()}, "
+            f"partition kwargs: {json.dumps(self.partition_kwargs)}]",
+        )
+        super().initialize()
+
+    def create_hash(self) -> str:
+        hash_dict = self.partition_config.to_dict()
+        hash_dict["partition_kwargs"] = self.partition_kwargs
+        return hashlib.sha256(json.dumps(hash_dict, sort_keys=True).encode()).hexdigest()[:32]
+
+    @abstractmethod
+    def run(self, json_path: str) -> t.Optional[str]:
+        pass
+
+    def get_path(self) -> Path:
+        return (Path(self.pipeline_context.work_dir) / "partitioned").resolve()
+
+
+@dataclass
+class ReformatNode(PipelineNode, ABC):
+    """
+    Encapsulated any logic to reformat the output List[Element]
+    content from partition before writing it
+    """
+
+    @abstractmethod
+    def run(self, elements_json: str) -> t.Optional[str]:
+        pass
+
+
+@dataclass
+class WriteNode(PipelineNode):
+    """
+    Encapsulated logic to write the final result to a downstream data connection
+    """
+
+    dest_doc_connector: BaseDestinationConnector
+
+    @abstractmethod
+    def run(self, json_paths: t.List[str]):
+        pass
+
+    def initialize(self):
+        logger.info(
+            f"Running write node to upload content. "
+            f"Destination connector: {self.dest_doc_connector.to_json(redact_sensitive=True)}]",
+        )
+        super().initialize()
+        self.dest_doc_connector.initialize()
+
+    def supported_multiprocessing(self) -> bool:
+        return False
+
+
+@dataclass
+class CopyNode(PipelineNode):
+    """
+    Encapsulated logic to copy the final result of the pipeline to the designated output location.
+    """
+
+    def initialize(self):
+        logger.info("Running copy node to move content to desired output location")
+        super().initialize()
+
+    @abstractmethod
+    def run(self, json_path: str):
+        pass
+
+
+@dataclass
+class PermissionsNode(PipelineNode):
+    """
+    Encapsulated logic to do operations on permissions related data.
+    """
+
+    def initialize(self):
+        logger.info("Running permissions node to cleanup the permissions folder")
+        super().initialize()
+
+    @abstractmethod
+    def run(self):
+        pass

unstructured_ingest/pipeline/partition.py

@@ -0,0 +1,60 @@
+import hashlib
+import json
+import typing as t
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from unstructured_ingest.connector.registry import create_ingest_doc_from_dict
+from unstructured_ingest.error import PartitionError
+from unstructured_ingest.logger import logger
+from unstructured_ingest.pipeline.interfaces import PartitionNode
+from unstructured_ingest.pipeline.utils import get_ingest_doc_hash
+
+
+@dataclass
+class Partitioner(PartitionNode):
+    @PartitionError.wrap
+    def run(self, ingest_doc_dict) -> Optional[str]:
+        try:
+            doc = create_ingest_doc_from_dict(ingest_doc_dict)
+            doc_filename_hash = get_ingest_doc_hash(ingest_doc_dict)
+            hashed_filename = hashlib.sha256(
+                f"{self.create_hash()}{doc_filename_hash}".encode(),
+            ).hexdigest()[:32]
+            self.pipeline_context.ingest_docs_map[hashed_filename] = ingest_doc_dict
+            doc_filename = f"{hashed_filename}.json"
+            json_path = (Path(self.get_path()) / doc_filename).resolve()
+            if (
+                not self.pipeline_context.reprocess
+                and json_path.is_file()
+                and json_path.stat().st_size
+            ):
+                logger.info(f"file exists: {json_path}, skipping partition")
+                return str(json_path)
+            partition_kwargs: t.Dict[str, t.Any] = {
+                "strategy": self.partition_config.strategy,
+                "encoding": self.partition_config.encoding,
+                "pdf_infer_table_structure": self.partition_config.pdf_infer_table_structure,
+                "languages": self.partition_config.ocr_languages,
+                "hi_res_model_name": self.partition_config.hi_res_model_name,
+            }
+            if self.partition_config.skip_infer_table_types:
+                partition_kwargs["skip_infer_table_types"] = (
+                    self.partition_config.skip_infer_table_types
+                )
+            if self.partition_config.additional_partition_args:
+                partition_kwargs.update(self.partition_config.additional_partition_args)
+            elements = doc.process_file(
+                partition_config=self.partition_config,
+                **partition_kwargs,
+            )
+            with open(json_path, "w", encoding="utf8") as output_f:
+                logger.info(f"writing partitioned content to {json_path}")
+                json.dump(elements, output_f, ensure_ascii=False, indent=2, sort_keys=True)
+            return str(json_path)
+        except Exception as e:
+            if self.pipeline_context.raise_on_error:
+                raise
+            logger.error(f"failed to partition doc: {ingest_doc_dict}, {e}", exc_info=True)
+            return None

unstructured_ingest/pipeline/permissions.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import PermissionsCleanupMixin, ProcessorConfig
+from unstructured_ingest.pipeline.interfaces import PermissionsNode
+
+
+@dataclass
+class PermissionsDataCleaner(PermissionsNode, PermissionsCleanupMixin):
+    processor_config: ProcessorConfig
+
+    def run(self):
+        self.cleanup_permissions()

unstructured_ingest/pipeline/pipeline.py

@@ -0,0 +1,117 @@
+import logging
+import multiprocessing as mp
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from dataclasses_json import DataClassJsonMixin
+
+from unstructured_ingest.connector.registry import create_ingest_doc_from_dict
+from unstructured_ingest.interfaces import BaseIngestDocBatch, BaseSingleIngestDoc
+from unstructured_ingest.logger import ingest_log_streaming_init, logger
+from unstructured_ingest.pipeline.copy import Copier
+from unstructured_ingest.pipeline.interfaces import (
+    DocFactoryNode,
+    PartitionNode,
+    PipelineContext,
+    ReformatNode,
+    SourceNode,
+    WriteNode,
+)
+from unstructured_ingest.pipeline.permissions import PermissionsDataCleaner
+from unstructured_ingest.pipeline.utils import get_ingest_doc_hash
+
+
+@dataclass
+class Pipeline(DataClassJsonMixin):
+    pipeline_context: PipelineContext
+    doc_factory_node: DocFactoryNode
+    source_node: SourceNode
+    partition_node: Optional[PartitionNode] = None
+    write_node: Optional[WriteNode] = None
+    reformat_nodes: "list[ReformatNode]" = field(default_factory=list)
+    permissions_node: Optional[PermissionsDataCleaner] = None
+
+    def initialize(self):
+        ingest_log_streaming_init(logging.DEBUG if self.pipeline_context.verbose else logging.INFO)
+
+    def get_nodes_str(self):
+        nodes = [self.doc_factory_node, self.source_node, self.partition_node]
+        nodes.extend(self.reformat_nodes)
+        if self.write_node:
+            nodes.append(self.write_node)
+        nodes.append(Copier(pipeline_context=self.pipeline_context))
+        return " -> ".join([node.__class__.__name__ for node in nodes])
+
+    def expand_batch_docs(self, dict_docs: "list[dict[str, Any]]") -> "list[dict[str, Any]]":
+        expanded_docs: list[dict[str, Any]] = []
+        for d in dict_docs:
+            doc = create_ingest_doc_from_dict(d)
+            if isinstance(doc, BaseSingleIngestDoc):
+                expanded_docs.append(doc.to_dict())
+            elif isinstance(doc, BaseIngestDocBatch):
+                expanded_docs.extend([single_doc.to_dict() for single_doc in doc.ingest_docs])
+            else:
+                raise ValueError(
+                    f"type of doc ({type(doc)}) is not a recognized type: "
+                    f"BaseSingleIngestDoc or BaseSingleIngestDoc"
+                )
+        return expanded_docs
+
+    def run(self):
+        logger.info(
+            f"running pipeline: {self.get_nodes_str()} "
+            f"with config: {self.pipeline_context.to_json()}",
+        )
+        self.initialize()
+        manager = mp.Manager()
+        self.pipeline_context.ingest_docs_map = manager.dict()
+        # -- Get the documents to be processed --
+        dict_docs = self.doc_factory_node()
+        dict_docs = [manager.dict(d) for d in dict_docs]
+        if not dict_docs:
+            logger.info("no docs found to process")
+            return
+        logger.info(
+            f"processing {len(dict_docs)} docs via "
+            f"{self.pipeline_context.num_processes} processes",
+        )
+        for doc in dict_docs:
+            self.pipeline_context.ingest_docs_map[get_ingest_doc_hash(doc)] = doc
+        fetched_filenames = self.source_node(iterable=dict_docs)
+        if self.source_node.read_config.download_only:
+            logger.info("stopping pipeline after downloading files")
+            return
+        if not fetched_filenames:
+            logger.info("No files to run partition over")
+            return
+        # -- To support batches ingest docs, expand those into the populated single ingest
+        # -- docs after downloading content
+        dict_docs = self.expand_batch_docs(dict_docs=dict_docs)
+        if self.partition_node is None:
+            raise ValueError("partition node not set")
+        partitioned_jsons = self.partition_node(iterable=dict_docs)
+        if not partitioned_jsons:
+            logger.info("No files to process after partitioning")
+            return
+        for reformat_node in self.reformat_nodes:
+            reformatted_jsons = reformat_node(iterable=partitioned_jsons)
+            if not reformatted_jsons:
+                logger.info(f"no files to process after {reformat_node.__class__.__name__}")
+                return
+            partitioned_jsons = reformatted_jsons
+
+        # -- Copy the final destination to the desired location --
+        copier = Copier(
+            pipeline_context=self.pipeline_context,
+        )
+        copier(iterable=partitioned_jsons)
+
+        if self.write_node:
+            logger.info(
+                f"uploading elements from {len(partitioned_jsons)} "
+                "document(s) to the destination"
+            )
+            self.write_node(iterable=partitioned_jsons)
+
+        if self.permissions_node:
+            self.permissions_node.cleanup_permissions()