unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/embed/interfaces.py

@@ -0,0 +1,117 @@
+import asyncio
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import numpy as np
+from pydantic import BaseModel
+
+
+class EmbeddingConfig(BaseModel):
+    pass
+
+
+@dataclass
+class BaseEncoder(ABC):
+    config: EmbeddingConfig
+
+    def initialize(self):
+        """Initializes the embedding encoder class. Should also validate the instance
+        is properly configured: e.g., embed a single a element"""
+
+    def wrap_error(self, e: Exception) -> Exception:
+        """Handle errors from the embedding service. Should raise a more informative error
+        if possible"""
+        return e
+
+    @staticmethod
+    def _add_embeddings_to_elements(
+        elements: list[dict], embeddings: list[list[float]]
+    ) -> list[dict]:
+        """
+        Add embeddings to elements.
+
+        Args:
+            elements (list[Element]): List of elements.
+            embeddings (list[list[float]]): List of embeddings.
+
+        Returns:
+            list[Element]: Elements with embeddings added.
+        """
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element["embeddings"] = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements
+
+
+@dataclass
+class BaseEmbeddingEncoder(BaseEncoder, ABC):
+
+    def initialize(self):
+        """Initializes the embedding encoder class. Should also validate the instance
+        is properly configured: e.g., embed a single a element"""
+
+    @property
+    def num_of_dimensions(self) -> tuple[int, ...]:
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def get_exemplary_embedding(self) -> list[float]:
+        return self.embed_query(query="Q")
+
+    @property
+    def is_unit_vector(self) -> bool:
+        """Denotes if the embedding vector is a unit vector."""
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    @abstractmethod
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        pass
+
+    @abstractmethod
+    def embed_query(self, query: str) -> list[float]:
+        pass
+
+    def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        results = []
+        for text in elements:
+            response = self.embed_query(query=text)
+            results.append(response)
+
+        return results
+
+
+@dataclass
+class AsyncBaseEmbeddingEncoder(BaseEncoder, ABC):
+
+    async def initialize(self):
+        """Initializes the embedding encoder class. Should also validate the instance
+        is properly configured: e.g., embed a single a element"""
+
+    @property
+    async def num_of_dimensions(self) -> tuple[int, ...]:
+        exemplary_embedding = await self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    async def get_exemplary_embedding(self) -> list[float]:
+        return await self.embed_query(query="Q")
+
+    @property
+    async def is_unit_vector(self) -> bool:
+        """Denotes if the embedding vector is a unit vector."""
+        exemplary_embedding = await self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    @abstractmethod
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        pass
+
+    @abstractmethod
+    async def embed_query(self, query: str) -> list[float]:
+        pass
+
+    async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        results = await asyncio.gather(*[self.embed_query(query=text) for text in elements])
+        return results

unstructured_ingest/embed/mixedbreadai.py

