unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/runner/writers/fsspec/azure.py

@@ -0,0 +1,24 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.fsspec.azure import (
+        AzureWriteConfig,
+        SimpleAzureBlobStorageConfig,
+    )
+
+
+@dataclass
+class AzureWriter(Writer):
+    connector_config: "SimpleAzureBlobStorageConfig"
+    write_config: "AzureWriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.fsspec.azure import (
+            AzureBlobStorageDestinationConnector,
+        )
+
+        return AzureBlobStorageDestinationConnector

unstructured_ingest/runner/writers/fsspec/box.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.fsspec.box import BoxWriteConfig, SimpleBoxConfig
+
+
+@dataclass
+class BoxWriter(Writer):
+    connector_config: "SimpleBoxConfig"
+    write_config: "BoxWriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.fsspec.box import (
+            BoxDestinationConnector,
+        )
+
+        return BoxDestinationConnector

unstructured_ingest/runner/writers/fsspec/dropbox.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.fsspec.dropbox import DropboxWriteConfig, SimpleDropboxConfig
+
+
+@dataclass
+class DropboxWriter(Writer):
+    connector_config: "SimpleDropboxConfig"
+    write_config: "DropboxWriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.fsspec.dropbox import (
+            DropboxDestinationConnector,
+        )
+
+        return DropboxDestinationConnector

unstructured_ingest/runner/writers/fsspec/gcs.py

@@ -0,0 +1,19 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.fsspec.gcs import GcsWriteConfig, SimpleGcsConfig
+
+
+@dataclass
+class GcsWriter(Writer):
+    connector_config: "SimpleGcsConfig"
+    write_config: "GcsWriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.fsspec.gcs import GcsDestinationConnector
+
+        return GcsDestinationConnector

unstructured_ingest/runner/writers/fsspec/s3.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.fsspec.s3 import S3WriteConfig, SimpleS3Config
+
+
+@dataclass
+class S3Writer(Writer):
+    connector_config: "SimpleS3Config"
+    write_config: "S3WriteConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.fsspec.s3 import (
+            S3DestinationConnector,
+        )
+
+        return S3DestinationConnector

unstructured_ingest/runner/writers/kafka.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.kafka import KafkaWriteConfig, SimpleKafkaConfig
+
+
+@dataclass
+class KafkaWriter(Writer):
+    write_config: "KafkaWriteConfig"
+    connector_config: "SimpleKafkaConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.kafka import (
+            KafkaDestinationConnector,
+        )
+
+        return KafkaDestinationConnector

unstructured_ingest/runner/writers/mongodb.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.mongodb import MongoDBWriteConfig, SimpleMongoDBConfig
+
+
+@dataclass
+class MongodbWriter(Writer):
+    write_config: "MongoDBWriteConfig"
+    connector_config: "SimpleMongoDBConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.mongodb import (
+            MongoDBDestinationConnector,
+        )
+
+        return MongoDBDestinationConnector

unstructured_ingest/runner/writers/opensearch.py

@@ -0,0 +1,26 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.elasticsearch import (
+        ElasticsearchWriteConfig,
+    )
+    from unstructured_ingest.connector.opensearch import (
+        SimpleOpenSearchConfig,
+    )
+
+
+@dataclass
+class OpenSearchWriter(Writer):
+    connector_config: "SimpleOpenSearchConfig"
+    write_config: "ElasticsearchWriteConfig"
+
+    def get_connector_cls(self) -> BaseDestinationConnector:
+        from unstructured_ingest.connector.opensearch import (
+            OpenSearchDestinationConnector,
+        )
+
+        return OpenSearchDestinationConnector

unstructured_ingest/runner/writers/pinecone.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.pinecone import PineconeWriteConfig, SimplePineconeConfig
+
+
+@dataclass
+class PineconeWriter(Writer):
+    write_config: "PineconeWriteConfig"
+    connector_config: "SimplePineconeConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.pinecone import (
+            PineconeDestinationConnector,
+        )
+
+        return PineconeDestinationConnector

