unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/processes/partitioner.py

@@ -0,0 +1,233 @@
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field, SecretStr
+
+from unstructured_ingest.error import UserError
+from unstructured_ingest.interfaces.process import BaseProcess
+from unstructured_ingest.logger import logger
+from unstructured_ingest.unstructured_api import call_api_async
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+
+class PartitionerConfig(BaseModel):
+    strategy: str = Field(
+        default="auto",
+        description="The method that will be used to process the documents. ",
+        examples=["fast", "hi_res", "auto"],
+    )
+    ocr_languages: Optional[list[str]] = Field(
+        default=None,
+        description="A list of language packs to specify which languages to use for OCR, "
+        "The appropriate Tesseract language pack needs to be installed.",
+        examples=["eng", "deu", "eng,deu"],
+    )
+    encoding: Optional[str] = Field(
+        default=None,
+        description="Text encoding to use when reading documents. "
+        "By default the encoding is detected automatically.",
+    )
+    additional_partition_args: Optional[dict[str, Any]] = Field(
+        default=None, description="Additional values to pass through to partition()"
+    )
+    skip_infer_table_types: Optional[list[str]] = Field(
+        default=None, description="Optional list of document data_types to skip table extraction on"
+    )
+    fields_include: list[str] = Field(
+        default_factory=lambda: ["element_id", "text", "type", "metadata", "embeddings"],
+        description="If set, include the specified top-level fields in an element.",
+    )
+    flatten_metadata: bool = Field(
+        default=False,
+        description="Results in flattened json elements. "
+        "Specifically, the metadata key values are brought to "
+        "the top-level of the element, and the `metadata` key itself is removed.",
+    )
+    metadata_exclude: list[str] = Field(
+        default_factory=list,
+        description="If set, drop the specified metadata fields if they exist.",
+    )
+    element_exclude: list[str] = Field(
+        default_factory=list,
+        description="If set, drop the specified element data_types, if they exist.",
+    )
+    metadata_include: list[str] = Field(
+        default_factory=list,
+        description="If set, include the specified metadata "
+        "fields if they exist and drop all other fields. ",
+    )
+    partition_endpoint: Optional[str] = Field(
+        default="https://api.unstructuredapp.io/general/v0/general",
+        description="If partitioning via api, use the following host.",
+    )
+    partition_by_api: bool = Field(
+        default=False,
+        description="Use a remote API to partition the files."
+        " Otherwise, use the function from partition.auto",
+    )
+    api_timeout_ms: Optional[int] = Field(
+        default=None, description="Timeout in milliseconds for all api call during partitioning."
+    )
+    api_key: Optional[SecretStr] = Field(
+        default=None, description="API Key for partition endpoint."
+    )
+    hi_res_model_name: Optional[str] = Field(
+        default=None, description="Model name for hi-res strategy."
+    )
+    raise_unsupported_filetype: bool = Field(
+        default=False, description="Raise an error if the file type is not supported"
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        if self.metadata_exclude and self.metadata_include:
+            raise ValueError(
+                "metadata_exclude and metadata_include are "
+                "mutually exclusive with each other. Cannot specify both."
+            )
+
+    def to_partition_kwargs(self) -> dict[str, Any]:
+        partition_kwargs: dict[str, Any] = {
+            "strategy": self.strategy,
+            "languages": self.ocr_languages,
+            "hi_res_model_name": self.hi_res_model_name,
+            "skip_infer_table_types": self.skip_infer_table_types,
+        }
+        # Don't inject information if None and allow default values in method to be used
+        partition_kwargs = {k: v for k, v in partition_kwargs.items() if v is not None}
+        if self.additional_partition_args:
+            partition_kwargs.update(self.additional_partition_args)
+        return partition_kwargs
+
+
+@dataclass
+class Partitioner(BaseProcess, ABC):
+    config: PartitionerConfig
+
+    def is_async(self) -> bool:
+        return self.config.partition_by_api
+
+    def postprocess(self, elements: list[dict]) -> list[dict]:
+        element_dicts = [e.copy() for e in elements]
+        if self.config.element_exclude:
+            element_dicts = list(
+                filter(
+                    lambda element: element["type"] not in self.config.element_exclude,
+                    element_dicts,
+                )
+            )
+        for elem in element_dicts:
+            if self.config.metadata_exclude:
+                ex_list = self.config.metadata_exclude
+                for ex in ex_list:
+                    if "." in ex:  # handle nested fields
+                        nested_fields = ex.split(".")
+                        current_elem = elem
+                        for f in nested_fields[:-1]:
+                            if f in current_elem:
+                                current_elem = current_elem[f]
+                        field_to_exclude = nested_fields[-1]
+                        if field_to_exclude in current_elem:
+                            current_elem.pop(field_to_exclude, None)
+                    else:  # handle top-level fields
+                        elem["metadata"].pop(ex, None)  # type: ignore[attr-defined]
+            elif self.config.metadata_include:
+                in_list = self.config.metadata_include
+                for k in list(elem["metadata"].keys()):  # type: ignore[attr-defined]
+                    if k not in in_list:
+                        elem["metadata"].pop(k, None)  # type: ignore[attr-defined]
+            in_list = self.config.fields_include
+            elem = {k: v for k, v in elem.items() if k in in_list}
+
+            if self.config.flatten_metadata and "metadata" in elem:
+                metadata = elem.pop("metadata")
+                elem.update(flatten_dict(metadata, keys_to_omit=["data_source_record_locator"]))
+        return element_dicts
+
+    @requires_dependencies(dependencies=["unstructured"])
+    def partition_locally(
+        self, filename: Path, metadata: Optional[dict] = None, **kwargs
+    ) -> list[dict]:
+        from unstructured.documents.elements import DataSourceMetadata
+        from unstructured.partition.auto import partition
+        from unstructured.staging.base import elements_to_dicts
+
+        @dataclass
+        class FileDataSourceMetadata(DataSourceMetadata):
+            filesize_bytes: Optional[int] = None
+
+        metadata = metadata or {}
+        logger.debug(f"using local partition with kwargs: {self.config.to_partition_kwargs()}")
+        logger.debug(f"partitioning file {filename} with metadata {metadata}")
+        try:
+            elements = partition(
+                filename=str(filename.resolve()),
+                data_source_metadata=FileDataSourceMetadata.from_dict(metadata),
+                **self.config.to_partition_kwargs(),
+            )
+        except ValueError as sdk_error:
+            if (
+                self.is_unstructured_error_unsupported_filetype(sdk_error=sdk_error)
+                and not self.config.raise_unsupported_filetype
+            ):
+                logger.warning(
+                    f"Unsupported file type for strategy {self.config.strategy}: {filename}"
+                )
+                return []
+            raise sdk_error
+        return self.postprocess(elements=elements_to_dicts(elements))
+
+    @requires_dependencies(dependencies=["unstructured_client"], extras="remote")
+    async def partition_via_api(
+        self, filename: Path, metadata: Optional[dict] = None, **kwargs
+    ) -> list[dict]:
+        metadata = metadata or {}
+        logger.debug(f"partitioning file {filename} with metadata: {metadata}")
+
+        elements = await call_api_async(
+            server_url=self.config.partition_endpoint,
+            api_key=self.config.api_key.get_secret_value(),
+            filename=filename,
+            api_parameters=self.config.to_partition_kwargs(),
+            timeout_ms=self.config.api_timeout_ms,
+        )
+
+        # Append the data source metadata the auto partition does for you
+        for element in elements:
+            element["metadata"]["data_source"] = metadata
+        return self.postprocess(elements=elements)
+
+    def is_unstructured_error_unsupported_filetype(self, sdk_error: ValueError) -> bool:
+        error_msg = sdk_error.args[0]
+        return (
+            "Invalid file" in error_msg
+            or "Unstructured schema" in error_msg
+            or "fast strategy is not available for image files" in error_msg
+        )
+
+    def is_client_error_unsupported_filetype(self, error: UserError) -> bool:
+        error_msg = error.args[0]
+        return "fast strategy is not available for image files" in error_msg or (
+            "file type" in error_msg.lower() and "is not supported" in error_msg.lower()
+        )
+
+    def run(self, filename: Path, metadata: Optional[dict] = None, **kwargs) -> list[dict]:
+        return self.partition_locally(filename, metadata=metadata, **kwargs)
+
+    async def run_async(
+        self, filename: Path, metadata: Optional[dict] = None, **kwargs
+    ) -> list[dict]:
+        try:
+            return await self.partition_via_api(filename, metadata=metadata, **kwargs)
+        except UserError as user_error:
+            if (
+                self.is_client_error_unsupported_filetype(error=user_error)
+                and not self.config.raise_unsupported_filetype
+            ):
+                logger.warning(
+                    f"Unsupported file type for strategy {self.config.strategy}: {filename}"
+                )
+                return []
+            raise user_error

unstructured_ingest/processes/uncompress.py

@@ -0,0 +1,61 @@
+from abc import ABC
+from copy import copy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+from uuid import NAMESPACE_DNS, uuid5
+
+from pydantic import BaseModel
+
+from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.interfaces.process import BaseProcess
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.compression import TAR_FILE_EXT, ZIP_FILE_EXT, uncompress_file
+
+
+class UncompressConfig(BaseModel):
+    pass
+
+
+@dataclass
+class Uncompressor(BaseProcess, ABC):
+    config: UncompressConfig = field(default_factory=UncompressConfig)
+
+    def is_async(self) -> bool:
+        return True
+
+    def run(self, file_data: FileData, **kwargs: Any) -> list[FileData]:
+        local_filepath = Path(file_data.local_download_path)
+        if local_filepath.suffix not in TAR_FILE_EXT + ZIP_FILE_EXT:
+            return [file_data]
+        new_path = uncompress_file(filename=str(local_filepath))
+        new_files = [i for i in Path(new_path).rglob("*") if i.is_file()]
+        responses = []
+        logger.debug(
+            "uncompressed {} files from original file {}: {}".format(
+                len(new_files), local_filepath, ", ".join([str(f) for f in new_files])
+            )
+        )
+        for f in new_files:
+            new_file_data = copy(file_data)
+            new_file_data.identifier = str(uuid5(NAMESPACE_DNS, str(f)))
+            new_file_data.local_download_path = str(f.resolve())
+            new_rel_download_path = str(f).replace(str(Path(local_filepath.parent)), "")[1:]
+            new_file_data.source_identifiers = SourceIdentifiers(
+                filename=f.name,
+                fullpath=str(file_data.source_identifiers.fullpath).replace(
+                    file_data.source_identifiers.filename, new_rel_download_path
+                ),
+                rel_path=(
+                    str(file_data.source_identifiers.rel_path).replace(
+                        file_data.source_identifiers.filename, new_rel_download_path
+                    )
+                    if file_data.source_identifiers.rel_path
+                    else None
+                ),
+            )
+            responses.append(new_file_data)
+        return responses
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> list[FileData]:
+        return self.run(file_data=file_data, **kwargs)

unstructured_ingest/processes/utils/__init__.py

@@ -0,0 +1,8 @@
+from .logging.connector import ConnectorLoggingMixin, LoggingConfig
+from .logging.sanitizer import DataSanitizer
+
+__all__ = [
+    "ConnectorLoggingMixin",
+    "DataSanitizer",
+    "LoggingConfig",
+]

unstructured_ingest/processes/utils/blob_storage.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.interfaces import UploadStager, UploadStagerConfig
+from unstructured_ingest.utils.data_prep import get_json_data, write_data
+
+
+class BlobStoreUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class BlobStoreUploadStager(UploadStager):
+    upload_stager_config: BlobStoreUploadStagerConfig = field(
+        default_factory=BlobStoreUploadStagerConfig
+    )
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        output_file = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
+        # Always save as json
+        data = get_json_data(elements_filepath)
+        write_data(path=output_file.with_suffix(".json"), data=data)
+        return output_file.with_suffix(".json")