unstructured-ingest 0.0.2.dev0__py3-none-any.whl → 0.0.4__py3-none-any.whl

unstructured_ingest/v2/processes/embedder.py

@@ -1,76 +1,135 @@
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
-from unstructured.documents.elements import Element
-from unstructured.embed.interfaces import BaseEmbeddingEncoder
-from unstructured.staging.base import elements_from_json
+from pydantic import BaseModel, Field, SecretStr
 
-from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin, enhanced_field
+from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces.process import BaseProcess
 
+if TYPE_CHECKING:
+    from unstructured.embed.interfaces import BaseEmbeddingEncoder
 
-@dataclass
-class EmbedderConfig(EnhancedDataClassJsonMixin):
-    embedding_provider: Optional[str] = None
-    embedding_api_key: Optional[str] = enhanced_field(default=None, sensitive=True)
-    embedding_model_name: Optional[str] = None
-    embedding_aws_access_key_id: Optional[str] = None
-    embedding_aws_secret_access_key: Optional[str] = None
-    embedding_aws_region: Optional[str] = None
-
-    def get_embedder(self) -> BaseEmbeddingEncoder:
+
+class EmbedderConfig(BaseModel):
+    embedding_provider: Optional[
+        Literal[
+            "langchain-openai",
+            "langchain-huggingface",
+            "langchain-aws-bedrock",
+            "langchain-vertexai",
+            "langchain-voyageai",
+            "octoai",
+        ]
+    ] = Field(default=None, description="Type of the embedding class to be used.")
+    embedding_api_key: Optional[SecretStr] = Field(
+        default=None,
+        description="API key for the embedding model, for the case an API key is needed.",
+    )
+    embedding_model_name: Optional[str] = Field(
+        default=None,
+        description="Embedding model name, if needed. "
+        "Chooses a particular LLM between different options, to embed with it.",
+    )
+    embedding_aws_access_key_id: Optional[str] = Field(
+        default=None, description="AWS access key used for AWS-based embedders, such as bedrock"
+    )
+    embedding_aws_secret_access_key: Optional[SecretStr] = Field(
+        default=None, description="AWS secret key used for AWS-based embedders, such as bedrock"
+    )
+    embedding_aws_region: Optional[str] = Field(
+        default="us-west-2", description="AWS region used for AWS-based embedders, such as bedrock"
+    )
+
+    @requires_dependencies(dependencies=["unstructured"], extras="embed-huggingface")
+    def get_huggingface_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.huggingface import (
+            HuggingFaceEmbeddingConfig,
+            HuggingFaceEmbeddingEncoder,
+        )
+
+        return HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig(**embedding_kwargs))
+
+    @requires_dependencies(dependencies=["unstructured"], extras="openai")
+    def get_openai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+
+        return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(**embedding_kwargs))
+
+    @requires_dependencies(dependencies=["unstructured"], extras="embed-octoai")
+    def get_octoai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+
+        return OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(**embedding_kwargs))
+
+    @requires_dependencies(dependencies=["unstructured"], extras="bedrock")
+    def get_bedrock_embedder(self) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.bedrock import BedrockEmbeddingConfig, BedrockEmbeddingEncoder
+
+        return BedrockEmbeddingEncoder(
+            config=BedrockEmbeddingConfig(
+                aws_access_key_id=self.embedding_aws_access_key_id,
+                aws_secret_access_key=self.embedding_aws_secret_access_key.get_secret_value(),
+                region_name=self.embedding_aws_region,
+            )
+        )
+
+    @requires_dependencies(dependencies=["unstructured"], extras="embed-vertexai")
+    def get_vertexai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.vertexai import (
+            VertexAIEmbeddingConfig,
+            VertexAIEmbeddingEncoder,
+        )
+
+        return VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(**embedding_kwargs))
+
+    @requires_dependencies(dependencies=["unstructured"], extras="embed-voyageai")
+    def get_voyageai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder
+
+        return VoyageAIEmbeddingEncoder(config=VoyageAIEmbeddingConfig(**embedding_kwargs))
+
+    def get_embedder(self) -> "BaseEmbeddingEncoder":
         kwargs: dict[str, Any] = {}
         if self.embedding_api_key:
-            kwargs["api_key"] = self.embedding_api_key
+            kwargs["api_key"] = self.embedding_api_key.get_secret_value()
         if self.embedding_model_name:
             kwargs["model_name"] = self.embedding_model_name
         # TODO make this more dynamic to map to encoder configs
         if self.embedding_provider == "langchain-openai":
-            from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+            return self.get_openai_embedder(embedding_kwargs=kwargs)
 
-            return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(**kwargs))
-        elif self.embedding_provider == "langchain-huggingface":
-            from unstructured.embed.huggingface import (
-                HuggingFaceEmbeddingConfig,
-                HuggingFaceEmbeddingEncoder,
-            )
+        if self.embedding_provider == "langchain-huggingface":
+            return self.get_huggingface_embedder(embedding_kwargs=kwargs)
 
-            return HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig(**kwargs))
-        elif self.embedding_provider == "octoai":
-            from unstructured.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
+        if self.embedding_provider == "octoai":
+            return self.get_octoai_embedder(embedding_kwargs=kwargs)
 