unstructured_ingest/runner/writers/qdrant.py

@@ -0,0 +1,19 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.qdrant import QdrantWriteConfig, SimpleQdrantConfig
+
+
+@dataclass
+class QdrantWriter(Writer):
+    write_config: "QdrantWriteConfig"
+    connector_config: "SimpleQdrantConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.qdrant import QdrantDestinationConnector
+
+        return QdrantDestinationConnector

unstructured_ingest/runner/writers/sql.py

@@ -0,0 +1,22 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.sql import SimpleSqlConfig
+    from unstructured_ingest.interfaces import WriteConfig
+
+
+@dataclass
+class SqlWriter(Writer):
+    write_config: "WriteConfig"
+    connector_config: "SimpleSqlConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.sql import (
+            SqlDestinationConnector,
+        )
+
+        return SqlDestinationConnector

unstructured_ingest/runner/writers/vectara.py

@@ -0,0 +1,22 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.vectara import SimpleVectaraConfig, VectaraWriteConfig
+
+
+@dataclass
+class VectaraWriter(Writer, EnhancedDataClassJsonMixin):
+    write_config: "VectaraWriteConfig"
+    connector_config: "SimpleVectaraConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.vectara import (
+            VectaraDestinationConnector,
+        )
+
+        return VectaraDestinationConnector

unstructured_ingest/runner/writers/weaviate.py

@@ -0,0 +1,21 @@
+import typing as t
+from dataclasses import dataclass
+
+from unstructured_ingest.interfaces import BaseDestinationConnector
+from unstructured_ingest.runner.writers.base_writer import Writer
+
+if t.TYPE_CHECKING:
+    from unstructured_ingest.connector.weaviate import SimpleWeaviateConfig, WeaviateWriteConfig
+
+
+@dataclass
+class WeaviateWriter(Writer):
+    write_config: "WeaviateWriteConfig"
+    connector_config: "SimpleWeaviateConfig"
+
+    def get_connector_cls(self) -> t.Type[BaseDestinationConnector]:
+        from unstructured_ingest.connector.weaviate import (
+            WeaviateDestinationConnector,
+        )
+
+        return WeaviateDestinationConnector

unstructured_ingest/utils/chunking.py

@@ -0,0 +1,56 @@
+import base64
+import hashlib
+import json
+import zlib
+from itertools import groupby
+
+
+def id_to_hash(element: dict, sequence_number: int) -> str:
+    """Calculates and assigns a deterministic hash as an ID.
+
+    The hash ID is based on element's text, sequence number on page,
+    page number and its filename.
+
+    Args:
+        sequence_number: index on page
+
+    Returns: new ID value
+    """
+    filename = element["metadata"].get("filename")
+    text = element["text"]
+    page_number = element["metadata"].get("page_number")
+    data = f"{filename}{text}{page_number}{sequence_number}"
+    element["element_id"] = hashlib.sha256(data.encode()).hexdigest()[:32]
+    return element["element_id"]
+
+
+def assign_and_map_hash_ids(elements: list[dict]) -> list[dict]:
+    # -- generate sequence number for each element on a page --
+    elements = elements.copy()
+    page_numbers = [e["metadata"].get("page_number") for e in elements]
+    page_seq_pairs = [
+        seq_on_page for page, group in groupby(page_numbers) for seq_on_page, _ in enumerate(group)
+    ]
+
+    # -- assign hash IDs to elements --
+    old_to_new_mapping = {
+        element["element_id"]: id_to_hash(element=element, sequence_number=seq_on_page_counter)
+        for element, seq_on_page_counter in zip(elements, page_seq_pairs)
+    }
+
+    # -- map old parent IDs to new ones --
+    for e in elements:
+        parent_id = e["metadata"].get("parent_id")
+        if not parent_id:
+            continue
+        e["metadata"]["parent_id"] = old_to_new_mapping[parent_id]
+
+    return elements
+
+
+def elements_from_base64_gzipped_json(raw_s: str) -> list[dict]:
+    decoded_b64_bytes = base64.b64decode(raw_s)
+    elements_json_bytes = zlib.decompress(decoded_b64_bytes)
+    elements_json_str = elements_json_bytes.decode("utf-8")
+    element_dicts = json.loads(elements_json_str)
+    return element_dicts

