unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/embed/bedrock.py

@@ -0,0 +1,323 @@
+import asyncio
+import json
+import os
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, AsyncIterable
+
+from pydantic import Field, SecretStr
+
+from unstructured_ingest.embed.interfaces import (
+    EMBEDDINGS_KEY,
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.error import (
+    ProviderError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+    is_internal_error,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+    class BedrockRuntimeClient(BaseClient):
+        def invoke_model(
+            self,
+            body: str,
+            modelId: str,
+            accept: str,
+            contentType: str,
+            inferenceProfileId: str = None,
+        ) -> dict:
+            pass
+
+    class AsyncBedrockRuntimeClient(BaseClient):
+        async def invoke_model(
+            self,
+            body: str,
+            modelId: str,
+            accept: str,
+            contentType: str,
+            inferenceProfileId: str = None,
+        ) -> dict:
+            pass
+
+    class BedrockClient(BaseClient):
+        def list_foundation_models(self, byOutputModality: str) -> dict:
+            pass
+
+
+def conform_query(query: str, provider: str) -> dict:
+    # replace newlines, which can negatively affect performance.
+    text = query.replace(os.linesep, " ")
+
+    # format input body for provider
+    input_body = {}
+    if provider == "cohere":
+        if "input_type" not in input_body:
+            input_body["input_type"] = "search_document"
+        input_body["texts"] = [text]
+    else:
+        # includes common provider == "amazon"
+        input_body["inputText"] = text
+    return input_body
+
+
+class BedrockEmbeddingConfig(EmbeddingConfig):
+    aws_access_key_id: SecretStr | None = Field(description="aws access key id", default=None)
+    aws_secret_access_key: SecretStr | None = Field(
+        description="aws secret access key", default=None
+    )
+    region_name: str = Field(
+        description="aws region name",
+        default_factory=lambda: (
+            os.getenv("BEDROCK_REGION_NAME") or
+            os.getenv("AWS_DEFAULT_REGION") or
+            "us-west-2"
+        )
+    )
+    endpoint_url: str | None = Field(description="custom bedrock endpoint url", default=None)
+    access_method: str = Field(
+        description="authentication method", default="credentials"
+    )  # "credentials" or "iam"
+    embedder_model_name: str = Field(
+        default="amazon.titan-embed-text-v1",
+        alias="model_name",
+        description="AWS Bedrock model name",
+    )
+    inference_profile_id: str | None = Field(
+        description="AWS Bedrock inference profile ID",
+        default_factory=lambda: os.getenv("BEDROCK_INFERENCE_PROFILE_ID"),
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        from botocore.exceptions import ClientError
+
+        if isinstance(e, ClientError):
+            # https://docs.aws.amazon.com/awssupport/latest/APIReference/CommonErrors.html
+            http_response = e.response
+            meta = http_response["ResponseMetadata"]
+            http_response_code = meta["HTTPStatusCode"]
+            error_code = http_response["Error"]["Code"]
+            if http_response_code == 400:
+                if error_code == "ValidationError":
+                    return UserError(http_response["Error"])
+                elif error_code == "ThrottlingException":
+                    return RateLimitError(http_response["Error"])
+                elif error_code == "NotAuthorized" or error_code == "AccessDeniedException":
+                    return UserAuthError(http_response["Error"])
+            if http_response_code == 403:
+                return UserAuthError(http_response["Error"])
+            if 400 <= http_response_code < 500:
+                return UserError(http_response["Error"])
+            if http_response_code >= 500:
+                return ProviderError(http_response["Error"])
+
+        logger.error(f"unhandled exception from bedrock: {e}", exc_info=True)
+        return e
+
+    def run_precheck(self) -> None:
+        # Validate access method and credentials configuration
+        if self.access_method == "credentials":
+            if not (self.aws_access_key_id and self.aws_secret_access_key):
+                raise ValueError(
+                    "Credentials access method requires aws_access_key_id and aws_secret_access_key"
+                )
+        elif self.access_method == "iam":
+            # For IAM, credentials are handled by AWS SDK
+            pass
+        else:
+            raise ValueError(
+                f"Invalid access_method: {self.access_method}. Must be 'credentials' or 'iam'"
+            )
+            
+        client = self.get_bedrock_client()
+        try:
+            model_info = client.list_foundation_models(byOutputModality="EMBEDDING")
+            summaries = model_info.get("modelSummaries", [])
+            model_ids = [m["modelId"] for m in summaries]
+            arns = [":".join(m["modelArn"]) for m in summaries]
+
+            if self.embedder_model_name not in model_ids and self.embedder_model_name not in arns:
+                raise UserError(
+                    "model '{}' not found either : {} or {}".format(
+                        self.embedder_model_name, ", ".join(model_ids), ", ".join(arns)
+                    )
+                )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+    def get_client_kwargs(self) -> dict:
+        kwargs = {
+            "region_name": self.region_name,
+        }
+        
+        if self.endpoint_url:
+            kwargs["endpoint_url"] = self.endpoint_url
+            
+        if self.access_method == "credentials":
+            if self.aws_access_key_id and self.aws_secret_access_key:
+                kwargs["aws_access_key_id"] = self.aws_access_key_id.get_secret_value()
+                kwargs["aws_secret_access_key"] = self.aws_secret_access_key.get_secret_value()
+            else:
+                raise ValueError(
+                    "Credentials access method requires aws_access_key_id and aws_secret_access_key"
+                )
+        elif self.access_method == "iam":
+            # For IAM, boto3 will use default credential chain (IAM roles, environment, etc.)
+            pass
+        else:
+            raise ValueError(
+                f"Invalid access_method: {self.access_method}. Must be 'credentials' or 'iam'"
+            )
+            
+        return kwargs
+
+    @requires_dependencies(
+        ["boto3"],
+        extras="bedrock",
+    )
+    def get_bedrock_client(self) -> "BedrockClient":
+        import boto3
+
+        bedrock_client = boto3.client(service_name="bedrock", **self.get_client_kwargs())
+
+        return bedrock_client
+
+    @requires_dependencies(
+        ["boto3", "numpy", "botocore"],
+        extras="bedrock",
+    )
+    def get_client(self) -> "BedrockRuntimeClient":
+        import boto3
+
+        bedrock_client = boto3.client(service_name="bedrock-runtime", **self.get_client_kwargs())
+
+        return bedrock_client
+
+    @requires_dependencies(
+        ["aioboto3"],
+        extras="bedrock",
+    )
+    @asynccontextmanager
+    async def get_async_client(self) -> AsyncIterable["AsyncBedrockRuntimeClient"]:
+        import aioboto3
+
+        session = aioboto3.Session()
+        async with session.client("bedrock-runtime", **self.get_client_kwargs()) as aws_bedrock:
+            yield aws_bedrock
+
+
+@dataclass
+class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: BedrockEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def embed_query(self, query: str) -> list[float]:
+        """Call out to Bedrock embedding endpoint."""
+        provider = self.config.embedder_model_name.split(".")[0]
+        body = conform_query(query=query, provider=provider)
+
+        bedrock_client = self.config.get_client()
+        # invoke bedrock API
+        try:
+            invoke_params = {
+                "body": json.dumps(body),
+                "modelId": self.config.embedder_model_name,
+                "accept": "application/json",
+                "contentType": "application/json",
+            }
+
+            # Add inference profile if configured
+            if self.config.inference_profile_id:
+                invoke_params["inferenceProfileId"] = self.config.inference_profile_id
+
+            response = bedrock_client.invoke_model(**invoke_params)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+        # format output based on provider
+        response_body = json.loads(response.get("body").read())
+        if provider == "cohere":
+            return response_body.get("embeddings")[0]
+        else:
+            # includes common provider == "amazon"
+            return response_body.get("embedding")
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        elements = elements.copy()
+        elements_with_text = [e for e in elements if e.get("text")]
+        if not elements_with_text:
+            return elements
+        embeddings = [self.embed_query(query=e["text"]) for e in elements_with_text]
+        for element, embedding in zip(elements_with_text, embeddings):
+            element[EMBEDDINGS_KEY] = embedding
+        return elements
+
+
+@dataclass
+class AsyncBedrockEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: BedrockEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    async def embed_query(self, query: str) -> list[float]:
+        """Call out to Bedrock embedding endpoint."""
+        provider = self.config.embedder_model_name.split(".")[0]
+        body = conform_query(query=query, provider=provider)
+        try:
+            async with self.config.get_async_client() as bedrock_client:
+                # invoke bedrock API
+                try:
+                    invoke_params = {
+                        "body": json.dumps(body),
+                        "modelId": self.config.embedder_model_name,
+                        "accept": "application/json",
+                        "contentType": "application/json",
+                    }
+
+                    # Add inference profile if configured
+                    if self.config.inference_profile_id:
+                        invoke_params["inferenceProfileId"] = self.config.inference_profile_id
+
+                    response = await bedrock_client.invoke_model(**invoke_params)
+                except Exception as e:
+                    raise self.wrap_error(e=e)
+                async with response.get("body") as client_response:
+                    response_body = await client_response.json()
+
+            # format output based on provider
+            if provider == "cohere":
+                return response_body.get("embeddings")[0]
+            else:
+                # includes common provider == "amazon"
+                return response_body.get("embedding")
+        except Exception as e:
+            raise ValueError(f"Error raised by inference endpoint: {e}")
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        elements = elements.copy()
+        elements_with_text = [e for e in elements if e.get("text")]
+        embeddings = await asyncio.gather(
+            *[self.embed_query(query=e.get("text", "")) for e in elements_with_text]
+        )
+        for element, embedding in zip(elements_with_text, embeddings):
+            element[EMBEDDINGS_KEY] = embedding
+        return elements

unstructured_ingest/embed/huggingface.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import Field
+
+from unstructured_ingest.embed.interfaces import (
+    EMBEDDINGS_KEY,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+
+class HuggingFaceEmbeddingConfig(EmbeddingConfig):
+    embedder_model_name: Optional[str] = Field(
+        default="all-MiniLM-L6-v2", alias="model_name", description="HuggingFace model name"
+    )
+    embedder_model_kwargs: Optional[dict] = Field(
+        default_factory=lambda: {"device": "cpu"},
+        alias="model_kwargs",
+        description="additional model parameters",
+    )
+    encode_kwargs: Optional[dict] = Field(
+        default_factory=lambda: {"normalize_embeddings": False},
+        description="additional embedding parameters",
+    )
+
+    @requires_dependencies(
+        ["sentence_transformers"],
+        extras="huggingface",
+    )
+    def get_client(self) -> "SentenceTransformer":
+        from sentence_transformers import SentenceTransformer
+
+        return SentenceTransformer(
+            model_name_or_path=self.embedder_model_name,
+            **self.embedder_model_kwargs,
+        )
+
+    def get_encoder_kwargs(self) -> dict:
+        encoder_kwargs = self.encode_kwargs or {}
+        encoder_kwargs["batch_size"] = self.batch_size
+        return encoder_kwargs
+
+
+@dataclass
+class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: HuggingFaceEmbeddingConfig = field(default_factory=HuggingFaceEmbeddingConfig)
+
+    def _embed_query(self, query: str) -> list[float]:
+        return self._embed_documents(texts=[query])[0]
+
+    def _embed_documents(self, texts: list[str]) -> list[list[float]]:
+        client = self.config.get_client()
+        embeddings = client.encode(texts, **self.config.get_encoder_kwargs())
+        return embeddings.tolist()
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        elements = elements.copy()
+        elements_with_text = [e for e in elements if e.get("text")]
+        if not elements_with_text:
+            return elements
+        embeddings = self._embed_documents([e["text"] for e in elements_with_text])
+        for element, embedding in zip(elements_with_text, embeddings):
+            element[EMBEDDINGS_KEY] = embedding
+        return elements

unstructured_ingest/embed/interfaces.py

@@ -0,0 +1,146 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+EMBEDDINGS_KEY = "embeddings"
+
+
+class EmbeddingConfig(BaseModel):
+    batch_size: Optional[int] = Field(
+        default=32, description="Optional batch size for embedding requests."
+    )
+
+
+@dataclass
+class BaseEncoder(ABC):
+    config: EmbeddingConfig
+
+    def precheck(self):
+        pass
+
+    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
+
+
+@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 dimension(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return len(exemplary_embedding)
+
+    def get_exemplary_embedding(self) -> list[float]:
+        return self.embed_query(query="Q")
+
+    @property
+    @requires_dependencies(["numpy"])
+    def is_unit_vector(self) -> bool:
+        """Denotes if the embedding vector is a unit vector."""
+        import numpy as np
+
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0, rtol=1e-03)
+
+    def get_client(self):
+        raise NotImplementedError
+
+    def embed_batch(self, client: Any, batch: list[str]) -> list[list[float]]:
+        raise NotImplementedError
+
+    def embed_documents(self, elements: list[dict]) -> list[dict]:
+        client = self.get_client()
+        elements = elements.copy()
+        elements_with_text = [e for e in elements if e.get("text")]
+        texts = [e["text"] for e in elements_with_text]
+        all_embeddings = []
+        try:
+            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
+                embeddings_batch = self.embed_batch(client=client, batch=batch)
+                all_embeddings.extend(embeddings_batch)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        for element, embedding in zip(elements_with_text, all_embeddings, strict=True):
+            element[EMBEDDINGS_KEY] = embedding
+        return elements
+
+    def _embed_query(self, query: str) -> list[float]:
+        client = self.get_client()
+        return self.embed_batch(client=client, batch=[query])[0]
+
+    def embed_query(self, query: str) -> list[float]:
+        try:
+            return self._embed_query(query=query)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+
+@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 dimension(self):
+        exemplary_embedding = await self.get_exemplary_embedding()
+        return len(exemplary_embedding)
+
+    async def get_exemplary_embedding(self) -> list[float]:
+        return await self.embed_query(query="Q")
+
+    @property
+    @requires_dependencies(["numpy"])
+    async def is_unit_vector(self) -> bool:
+        """Denotes if the embedding vector is a unit vector."""
+        import numpy as np
+
+        exemplary_embedding = await self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0, rtol=1e-03)
+
+    def get_client(self):
+        raise NotImplementedError
+
+    async def embed_batch(self, client: Any, batch: list[str]) -> list[list[float]]:
+        raise NotImplementedError
+
+    async def embed_documents(self, elements: list[dict]) -> list[dict]:
+        client = self.get_client()
+        elements = elements.copy()
+        elements_with_text = [e for e in elements if e.get("text")]
+        texts = [e["text"] for e in elements_with_text]
+        all_embeddings = []
+        try:
+            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
+                embeddings_batch = await self.embed_batch(client=client, batch=batch)
+                all_embeddings.extend(embeddings_batch)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        for element, embedding in zip(elements_with_text, all_embeddings, strict=True):
+            element[EMBEDDINGS_KEY] = embedding
+        return elements
+
+    async def _embed_query(self, query: str) -> list[float]:
+        client = self.get_client()
+        embeddings = await self.embed_batch(client=client, batch=[query])
+        return embeddings[0]
+
+    async def embed_query(self, query: str) -> list[float]:
+        try:
+            return await self._embed_query(query=query)
+        except Exception as e:
+            raise self.wrap_error(e=e)

unstructured_ingest/embed/mixedbreadai.py

@@ -0,0 +1,134 @@
+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"
+TIMEOUT = 60
+MAX_RETRIES = 3
+ENCODING_FORMAT = "float"
+TRUNCATION_STRATEGY = "end"
+
+
+if TYPE_CHECKING:
+    from mixedbread import AsyncMixedbread, Mixedbread
+
+
+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")),
+        description="API key for Mixedbread AI",
+    )
+
+    embedder_model_name: str = Field(
+        default="mixedbread-ai/mxbai-embed-large-v1",
+        alias="model_name",
+        description="Mixedbread AI model name",
+    )
+
+    @requires_dependencies(
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
+    )
+    def get_client(self) -> "Mixedbread":
+        """
+        Create the Mixedbread AI client.
+
+        Returns:
+            Mixedbread: Initialized client.
+        """
+        from mixedbread import Mixedbread
+
+        return Mixedbread(
+            api_key=self.api_key.get_secret_value(),
+            max_retries=MAX_RETRIES,
+        )
+
+    @requires_dependencies(
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
+    )
+    def get_async_client(self) -> "AsyncMixedbread":
+        from mixedbread import AsyncMixedbread
+
+        return AsyncMixedbread(
+            api_key=self.api_key.get_secret_value(),
+            max_retries=MAX_RETRIES,
+        )
+
+
+@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_query(query="Q")
+
+    @requires_dependencies(
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
+    )
+    def get_client(self) -> "Mixedbread":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "Mixedbread", batch: list[str]) -> list[list[float]]:
+        response = client.embed(
+            model=self.config.embedder_model_name,
+            input=batch,
+            normalized=True,
+            encoding_format=ENCODING_FORMAT,
+            extra_headers={"User-Agent": USER_AGENT},
+            timeout=TIMEOUT,
+        )
+        return [datum.embedding for datum in response.data]
+
+
+@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."""
+        return await self.embed_query(query="Q")
+
+    @requires_dependencies(
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
+    )
+    def get_client(self) -> "AsyncMixedbread":
+        return self.config.get_async_client()
+
+    async def embed_batch(self, client: "AsyncMixedbread", batch: list[str]) -> list[list[float]]:
+        response = await client.embed(
+            model=self.config.embedder_model_name,
+            input=batch,
+            normalized=True,
+            encoding_format=ENCODING_FORMAT,
+            extra_headers={"User-Agent": USER_AGENT},
+            timeout=TIMEOUT,
+        )
+        return [datum.embedding for datum in response.data]