@@ -0,0 +1,233 @@
+import asyncio
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field, SecretStr
+
+from unstructured_ingest.embed.interfaces import (
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+USER_AGENT = "@mixedbread-ai/unstructured"
+BATCH_SIZE = 128
+TIMEOUT = 60
+MAX_RETRIES = 3
+ENCODING_FORMAT = "float"
+TRUNCATION_STRATEGY = "end"
+
+
+if TYPE_CHECKING:
+    from mixedbread_ai.client import AsyncMixedbreadAI, MixedbreadAI
+    from mixedbread_ai.core import RequestOptions
+
+
+class MixedbreadAIEmbeddingConfig(EmbeddingConfig):
+    """
+    Configuration class for Mixedbread AI Embedding Encoder.
+
+    Attributes:
+        api_key (str): API key for accessing Mixedbread AI..
+        embedder_model_name (str): Name of the model to use for embeddings.
+    """
+
+    api_key: SecretStr = Field(
+        default_factory=lambda: SecretStr(os.environ.get("MXBAI_API_KEY")),
+    )
+
+    embedder_model_name: str = Field(
+        default="mixedbread-ai/mxbai-embed-large-v1", alias="model_name"
+    )
+
+    @requires_dependencies(
+        ["mixedbread_ai"],
+        extras="embed-mixedbreadai",
+    )
+    def get_client(self) -> "MixedbreadAI":
+        """
+        Create the Mixedbread AI client.
+
+        Returns:
+            MixedbreadAI: Initialized client.
+        """
+        from mixedbread_ai.client import MixedbreadAI
+
+        return MixedbreadAI(
+            api_key=self.api_key.get_secret_value(),
+        )
+
+    @requires_dependencies(
+        ["mixedbread_ai"],
+        extras="embed-mixedbreadai",
+    )
+    def get_async_client(self) -> "AsyncMixedbreadAI":
+        from mixedbread_ai.client import AsyncMixedbreadAI
+
+        return AsyncMixedbreadAI(
+            api_key=self.api_key.get_secret_value(),
+        )
+
+
+@dataclass
+class MixedbreadAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    """
+    Embedding encoder for Mixedbread AI.
+
+    Attributes:
+        config (MixedbreadAIEmbeddingConfig): Configuration for the embedding encoder.
+    """
+
+    config: MixedbreadAIEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> list[float]:
+        """Get an exemplary embedding to determine dimensions and unit vector status."""
+        return self._embed(["Q"])[0]
+
+    @requires_dependencies(
+        ["mixedbread_ai"],
+        extras="embed-mixedbreadai",
+    )
+    def get_request_options(self) -> "RequestOptions":
+        from mixedbread_ai.core import RequestOptions
+
+        return RequestOptions(
+            max_retries=MAX_RETRIES,
+            timeout_in_seconds=TIMEOUT,
+            additional_headers={"User-Agent": USER_AGENT},
+        )
+
+    def _embed(self, texts: list[str]) -> list[list[float]]:
+        """
+        Embed a list of texts using the Mixedbread AI API.
+
+        Args:
+            texts (list[str]): List of texts to embed.
+
+        Returns:
+            list[list[float]]: List of embeddings.
+        """
+        batch_size = BATCH_SIZE
+        batch_itr = range(0, len(texts), batch_size)
+
+        responses = []
+        client = self.config.get_client()
+        for i in batch_itr:
+            batch = texts[i : i + batch_size]
+            response = client.embeddings(
+                model=self.config.embedder_model_name,
+                normalized=True,
+                encoding_format=ENCODING_FORMAT,
+                truncation_strategy=TRUNCATION_STRATEGY,
+                request_options=self.get_request_options(),
+                input=batch,
+            )
+            responses.append(response)
+        return [item.embedding for response in responses for item in response.data]
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        """
+        Embed a list of document elements.
+
+        Args:
+            elements (list[Element]): List of document elements.
+
+        Returns:
+            list[Element]: Elements with embeddings.
+        """
+        embeddings = self._embed([e.get("text", "") for e in elements])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    def embed_query(self, query: str) -> list[float]:
+        """
+        Embed a query string.
+
+        Args:
+            query (str): Query string to embed.
+
+        Returns:
+            list[float]: Embedding of the query.
+        """
+        return self._embed([query])[0]
+
+
+@dataclass
+class AsyncMixedbreadAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+
+    config: MixedbreadAIEmbeddingConfig
+
+    async def get_exemplary_embedding(self) -> list[float]:
+        """Get an exemplary embedding to determine dimensions and unit vector status."""
+        embedding = await self._embed(["Q"])
+        return embedding[0]
+
+    @requires_dependencies(
+        ["mixedbread_ai"],
+        extras="embed-mixedbreadai",
+    )
+    def get_request_options(self) -> "RequestOptions":
+        from mixedbread_ai.core import RequestOptions
+
+        return RequestOptions(
+            max_retries=MAX_RETRIES,
+            timeout_in_seconds=TIMEOUT,
+            additional_headers={"User-Agent": USER_AGENT},
+        )
+
+    async def _embed(self, texts: list[str]) -> list[list[float]]:
+        """
+        Embed a list of texts using the Mixedbread AI API.
+
+        Args:
+            texts (list[str]): List of texts to embed.
+
+        Returns:
+            list[list[float]]: List of embeddings.
+        """
+        batch_size = BATCH_SIZE
+        batch_itr = range(0, len(texts), batch_size)
+
+        client = self.config.get_async_client()
+        tasks = []
+        for i in batch_itr:
+            batch = texts[i : i + batch_size]
+            tasks.append(
+                client.embeddings(
+                    model=self.config.embedder_model_name,
+                    normalized=True,
+                    encoding_format=ENCODING_FORMAT,
+                    truncation_strategy=TRUNCATION_STRATEGY,
+                    request_options=self.get_request_options(),
+                    input=batch,
+                )
+            )
+        responses = await asyncio.gather(*tasks)
+        return [item.embedding for response in responses for item in response.data]
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        """
+        Embed a list of document elements.
+
+        Args:
+            elements (list[Element]): List of document elements.
+
+        Returns:
+            list[Element]: Elements with embeddings.
+        """
+        embeddings = await self._embed([e.get("text", "") for e in elements])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    async def embed_query(self, query: str) -> list[float]:
+        """
+        Embed a query string.
+
+        Args:
+            query (str): Query string to embed.
+
+        Returns:
+            list[float]: Embedding of the query.
+        """
+        embedding = await self._embed([query])
+        return embedding[0]

unstructured_ingest/embed/octoai.py

@@ -0,0 +1,130 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+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,
+    QuotaError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+)
+
+if TYPE_CHECKING:
+    from openai import AsyncOpenAI, OpenAI
+
+
+class OctoAiEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    embedder_model_name: str = Field(default="thenlper/gte-large", alias="model_name")
+    base_url: str = Field(default="https://text.octoai.run/v1")
+
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://platform.openai.com/docs/guides/error-codes/api-errors
+        from openai import APIStatusError
+
+        if not isinstance(e, APIStatusError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        error_code = e.code
+        if 400 <= e.status_code < 500:
+            # user error
+            if e.status_code == 401:
+                return UserAuthError(e.message)
+            if e.status_code == 429:
+                # 429 indicates rate limit exceeded and quote exceeded
+                if error_code == "insufficient_quota":
+                    return QuotaError(e.message)
+                else:
+                    return RateLimitError(e.message)
+            return UserError(e.message)
+        if e.status_code >= 500:
+            return ProviderError(e.message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
+    @requires_dependencies(
+        ["openai", "tiktoken"],
+        extras="embed-octoai",
+    )
+    def get_client(self) -> "OpenAI":
+        """Creates an OpenAI python client to embed elements. Uses the OpenAI SDK."""
+        from openai import OpenAI
+
+        return OpenAI(api_key=self.api_key.get_secret_value(), base_url=self.base_url)
+
+    @requires_dependencies(
+        ["openai", "tiktoken"],
+        extras="embed-octoai",
+    )
+    def get_async_client(self) -> "AsyncOpenAI":
+        """Creates an OpenAI python client to embed elements. Uses the OpenAI SDK."""
+        from openai import AsyncOpenAI
+
+        return AsyncOpenAI(api_key=self.api_key.get_secret_value(), base_url=self.base_url)
+
+
+@dataclass
+class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OctoAiEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query: str):
+        try:
+            client = self.config.get_client()
+            response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.data[0].embedding
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        texts = [e.get("text", "") for e in elements]
+        try:
+            client = self.config.get_client()
+            response = client.embeddings.create(input=texts, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        embeddings = [data.embedding for data in response.data]
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+
+@dataclass
+class AsyncOctoAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: OctoAiEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query: str):
+        client = self.config.get_async_client()
+        try:
+            response = await client.embeddings.create(
+                input=query, model=self.config.embedder_model_name
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.data[0].embedding
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        texts = [e.get("text", "") for e in elements]
+        client = self.config.get_async_client()
+        try:
+            response = await client.embeddings.create(
+                input=texts, model=self.config.embedder_model_name
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        embeddings = [data.embedding for data in response.data]
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings

unstructured_ingest/embed/openai.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+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,
+    QuotaError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+)
+
+if TYPE_CHECKING:
+    from openai import AsyncOpenAI, OpenAI
+
+
+class OpenAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    embedder_model_name: str = Field(default="text-embedding-ada-002", alias="model_name")
+
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://platform.openai.com/docs/guides/error-codes/api-errors
+        from openai import APIStatusError
+
+        if not isinstance(e, APIStatusError):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            raise e
+        error_code = e.code
+        if 400 <= e.status_code < 500:
+            # user error
+            if e.status_code == 401:
+                return UserAuthError(e.message)
+            if e.status_code == 429:
+                # 429 indicates rate limit exceeded and quote exceeded
+                if error_code == "insufficient_quota":
+                    return QuotaError(e.message)
+                else:
+                    return RateLimitError(e.message)
+            return UserError(e.message)
+        if e.status_code >= 500:
+            return ProviderError(e.message)
+        logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+        return e
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_client(self) -> "OpenAI":
+        from openai import OpenAI
+
+        return OpenAI(api_key=self.api_key.get_secret_value())
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_async_client(self) -> "AsyncOpenAI":
+        from openai import AsyncOpenAI
+
+        return AsyncOpenAI(api_key=self.api_key.get_secret_value())
+
+
+@dataclass
+class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OpenAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query: str) -> list[float]:
+
+        client = self.config.get_client()
+        try:
+            response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.data[0].embedding
+
+    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
+
+
+@dataclass
+class AsyncOpenAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: OpenAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query: str) -> list[float]:
+        client = self.config.get_async_client()
+        try:
+            response = await client.embeddings.create(
+                input=query, model=self.config.embedder_model_name
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return response.data[0].embedding
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        client = self.config.get_async_client()
+        texts = [e.get("text", "") for e in elements]
+        try:
+            response = await client.embeddings.create(
+                input=texts, model=self.config.embedder_model_name
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        embeddings = [data.embedding for data in response.data]
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings

unstructured_ingest/embed/togetherai.py

@@ -0,0 +1,106 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+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 (
+    RateLimitError as CustomRateLimitError,
+)
+from unstructured_ingest.v2.errors import (
+    UserAuthError,
+    UserError,
+)
+
+if TYPE_CHECKING:
+    from together import AsyncTogether, Together
+
+
+class TogetherAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    embedder_model_name: str = Field(
+        default="togethercomputer/m2-bert-80M-8k-retrieval", alias="model_name"
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        # https://docs.together.ai/docs/error-codes
+        from together.error import AuthenticationError, RateLimitError, TogetherException
+
+        if not isinstance(e, TogetherException):
+            logger.error(f"unhandled exception from openai: {e}", exc_info=True)
+            return e
+        message = e.args[0]
+        if isinstance(e, AuthenticationError):
+            return UserAuthError(message)
+        if isinstance(e, RateLimitError):
+            return CustomRateLimitError(message)
+        return UserError(message)
+
+    @requires_dependencies(["together"], extras="togetherai")
+    def get_client(self) -> "Together":
+        from together import Together
+
+        return Together(api_key=self.api_key.get_secret_value())
+
+    @requires_dependencies(["together"], extras="togetherai")
+    def get_async_client(self) -> "AsyncTogether":
+        from together import AsyncTogether
+
+        return AsyncTogether(api_key=self.api_key.get_secret_value())
+
+
+@dataclass
+class TogetherAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: TogetherAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query: str) -> list[float]:
+        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])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        client = self.config.get_client()
+        try:
+            outputs = client.embeddings.create(
+                model=self.config.embedder_model_name, input=elements
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return [outputs.data[i].embedding for i in range(len(elements))]
+
+
+@dataclass
+class AsyncTogetherAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: TogetherAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query: str) -> list[float]:
+        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])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    async def _embed_documents(self, elements: list[str]) -> list[list[float]]:
+        client = self.config.get_async_client()
+        try:
+            outputs = await client.embeddings.create(
+                model=self.config.embedder_model_name, input=elements
+            )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        return [outputs.data[i].embedding for i in range(len(elements))]