unstructured_ingest/utils/compression.py

@@ -0,0 +1,118 @@
+import copy
+import os
+import sys
+import tarfile
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+from unstructured_ingest.connector.local import LocalSourceConnector, SimpleLocalConfig
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    ProcessorConfig,
+    ReadConfig,
+)
+from unstructured_ingest.logger import logger
+
+ZIP_FILE_EXT = [".zip"]
+TAR_FILE_EXT = [".tar", ".tar.gz", ".tgz"]
+
+
+def uncompress_file(filename: str, path: Optional[str] = None) -> str:
+    """
+    Takes in a compressed zip or tar file and decompresses it
+    """
+    # Create path if it doesn't already exist
+    if path:
+        Path(path).mkdir(parents=True, exist_ok=True)
+
+    if any(filename.endswith(ext) for ext in ZIP_FILE_EXT):
+        return uncompress_zip_file(zip_filename=filename, path=path)
+    elif any(filename.endswith(ext) for ext in TAR_FILE_EXT):
+        return uncompress_tar_file(tar_filename=filename, path=path)
+    else:
+        raise ValueError(
+            "filename {} not a recognized compressed extension: {}".format(
+                filename,
+                ", ".join(ZIP_FILE_EXT + TAR_FILE_EXT),
+            ),
+        )
+
+
+def uncompress_zip_file(zip_filename: str, path: Optional[str] = None) -> str:
+    head, tail = os.path.split(zip_filename)
+    for ext in ZIP_FILE_EXT:
+        if tail.endswith(ext):
+            tail = tail[: -(len(ext))]
+            break
+    path = path if path else os.path.join(head, f"{tail}-zip-uncompressed")
+    logger.info(f"extracting zip {zip_filename} -> {path}")
+    with zipfile.ZipFile(zip_filename) as zfile:
+        zfile.extractall(path=path)
+    return path
+
+
+def uncompress_tar_file(tar_filename: str, path: Optional[str] = None) -> str:
+    head, tail = os.path.split(tar_filename)
+    for ext in TAR_FILE_EXT:
+        if tail.endswith(ext):
+            tail = tail[: -(len(ext))]
+            break
+
+    path = path if path else os.path.join(head, f"{tail}-tar-uncompressed")
+    logger.info(f"extracting tar {tar_filename} -> {path}")
+    # NOTE: "r:*" mode opens both compressed (e.g ".tar.gz") and uncompressed ".tar" archives
+    with tarfile.open(tar_filename, "r:*") as tfile:
+        # NOTE(robinson): Mitigate against malicious content being extracted from the tar file.
+        # This was added in Python 3.12
+        # Ref: https://docs.python.org/3/library/tarfile.html#extraction-filters
+        if sys.version_info >= (3, 12):
+            tfile.extraction_filter = tarfile.tar_filter
+        else:
+            logger.warning(
+                "Extraction filtering for tar files is available for Python 3.12 and above. "
+                "Consider upgrading your Python version to improve security. "
+                "See https://docs.python.org/3/library/tarfile.html#extraction-filters"
+            )
+        tfile.extractall(path=path)
+    return path
+
+
+@dataclass
+class CompressionSourceConnectorMixin:
+    processor_config: ProcessorConfig
+    read_config: ReadConfig
+    connector_config: BaseConnectorConfig
+
+    def process_compressed_doc(self, doc: BaseSingleIngestDoc) -> List[BaseSingleIngestDoc]:
+        """
+        Utility function which helps process compressed files. Extracts the contents and returns
+        generated ingest docs via local source connector
+        """
+        # Download the raw file to local
+        doc.get_file()
+        path = uncompress_file(filename=str(doc.filename))
+        new_read_configs = copy.copy(self.read_config)
+        new_process_configs = copy.copy(self.processor_config)
+        relative_path = path.replace(self.read_config.download_dir, "")
+
+        if self.processor_config.output_dir.endswith(os.sep):
+            new_process_configs.output_dir = f"{self.processor_config.output_dir}{relative_path}"
+        else:
+            new_process_configs.output_dir = (
+                f"{self.processor_config.output_dir}{os.sep}{relative_path}"
+            )
+
+        local_connector = LocalSourceConnector(
+            connector_config=SimpleLocalConfig(
+                input_path=path,
+                recursive=True,
+            ),
+            read_config=new_read_configs,
+            processor_config=new_process_configs,
+        )
+        logger.info(f"created local source connector: {local_connector.to_json()}")
+        local_connector.initialize()
+        return local_connector.get_ingest_docs()

