unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/embed/vertexai.py

@@ -0,0 +1,126 @@
+# type: ignore
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, Any, Optional
+
+from pydantic import Field, Secret, ValidationError
+from pydantic.functional_validators import BeforeValidator
+
+from unstructured_ingest.embed.interfaces import (
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import UserAuthError
+
+if TYPE_CHECKING:
+    from vertexai.language_models import TextEmbeddingModel
+
+
+def conform_string_to_dict(value: Any) -> dict:
+    if isinstance(value, dict):
+        return value
+    if isinstance(value, str):
+        return json.loads(value)
+    raise ValidationError(f"Input could not be mapped to a valid dict: {value}")
+
+
+ApiKeyType = Secret[Annotated[dict, BeforeValidator(conform_string_to_dict)]]
+
+
+class VertexAIEmbeddingConfig(EmbeddingConfig):
+    api_key: ApiKeyType
+    embedder_model_name: Optional[str] = Field(
+        default="textembedding-gecko@001", alias="model_name"
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        from google.auth.exceptions import GoogleAuthError
+
+        if isinstance(e, GoogleAuthError):
+            return UserAuthError(e)
+        return e
+
+    def register_application_credentials(self):
+        # TODO look into passing credentials in directly, rather than via env var and tmp file
+        application_credentials_path = Path("/tmp") / "google-vertex-app-credentials.json"
+        with application_credentials_path.open("w+") as credentials_file:
+            json.dump(self.api_key.get_secret_value(), credentials_file)
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str(application_credentials_path)
+
+    @requires_dependencies(
+        ["vertexai"],
+        extras="embed-vertexai",
+    )
+    def get_client(self) -> "TextEmbeddingModel":
+        """Creates a VertexAI python client to embed elements."""
+        from vertexai.language_models import TextEmbeddingModel
+
+        self.register_application_credentials()
+        return TextEmbeddingModel.from_pretrained(self.embedder_model_name)
+
+
+@dataclass
+class VertexAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: VertexAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query):
+        return self._embed_documents(elements=[query])[0]
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = self._embed_documents([e.get("text", "") for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    @requires_dependencies(
+        ["vertexai"],
+        extras="embed-vertexai",
+    )
+    def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        from vertexai.language_models import TextEmbeddingInput
+
+        inputs = [TextEmbeddingInput(text=element) for element in elements]
+        try:
+            client = self.config.get_client()
+            embeddings = client.get_embeddings(inputs)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return [e.values for e in embeddings]
+
+
+@dataclass
+class AsyncVertexAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: VertexAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query):
+        embedding = await self._embed_documents(elements=[query])
+        return embedding[0]
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = await self._embed_documents([e.get("text", "") for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    @requires_dependencies(
+        ["vertexai"],
+        extras="embed-vertexai",
+    )
+    async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        from vertexai.language_models import TextEmbeddingInput
+
+        inputs = [TextEmbeddingInput(text=element) for element in elements]
+        try:
+            client = self.config.get_client()
+            embeddings = await client.get_embeddings_async(inputs)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return [e.values for e in embeddings]

unstructured_ingest/embed/voyageai.py

@@ -0,0 +1,130 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import Field, SecretStr
+
+from unstructured_ingest.embed.interfaces import (
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.errors import (
+    ProviderError,
+    UserAuthError,
+    UserError,
+)
+from unstructured_ingest.v2.errors import (
+    RateLimitError as CustomRateLimitError,
+)
+
+if TYPE_CHECKING:
+    from voyageai import AsyncClient as AsyncVoyageAIClient
+    from voyageai import Client as VoyageAIClient
+
+
+class VoyageAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    embedder_model_name: str = Field(default="voyage-3", alias="model_name")
+    batch_size: Optional[int] = Field(default=None)
+    truncation: Optional[bool] = Field(default=None)
+    max_retries: int = 0
+    timeout_in_seconds: Optional[int] = None
+
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://docs.voyageai.com/docs/error-codes
+        from voyageai.error import AuthenticationError, RateLimitError, VoyageError
+
+        if not isinstance(e, VoyageError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        http_code = e.http_status
+        message = e.user_message
+        if isinstance(e, AuthenticationError):
+            return UserAuthError(message)
+        if isinstance(e, RateLimitError):
+            return CustomRateLimitError(message)
+        if 400 <= http_code < 500:
+            return UserError(message)
+        if http_code >= 500:
+            return ProviderError(message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
+    @requires_dependencies(
+        ["voyageai"],
+        extras="embed-voyageai",
+    )
+    def get_client(self) -> "VoyageAIClient":
+        """Creates a VoyageAI python client to embed elements."""
+        from voyageai import Client as VoyageAIClient
+
+        client = VoyageAIClient(
+            api_key=self.api_key.get_secret_value(),
+            max_retries=self.max_retries,
+            timeout=self.timeout_in_seconds,
+        )
+        return client
+
+    @requires_dependencies(
+        ["voyageai"],
+        extras="embed-voyageai",
+    )
+    def get_async_client(self) -> "AsyncVoyageAIClient":
+        """Creates a VoyageAI python client to embed elements."""
+        from voyageai import AsyncClient as AsyncVoyageAIClient
+
+        client = AsyncVoyageAIClient(
+            api_key=self.api_key.get_secret_value(),
+            max_retries=self.max_retries,
+            timeout=self.timeout_in_seconds,
+        )
+        return client
+
+
+@dataclass
+class VoyageAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: VoyageAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        client: VoyageAIClient = self.config.get_client()
+        try:
+            response = client.embed(texts=elements, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.embeddings
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = self._embed_documents([e.get("text", "") for e in elements])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    def embed_query(self, query: str) -> list[float]:
+        return self._embed_documents(elements=[query])[0]
+
+
+@dataclass
+class AsyncVoyageAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: VoyageAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        client = self.config.get_async_client()
+        try:
+            response = await client.embed(texts=elements, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.embeddings
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        embeddings = await self._embed_documents([e.get("text", "") for e in elements])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    async def embed_query(self, query: str) -> list[float]:
+        embedding = await self._embed_documents(elements=[query])
+        return embedding[0]

unstructured_ingest/enhanced_dataclass/__init__.py

@@ -0,0 +1,4 @@
+from .dataclasses import enhanced_field
+from .json_mixin import EnhancedDataClassJsonMixin
+
+__all__ = ["enhanced_field", "EnhancedDataClassJsonMixin"]

unstructured_ingest/enhanced_dataclass/core.py

@@ -0,0 +1,99 @@
+import _thread
+import copy
+import functools
+from dataclasses import fields
+
+from dataclasses_json.core import (
+    Collection,
+    Enum,
+    Mapping,
+    _encode_overrides,
+    _handle_undefined_parameters_safe,
+    _user_overrides_or_exts,
+    is_dataclass,
+)
+
+
+def _recursive_repr(user_function):
+    # Copied from dataclasses as this method isn't exposed for importing
+    repr_running = set()
+
+    @functools.wraps(user_function)
+    def wrapper(self):
+        key = id(self), _thread.get_ident()
+        if key in repr_running:
+            return "..."
+        repr_running.add(key)
+        try:
+            result = user_function(self)
+        finally:
+            repr_running.discard(key)
+        return result
+
+    return wrapper
+
+
+def _asdict(
+    obj,
+    encode_json=False,
+    redact_sensitive=False,
+    redacted_text="***REDACTED***",
+    apply_name_overload: bool = True,
+):
+    """
+    A re-implementation of `asdict` (based on the original in the `dataclasses`
+    source) to support arbitrary Collection and Mapping types.
+    """
+    if is_dataclass(obj):
+        result = []
+        overrides = _user_overrides_or_exts(obj)
+        for field in fields(obj):
+            if overrides[field.name].encoder:
+                value = getattr(obj, field.name)
+            else:
+                value = _asdict(
+                    getattr(obj, field.name),
+                    encode_json=encode_json,
+                    redact_sensitive=redact_sensitive,
+                    redacted_text=redacted_text,
+                    apply_name_overload=apply_name_overload,
+                )
+            if getattr(field, "sensitive", False) and redact_sensitive and value:
+                value = redacted_text
+            if getattr(field, "overload_name", None) and apply_name_overload:
+                overload_name = getattr(field, "overload_name")
+                result.append((overload_name, value))
+            else:
+                result.append((field.name, value))
+
+        result = _handle_undefined_parameters_safe(cls=obj, kvs=dict(result), usage="to")
+        return _encode_overrides(
+            dict(result), _user_overrides_or_exts(obj), encode_json=encode_json
+        )
+    elif isinstance(obj, Mapping):
+        return {
+            _asdict(
+                k,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            ): _asdict(
+                v,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            )
+            for k, v in obj.items()
+        }
+    elif isinstance(obj, Collection) and not isinstance(obj, (str, bytes, Enum)):
+        return [
+            _asdict(
+                v,
+                encode_json=encode_json,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+            )
+            for v in obj
+        ]
+    else:
+        return copy.deepcopy(obj)

unstructured_ingest/enhanced_dataclass/dataclasses.py

@@ -0,0 +1,54 @@
+import typing as t
+from dataclasses import MISSING, Field
+
+from unstructured_ingest.enhanced_dataclass.core import _recursive_repr
+
+
+class EnhancedField(Field):
+    def __init__(self, *args, sensitive=False, overload_name: t.Optional[str] = None):
+        super().__init__(*args)
+        self.sensitive = sensitive
+        self.overload_name = overload_name
+
+    @_recursive_repr
+    def __repr__(self):
+        # Support for kw_only added in 3.10, to support as low as 3.8, need to dynamically map
+        fields_array = [
+            f"name={self.name!r}",
+            f"type={self.type!r}",
+            f"default={self.default!r}",
+            f"default_factory={self.default_factory!r}",
+            f"init={self.init!r}",
+            f"repr={self.repr!r}",
+            f"hash={self.hash!r}",
+            f"compare={self.compare!r}",
+            f"metadata={self.metadata!r}",
+            f"sensitive={self.sensitive!r}",
+            f"overload_name={self.overload_name!r}",
+            f"_field_type={self._field_type}",
+        ]
+        if kw_only := getattr(self, "kw_only", None):
+            fields_array.append(f"kw_only={kw_only!r}")
+        return "Field({})".format(",".join(fields_array))
+
+
+def enhanced_field(
+    *,
+    default=MISSING,
+    default_factory=MISSING,
+    init: bool = True,
+    repr: bool = True,
+    hash=None,
+    compare: bool = True,
+    metadata=None,
+    kw_only=MISSING,
+    sensitive: bool = False,
+    overload_name: t.Optional[str] = None,
+):
+    if default is not MISSING and default_factory is not MISSING:
+        raise ValueError("cannot specify both default and default_factory")
+    args = [default, default_factory, init, repr, hash, compare, metadata]
+    # Support for kw_only added in 3.10, to support as low as 3.8, need to dynamically map
+    if "kw_only" in EnhancedField.__slots__:
+        args.append(kw_only)
+    return EnhancedField(*args, sensitive=sensitive, overload_name=overload_name)

unstructured_ingest/enhanced_dataclass/json_mixin.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import json
+from dataclasses import InitVar, fields
+from typing import Any, Callable, Optional, Type, TypeVar, Union
+
+import dataclasses_json.core as dataclasses_json_core
+from dataclasses_json import DataClassJsonMixin
+
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+
+A = TypeVar("A", bound="EnhancedDataClassJsonMixin")
+
+# Monkey-patch _decode_dataclass class to support name override
+og_decode_dataclass = dataclasses_json_core._decode_dataclass
+
+
+def custom_decode_dataclass(cls, kvs, infer_missing):
+    dataclass_fields = fields(cls)
+    for f in [
+        field
+        for field in dataclass_fields
+        if hasattr(field, "overload_name") and getattr(field, "overload_name", None)
+    ]:
+        field_name = f.name
+        overload_name = getattr(f, "overload_name")
+        if isinstance(kvs, dict) and overload_name in kvs:
+            kvs[field_name] = kvs.pop(overload_name)
+    return og_decode_dataclass(cls, kvs, infer_missing)
+
+
+dataclasses_json_core._decode_dataclass = custom_decode_dataclass
+
+
+class EnhancedDataClassJsonMixin(DataClassJsonMixin):
+    """A mixin class extending DataClassJsonMixin.
+
+    This class extends the functionality of DataClassJsonMixin to provide enhanced functionality
+    for JSON serialization and deserialization. It introduces options for redacting sensitive
+    information, custom encoding, and more advanced schema handling.
+
+    Attributes:
+        N/A (No additional attributes)
+
+    Methods:
+        to_json: Serialize the object to JSON format with customizable options.
+        from_dict: Deserialize a dictionary into an object of this class.
+        to_dict: Convert the object to a dictionary with customizable options.
+        schema: Generate a schema for validating and parsing JSON data based on this class.
+    """
+
+    @classmethod
+    def check_init_var(cls):
+        ann = cls.__dict__.get("__annotations__", {})
+        init_vars = {k: v for k, v in ann.items() if isinstance(v, InitVar)}
+        if init_vars:
+            raise TypeError(
+                "Class {} has the following fields defined with an InitVar which "
+                "cannot be used with EnhancedDataClassJsonMixin: {}".format(
+                    cls.__name__, ", ".join(init_vars.keys())
+                )
+            )
+
+    def to_json(
+        self,
+        *,
+        skipkeys: bool = False,
+        ensure_ascii: bool = True,
+        check_circular: bool = True,
+        allow_nan: bool = True,
+        indent: Optional[Union[int, str]] = None,
+        separators: Optional[tuple[str, str]] = None,
+        default: Optional[Callable[..., Any]] = None,
+        sort_keys: bool = False,
+        redact_sensitive: bool = False,
+        redacted_text: str = "***REDACTED***",
+        apply_name_overload: bool = True,
+        **kw: Any,
+    ) -> str:
+        self.check_init_var()
+        return json.dumps(
+            self.to_dict(
+                encode_json=False,
+                redact_sensitive=redact_sensitive,
+                redacted_text=redacted_text,
+                apply_name_overload=apply_name_overload,
+            ),
+            cls=dataclasses_json_core._ExtendedEncoder,
+            skipkeys=skipkeys,
+            ensure_ascii=ensure_ascii,
+            check_circular=check_circular,
+            allow_nan=allow_nan,
+            indent=indent,
+            separators=separators,
+            default=default,
+            sort_keys=sort_keys,
+            **kw,
+        )
+
+    @classmethod
+    def from_dict(
+        cls: Type[A],
+        kvs: dataclasses_json_core.Json,
+        *,
+        infer_missing=False,
+        apply_name_overload=False,
+    ) -> A:
+        cls.check_init_var()
+        return dataclasses_json_core._decode_dataclass(cls, kvs, infer_missing)
+
+    def to_dict(
+        self,
+        encode_json: bool = False,
+        redact_sensitive: bool = False,
+        redacted_text: str = "***REDACTED***",
+        apply_name_overload: bool = True,
+    ) -> dict[str, dataclasses_json_core.Json]:
+        self.check_init_var()
+        return _asdict(
+            self,
+            encode_json=encode_json,
+            redact_sensitive=redact_sensitive,
+            redacted_text=redacted_text,
+            apply_name_overload=apply_name_overload,
+        )

unstructured_ingest/error.py

@@ -0,0 +1,49 @@
+from abc import ABC
+from functools import wraps
+
+
+class CustomError(Exception, ABC):
+    error_string: str
+
+    @classmethod
+    def wrap(cls, f):
+        """
+        Provides a wrapper for a function that catches any exception and
+        re-raises it as the customer error. If the exception itself is already an instance
+        of the custom error, re-raises original error.
+        """
+
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            try:
+                return f(*args, **kwargs)
+            except BaseException as error:
+                if not isinstance(error, cls) and not issubclass(type(error), cls):
+                    raise cls(cls.error_string.format(str(error))) from error
+                raise
+
+        return wrapper
+
+
+class SourceConnectionError(CustomError):
+    error_string = "Error in getting data from upstream data source: {}"
+
+
+class SourceConnectionNetworkError(SourceConnectionError):
+    error_string = "Error in connecting to upstream data source: {}"
+
+
+class DestinationConnectionError(CustomError):
+    error_string = "Error in connecting to downstream data source: {}"
+
+
+class EmbeddingEncoderConnectionError(CustomError):
+    error_string = "Error in connecting to the embedding model provider: {}"
+
+
+class WriteError(CustomError):
+    error_string = "Error in writing to downstream data source: {}"
+
+
+class PartitionError(CustomError):
+    error_string = "Error in partitioning content: {}"

unstructured_ingest/ingest_backoff/__init__.py

@@ -0,0 +1,3 @@
+from ._wrapper import RetryHandler
+
+__all__ = ["RetryHandler"]

unstructured_ingest/ingest_backoff/_common.py

@@ -0,0 +1,102 @@
+import logging
+import sys
+import traceback
+
+
+# Default startup handler
+def _log_start(details, logger, log_level):
+    max_tried = details.get("max_tries")
+    max_time = details.get("max_time")
+    if max_tried is not None and max_time is not None:
+        s = "%.1fs or %d tries"
+        s_args = [max_time, max_tried]
+    elif max_tried is not None:
+        s = "%d tries"
+        s_args = [max_tried]
+    else:
+        s = "%.1fs"
+        s_args = [max_time]
+    exception = details.get("exception")
+    if isinstance(exception, tuple):
+        exception = list(exception)
+    elif not isinstance(exception, list):
+        exception = [exception]
+    exception_s = ", ".join([e.__name__ for e in exception])
+    if log_level >= logging.INFO:
+        msg = f"Attempting %s(...), will retry for {s} given these issues: %s"
+        log_args = [details["target"].__name__] + s_args + [exception_s]
+    else:
+        msg = f"Attempting %s(%s), will retry for {s} given these issues: %s"
+        target_input_list = []
+        if args := details.get("args"):
+            target_input_list.extend([str(d) for d in args])
+        if kwargs := details.get("kwargs"):
+            target_input_list.extend([f"{k}={str(v)}" for k, v in kwargs.items()])
+        target_input = ", ".join(target_input_list) if target_input_list else ""
+        log_args = (
+            [
+                details["target"].__name__,
+                target_input,
+            ]
+            + s_args
+            + [exception_s]
+        )
+    logger.log(log_level, msg, *log_args)
+
+
+# Default backoff handler
+def _log_backoff(details, logger, log_level):
+    if log_level >= logging.INFO:
+        msg = "Backing off %s(...) for %.1fs (%s)"
+        log_args = [details["target"].__name__, details["tries"]]
+    else:
+        msg = "Backing off %.1fs seconds after %d tries calling function %s(%s) -> %s"
+        target_input_list = []
+        if args := details.get("args"):
+            target_input_list.extend([str(d) for d in args])
+        if kwargs := details.get("kwargs"):
+            target_input_list.extend([f"{k}={str(v)}" for k, v in kwargs.items()])
+        target_input = ", ".join(target_input_list) if target_input_list else ""
+        log_args = [
+            details["wait"],
+            details["tries"],
+            details["target"].__name__,
+            target_input,
+        ]
+    exc_typ, exc, _ = sys.exc_info()
+    if exc is not None:
+        exc_fmt = traceback.format_exception_only(exc_typ, exc)[-1]
+        log_args.append(exc_fmt.rstrip("\n"))
+    else:
+        log_args.append(str(details["value"]))
+    logger.log(log_level, msg, *log_args)
+
+
+# Default giveup handler
+def _log_giveup(details, logger, log_level):
+    if log_level >= logging.INFO:
+        msg = "Giving up %s(...) after %.1fs (%s)"
+        log_args = [details["target"].__name__, details["tries"]]
+    else:
+        msg = "Giving up after %d tries (%.1fs) calling function %s(%s) -> %s"
+        target_input_list = []
+        if args := details.get("args"):
+            target_input_list.extend([str(d) for d in args])
+        if kwargs := details.get("kwargs"):
+            target_input_list.extend([f"{k}={str(v)}" for k, v in kwargs.items()])
+        target_input = ", ".join(target_input_list) if target_input_list else "..."
+        log_args = [
+            details["tries"],
+            details["wait"],
+            details["target"].__name__,
+            target_input,
+        ]
+
+    exc_typ, exc, _ = sys.exc_info()
+    if exc is not None:
+        exc_fmt = traceback.format_exception_only(exc_typ, exc)[-1]
+        log_args.append(exc_fmt.rstrip("\n"))
+    else:
+        log_args.append(details["value"])
+
+    logger.log(log_level, msg, *log_args)