-            return OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(**kwargs))
-        elif self.embedding_provider == "langchain-aws-bedrock":
-            from unstructured.embed.bedrock import BedrockEmbeddingConfig, BedrockEmbeddingEncoder
+        if self.embedding_provider == "langchain-aws-bedrock":
+            return self.get_bedrock_embedder()
 
-            return BedrockEmbeddingEncoder(
-                config=BedrockEmbeddingConfig(
-                    aws_access_key_id=self.embedding_aws_access_key_id,
-                    aws_secret_access_key=self.embedding_aws_secret_access_key,
-                    region_name=self.embedding_aws_region,
-                )
-            )
-        elif self.embedding_provider == "langchain-vertexai":
-            from unstructured.embed.vertexai import (
-                VertexAIEmbeddingConfig,
-                VertexAIEmbeddingEncoder,
-            )
+        if self.embedding_provider == "langchain-vertexai":
+            return self.get_vertexai_embedder(embedding_kwargs=kwargs)
 
-            return VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(**kwargs))
-        else:
-            raise ValueError(f"{self.embedding_provider} not a recognized encoder")
+        if self.embedding_provider == "langchain-voyageai":
+            return self.get_voyageai_embedder(embedding_kwargs=kwargs)
+
+        raise ValueError(f"{self.embedding_provider} not a recognized encoder")
 
 
 @dataclass
 class Embedder(BaseProcess, ABC):
     config: EmbedderConfig
 
-    def run(self, elements_filepath: Path, **kwargs: Any) -> list[Element]:
+    @requires_dependencies(dependencies=["unstructured"])
+    def run(self, elements_filepath: Path, **kwargs: Any) -> list[dict]:
+        from unstructured.staging.base import elements_from_json
+
         # TODO update base embedder classes to support async
         embedder = self.config.get_embedder()
         elements = elements_from_json(filename=str(elements_filepath))
         if not elements:
-            return elements
-        return embedder.embed_documents(elements=elements)
+            return [e.to_dict() for e in elements]
+        embedded_elements = embedder.embed_documents(elements=elements)
+        return [e.to_dict() for e in embedded_elements]

unstructured_ingest/v2/processes/filter.py

@@ -0,0 +1,60 @@
+import fnmatch
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+from pydantic import BaseModel, Field
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.process import BaseProcess
+from unstructured_ingest.v2.logger import logger
+
+
+class FiltererConfig(BaseModel):
+    file_glob: Optional[list[str]] = Field(
+        default=None,
+        description="file globs to limit which types of " "files are accepted",
+        examples=["*.pdf", "*.html"],
+    )
+    max_file_size: Optional[int] = Field(
+        default=None, description="Max file size to process in bytes"
+    )
+
+
+@dataclass
+class Filterer(BaseProcess, ABC):
+    config: FiltererConfig = field(default_factory=lambda: FiltererConfig())
+    filters: list[Callable[[FileData], bool]] = field(init=False, default_factory=list)
+
+    def __post_init__(self):
+        # Populate the filters based on values in config
+        if self.config.file_glob is not None:
+            self.filters.append(self.glob_filter)
+        if self.config.max_file_size:
+            self.filters.append(self.file_size_filter)
+
+    def is_async(self) -> bool:
+        return False
+
+    def file_size_filter(self, file_data: FileData) -> bool:
+        if filesize_bytes := file_data.metadata.filesize_bytes:
+            return filesize_bytes <= self.config.max_file_size
+        return True
+
+    def glob_filter(self, file_data: FileData) -> bool:
+        patterns = self.config.file_glob
+        path = file_data.source_identifiers.fullpath
+        for pattern in patterns:
+            if fnmatch.filter([path], pattern):
+                return True
+        logger.debug(f"The file {path!r} is discarded as it does not match any given glob.")
+        return False
+
+    def run(self, file_data: FileData, **kwargs: Any) -> Optional[FileData]:
+        for filter in self.filters:
+            if not filter(file_data):
+                logger.debug(
+                    f"filtered out file data due to {filter.__name__}: {file_data.identifier}"
+                )
+                return None
+        return file_data

unstructured_ingest/v2/processes/partitioner.py

@@ -1,14 +1,13 @@
 import asyncio
 from abc import ABC
-from dataclasses import dataclass, field, fields
+from dataclasses import dataclass, fields
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
-from unstructured.documents.elements import DataSourceMetadata
-from unstructured.staging.base import elements_to_dicts, flatten_dict
+from pydantic import BaseModel, Field, SecretStr
 
-from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin
-from unstructured_ingest.enhanced_dataclass.dataclasses import enhanced_field
+from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces.process import BaseProcess
 from unstructured_ingest.v2.logger import logger
 
@@ -17,25 +16,65 @@ if TYPE_CHECKING:
     from unstructured_client.models.shared import PartitionParameters
 
 
-@dataclass
-class PartitionerConfig(EnhancedDataClassJsonMixin):
-    strategy: str = "auto"
-    ocr_languages: Optional[list[str]] = None
-    encoding: Optional[str] = None
-    additional_partition_args: Optional[dict[str, Any]] = None
-    skip_infer_table_types: Optional[list[str]] = None
-    fields_include: list[str] = field(
+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 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.",
     )