unstructured_ingest/utils/data_prep.py

@@ -0,0 +1,200 @@
+import itertools
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Generator, Iterable, Optional, Sequence, TypeVar, cast
+
+import ndjson
+import pandas as pd
+
+from unstructured_ingest.v2.logger import logger
+
+DATE_FORMATS = ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d+%H:%M:%S", "%Y-%m-%dT%H:%M:%S%z")
+
+T = TypeVar("T")
+IterableT = Iterable[T]
+
+
+def split_dataframe(df: pd.DataFrame, chunk_size: int = 100) -> Generator[pd.DataFrame, None, None]:
+    num_chunks = len(df) // chunk_size + 1
+    for i in range(num_chunks):
+        yield df[i * chunk_size : (i + 1) * chunk_size]
+
+
+def batch_generator(iterable: IterableT, batch_size: int = 100) -> IterableT:
+    """A helper function to break an iterable into batches of size batch_size."""
+    it = iter(iterable)
+    chunk = tuple(itertools.islice(it, batch_size))
+    while chunk:
+        yield chunk
+        chunk = tuple(itertools.islice(it, batch_size))
+
+
+def generator_batching_wbytes(
+    iterable: IterableT,
+    batch_size_limit_bytes: Optional[int] = None,
+    max_batch_size: Optional[int] = None,
+) -> IterableT:
+    if not batch_size_limit_bytes and not max_batch_size:
+        return iterable
+    """A helper function to break an iterable into chunks of specified bytes."""
+    current_batch, current_batch_size = [], 0
+
+    for item in iterable:
+        item_size_bytes = len(json.dumps(item).encode("utf-8"))
+        if batch_size_limit_bytes and current_batch_size + item_size_bytes > batch_size_limit_bytes:
+            yield current_batch
+            current_batch, current_batch_size = [item], item_size_bytes
+            continue
+        if max_batch_size and len(current_batch) + 1 > max_batch_size:
+            yield current_batch
+            current_batch, current_batch_size = [item], item_size_bytes
+            continue
+
+        current_batch.append(item)
+        current_batch_size += item_size_bytes
+
+    if current_batch:
+        yield current_batch
+
+
+def flatten_dict(
+    dictionary: dict[str, Any],
+    parent_key: str = "",
+    separator: str = "_",
+    flatten_lists: bool = False,
+    remove_none: bool = False,
+    keys_to_omit: Optional[Sequence[str]] = None,
+) -> dict[str, Any]:
+    """Flattens a nested dictionary into a single level dictionary.
+
+    keys_to_omit is a list of keys that don't get flattened. If omitting a nested key, format as
+    {parent_key}{separator}{key}. If flatten_lists is True, then lists and tuples are flattened as
+    well. If remove_none is True, then None keys/values are removed from the flattened
+    dictionary.
+    """
+    keys_to_omit = keys_to_omit if keys_to_omit else []
+    flattened_dict: dict[str, Any] = {}
+    for key, value in dictionary.items():
+        new_key = f"{parent_key}{separator}{key}" if parent_key else key
+        if new_key in keys_to_omit:
+            flattened_dict[new_key] = value
+        elif value is None and remove_none:
+            continue
+        elif isinstance(value, dict):
+            value = cast("dict[str, Any]", value)
+            flattened_dict.update(
+                flatten_dict(
+                    value, new_key, separator, flatten_lists, remove_none, keys_to_omit=keys_to_omit
+                ),
+            )
+        elif isinstance(value, (list, tuple)) and flatten_lists:
+            value = cast("list[Any] | tuple[Any]", value)
+            for index, item in enumerate(value):
+                flattened_dict.update(
+                    flatten_dict(
+                        {f"{new_key}{separator}{index}": item},
+                        "",
+                        separator,
+                        flatten_lists,
+                        remove_none,
+                        keys_to_omit=keys_to_omit,
+                    )
+                )
+        else:
+            flattened_dict[new_key] = value
+
+    return flattened_dict
+
+
+def validate_date_args(date: Optional[str] = None) -> bool:
+    """Validate whether the provided date string satisfies any of the supported date formats.
+
+    Used by unstructured/ingest/connector/biomed.py
+
+    Returns `True` if the date string satisfies any of the supported formats, otherwise raises
+    `ValueError`.
+
+    Supported Date Formats:
+        - 'YYYY-MM-DD'
+        - 'YYYY-MM-DDTHH:MM:SS'
+        - 'YYYY-MM-DD+HH:MM:SS'
+        - 'YYYY-MM-DDTHH:MM:SS±HHMM'
+    """
+    if not date:
+        raise ValueError("The argument date is None.")
+
+    for format in DATE_FORMATS:
+        try:
+            datetime.strptime(date, format)
+            return True
+        except ValueError:
+            pass
+
+    raise ValueError(
+        f"The argument {date} does not satisfy the format:"
+        f" YYYY-MM-DD or YYYY-MM-DDTHH:MM:SS or YYYY-MM-DD+HH:MM:SS or YYYY-MM-DDTHH:MM:SS±HHMM",
+    )
+
+
+def get_data_by_suffix(path: Path) -> list[dict]:
+    with path.open() as f:
+        if path.suffix == ".json":
+            return json.load(f)
+        elif path.suffix == ".ndjson":
+            return ndjson.load(f)
+        elif path.suffix == ".csv":
+            df = pd.read_csv(path)
+            return df.to_dict(orient="records")
+        elif path.suffix == ".parquet":
+            df = pd.read_parquet(path)
+            return df.to_dict(orient="records")
+        else:
+            raise ValueError(f"Unsupported file type: {path}")
+
+
+def get_data(path: Path) -> list[dict]:
+    try:
+        return get_data_by_suffix(path=path)
+    except Exception as e:
+        logger.warning(f"failed to read {path} by extension: {e}")
+    # Fall back
+    with path.open() as f:
+        try:
+            return json.load(f)
+        except Exception as e:
+            logger.warning(f"failed to read {path} as json: {e}")
+        try:
+            return ndjson.load(f)
+        except Exception as e:
+            logger.warning(f"failed to read {path} as ndjson: {e}")
+        try:
+            df = pd.read_csv(path)
+            return df.to_dict(orient="records")
+        except Exception as e:
+            logger.warning(f"failed to read {path} as csv: {e}")
+        try:
+            df = pd.read_parquet(path)
+            return df.to_dict(orient="records")
+        except Exception as e:
+            logger.warning(f"failed to read {path} as parquet: {e}")
+
+    raise IOError(f"File could not be parsed: {path}")
+
+
+def get_data_df(path: Path) -> pd.DataFrame:
+    with path.open() as f:
+        if path.suffix == ".json":
+            data = json.load(f)
+            return pd.DataFrame(data=data)
+        elif path.suffix == ".ndjson":
+            data = ndjson.load(f)
+            return pd.DataFrame(data=data)
+        elif path.suffix == ".csv":
+            df = pd.read_csv(path)
+            return df
+        elif path.suffix == ".parquet":
+            df = pd.read_parquet(path)
+            return df
+        else:
+            raise ValueError(f"Unsupported file type: {path}")