-    flatten_metadata: bool = False
-    metadata_exclude: list[str] = field(default_factory=list)
-    metadata_include: list[str] = field(default_factory=list)
-    partition_endpoint: Optional[str] = "https://api.unstructured.io/general/v0/general"
-    partition_by_api: bool = False
-    api_key: Optional[str] = enhanced_field(default=None, sensitive=True)
-    hi_res_model_name: Optional[str] = None
-
-    def __post_init__(self):
+    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.unstructured.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_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."
+    )
+
+    def model_post_init(self, __context: Any) -> None:
         if self.metadata_exclude and self.metadata_include:
             raise ValueError(
                 "metadata_exclude and metadata_include are "
@@ -93,16 +132,23 @@ class Partitioner(BaseProcess, ABC):
                 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[DataSourceMetadata] = None, **kwargs
+        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
 
         logger.debug(f"Using local partition with kwargs: {self.config.to_partition_kwargs()}")
-        logger.debug(f"partitioning file {filename} with metadata {metadata.to_dict()}")
+        logger.debug(f"partitioning file {filename} with metadata {metadata}")
         elements = partition(
             filename=str(filename.resolve()),
-            data_source_metadata=metadata,
+            data_source_metadata=FileDataSourceMetadata.from_dict(metadata),
             **self.config.to_partition_kwargs(),
         )
         return self.postprocess(elements=elements_to_dicts(elements))
@@ -138,29 +184,29 @@ class Partitioner(BaseProcess, ABC):
         partition_params = PartitionParameters(**filtered_partition_request)
         return partition_params
 
+    @requires_dependencies(dependencies=["unstructured_client"], extras="remote")
     async def partition_via_api(
-        self, filename: Path, metadata: Optional[DataSourceMetadata] = None, **kwargs
+        self, filename: Path, metadata: Optional[dict] = None, **kwargs
     ) -> list[dict]:
         from unstructured_client import UnstructuredClient
 
-        logger.debug(f"partitioning file {filename} with metadata: {metadata.to_dict()}")
+        logger.debug(f"partitioning file {filename} with metadata: {metadata}")
         client = UnstructuredClient(
-            server_url=self.config.partition_endpoint, api_key_auth=self.config.api_key
+            server_url=self.config.partition_endpoint,
+            api_key_auth=self.config.api_key.get_secret_value(),
         )
         partition_params = self.create_partition_parameters(filename=filename)
         resp = await self.call_api(client=client, request=partition_params)
         elements = resp.elements or []
         # Append the data source metadata the auto partition does for you
         for element in elements:
-            element["metadata"]["data_source"] = metadata.to_dict()
+            element["metadata"]["data_source"] = metadata
         return self.postprocess(elements=elements)
 
-    def run(
-        self, filename: Path, metadata: Optional[DataSourceMetadata] = None, **kwargs
-    ) -> list[dict]:
+    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[DataSourceMetadata] = None, **kwargs
+        self, filename: Path, metadata: Optional[dict] = None, **kwargs
     ) -> list[dict]:
         return await self.partition_via_api(filename, metadata=metadata, **kwargs)

unstructured_ingest/v2/processes/uncompress.py

@@ -4,14 +4,14 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
-from unstructured_ingest.enhanced_dataclass import EnhancedDataClassJsonMixin
+from pydantic import BaseModel
+
 from unstructured_ingest.utils.compression import TAR_FILE_EXT, ZIP_FILE_EXT, uncompress_file
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.interfaces.process import BaseProcess
 
 
-@dataclass
-class UncompressConfig(EnhancedDataClassJsonMixin):
+class UncompressConfig(BaseModel):
     pass
 
 

unstructured_ingest/v2/utils.py

@@ -0,0 +1,45 @@
+import json
+from datetime import datetime
+from inspect import isclass
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic.types import _SecretBase
+
+
+def is_secret(value: Any) -> bool:
+    # Case Secret[int]
+    if hasattr(value, "__origin__") and hasattr(value, "__args__"):
+        origin = value.__origin__
+        return isclass(origin) and issubclass(origin, _SecretBase)
+    # Case SecretStr
+    return isclass(value) and issubclass(value, _SecretBase)
+
+
+def serialize_base_model(model: BaseModel) -> dict:
+    # To get the full serialized dict regardless of if values are marked as Secret
+    model_dict = model.dict()
+    for k, v in model_dict.items():
+        if isinstance(v, _SecretBase):
+            secret_value = v.get_secret_value()
+            if isinstance(secret_value, BaseModel):
+                model_dict[k] = serialize_base_model(model=secret_value)
+            else:
+                model_dict[k] = secret_value
+
+    return model_dict
+
+
+def serialize_base_model_json(model: BaseModel, **json_kwargs) -> str:
+    model_dict = serialize_base_model(model=model)
+
+    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))
+
+    # Support json dumps kwargs such as sort_keys
+    return json.dumps(model_dict, default=json_serial, **json_kwargs)