unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/embed/octoai.py

@@ -0,0 +1,133 @@
+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.error import (
+    ProviderError,
+    QuotaError,
+    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 openai import AsyncOpenAI, OpenAI
+
+
+class OctoAiEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr = Field(description="API key for OctoAI")
+    embedder_model_name: str = Field(
+        default="thenlper/gte-large", alias="model_name", description="octoai model name"
+    )
+    base_url: str = Field(
+        default="https://text.octoai.run/v1", description="optional override for the base url"
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        # 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
+
+    def run_precheck(self) -> None:
+        client = self.get_client()
+        try:
+            models = [m.id for m in list(client.models.list())]
+            if self.embedder_model_name not in models:
+                raise UserError(
+                    "model '{}' not found: {}".format(self.embedder_model_name, ", ".join(models))
+                )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+    @requires_dependencies(
+        ["openai", "tiktoken"],
+        extras="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="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 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):
+        client = self.get_client()
+        response = client.embeddings.create(input=query, model=self.config.embedder_model_name)
+        return response.data[0].embedding
+
+    def get_client(self) -> "OpenAI":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "OpenAI", batch: list[str]) -> list[list[float]]:
+        response = client.embeddings.create(input=batch, model=self.config.embedder_model_name)
+        return [data.embedding for data in response.data]
+
+
+@dataclass
+class AsyncOctoAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: OctoAiEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "AsyncOpenAI":
+        return self.config.get_async_client()
+
+    async def embed_batch(self, client: "AsyncOpenAI", batch: list[str]) -> list[list[float]]:
+        response = await client.embeddings.create(
+            input=batch, model=self.config.embedder_model_name
+        )
+        return [data.embedding for data in response.data]

unstructured_ingest/embed/openai.py

@@ -0,0 +1,142 @@
+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.error import (
+    ProviderError,
+    QuotaError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+    is_internal_error,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.tls import ssl_context_with_optional_ca_override
+
+if TYPE_CHECKING:
+    from openai import AsyncOpenAI, OpenAI
+
+
+class OpenAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr = Field(description="API key for OpenAI")
+    embedder_model_name: str = Field(
+        default="text-embedding-ada-002", alias="model_name", description="OpenAI model name"
+    )
+    base_url: Optional[str] = Field(default=None, description="optional override for the base url")
+
+    @requires_dependencies(["openai"], extras="openai")
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        # 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_models(self) -> Optional[list[str]]:
+        # In case the list model endpoint isn't exposed, don't break
+        from openai import APIStatusError
+
+        client = self.get_client()
+        try:
+            models = [m.id for m in list(client.models.list())]
+            return models
+        except APIStatusError as e:
+            if e.status_code == 404:
+                return None
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+    def run_precheck(self) -> None:
+        try:
+            models = self.get_models()
+            if models is None:
+                return
+            if self.embedder_model_name not in models:
+                raise UserError(
+                    "model '{}' not found: {}".format(self.embedder_model_name, ", ".join(models))
+                )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_client(self) -> "OpenAI":
+        from openai import DefaultHttpxClient, OpenAI
+
+        client = DefaultHttpxClient(verify=ssl_context_with_optional_ca_override())
+        return OpenAI(
+            api_key=self.api_key.get_secret_value(), http_client=client, base_url=self.base_url
+        )
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_async_client(self) -> "AsyncOpenAI":
+        from openai import AsyncOpenAI, DefaultAsyncHttpxClient
+
+        client = DefaultAsyncHttpxClient(verify=ssl_context_with_optional_ca_override())
+        return AsyncOpenAI(
+            api_key=self.api_key.get_secret_value(), http_client=client, base_url=self.base_url
+        )
+
+
+@dataclass
+class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OpenAIEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "OpenAI":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "OpenAI", batch: list[str]) -> list[list[float]]:
+        response = client.embeddings.create(input=batch, model=self.config.embedder_model_name)
+        return [data.embedding for data in response.data]
+
+
+@dataclass
+class AsyncOpenAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: OpenAIEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "AsyncOpenAI":
+        return self.config.get_async_client()
+
+    async def embed_batch(self, client: "AsyncOpenAI", batch: list[str]) -> list[list[float]]:
+        response = await client.embeddings.create(
+            input=batch, model=self.config.embedder_model_name
+        )
+        return [data.embedding for data in response.data]

unstructured_ingest/embed/togetherai.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from pydantic import Field, SecretStr
+
+from unstructured_ingest.embed.interfaces import (
+    AsyncBaseEmbeddingEncoder,
+    BaseEmbeddingEncoder,
+    EmbeddingConfig,
+)
+from unstructured_ingest.error import (
+    ProviderError,
+    UserAuthError,
+    UserError,
+    is_internal_error,
+)
+from unstructured_ingest.error import (
+    RateLimitError as CustomRateLimitError,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from together import AsyncTogether, Together
+
+
+class TogetherAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr = Field(description="API key for Together AI")
+    embedder_model_name: str = Field(
+        default="togethercomputer/m2-bert-80M-32k-retrieval",
+        alias="model_name",
+        description="Together AI model name",
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        # 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 together: {e}", exc_info=True)
+            return e
+        message = e.args[0]
+        if isinstance(e, AuthenticationError):
+            return UserAuthError(message)
+        if isinstance(e, RateLimitError):
+            return CustomRateLimitError(message)
+
+        status_code = getattr(e, "status_code", None)
+        if status_code is not None:
+            if 400 <= status_code < 500:
+                return UserError(message)
+            if status_code >= 500:
+                return ProviderError(message)
+        logger.error(f"unhandled exception from together: {e}", exc_info=True)
+        return e
+
+    def run_precheck(self) -> None:
+        client = self.get_client()
+        try:
+            models = [m.id for m in list(client.models.list())]
+            if self.embedder_model_name not in models:
+                raise UserError(
+                    "model '{}' not found: {}".format(self.embedder_model_name, ", ".join(models))
+                )
+        except Exception as e:
+            raise self.wrap_error(e=e)
+
+    @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 precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "Together":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "Together", batch: list[str]) -> list[list[float]]:
+        outputs = client.embeddings.create(model=self.config.embedder_model_name, input=batch)
+        return [outputs.data[i].embedding for i in range(len(batch))]
+
+
+@dataclass
+class AsyncTogetherAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: TogetherAIEmbeddingConfig
+
+    def precheck(self):
+        self.config.run_precheck()
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "AsyncTogether":
+        return self.config.get_async_client()
+
+    async def embed_batch(self, client: Any, batch: list[str]) -> list[list[float]]:
+        outputs = await client.embeddings.create(model=self.config.embedder_model_name, input=batch)
+        return [outputs.data[i].embedding for i in range(len(batch))]

unstructured_ingest/embed/vertexai.py

@@ -0,0 +1,109 @@
+# 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.error import UserAuthError, is_internal_error
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+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 = Field(description="API key for Vertex AI")
+    embedder_model_name: Optional[str] = Field(
+        default="text-embedding-005", alias="model_name", description="Vertex AI model name"
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        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="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 get_client(self) -> "TextEmbeddingModel":
+        return self.config.get_client()
+
+    @requires_dependencies(
+        ["vertexai"],
+        extras="embed-vertexai",
+    )
+    def embed_batch(self, client: "TextEmbeddingModel", batch: list[str]) -> list[list[float]]:
+        from vertexai.language_models import TextEmbeddingInput
+
+        inputs = [TextEmbeddingInput(text=text) for text in batch]
+        response = client.get_embeddings(inputs)
+        return [e.values for e in response]
+
+
+@dataclass
+class AsyncVertexAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: VertexAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "TextEmbeddingModel":
+        return self.config.get_client()
+
+    @requires_dependencies(
+        ["vertexai"],
+        extras="embed-vertexai",
+    )
+    async def embed_batch(self, client: Any, batch: list[str]) -> list[list[float]]:
+        from vertexai.language_models import TextEmbeddingInput
+
+        inputs = [TextEmbeddingInput(text=text) for text in batch]
+        response = await client.get_embeddings_async(inputs)
+        return [e.values for e in response]

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.error import ProviderError, UserAuthError, UserError, is_internal_error
+from unstructured_ingest.error import (
+    RateLimitError as CustomRateLimitError,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from voyageai import AsyncClient as AsyncVoyageAIClient
+    from voyageai import Client as VoyageAIClient
+
+
+class VoyageAIEmbeddingConfig(EmbeddingConfig):
+    batch_size: int = Field(
+        default=32,
+        le=128,
+        description="Batch size for embedding requests. VoyageAI has a limit of 128.",
+    )
+    api_key: SecretStr = Field(description="API key for VoyageAI")
+    embedder_model_name: str = Field(
+        default="voyage-3", alias="model_name", description="VoyageAI model name"
+    )
+    max_retries: int = Field(default=0, description="Max retries for embedding requests.")
+    timeout_in_seconds: Optional[int] = Field(
+        default=None, description="Optional timeout in seconds for embedding requests."
+    )
+
+    def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
+        # 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 get_client(self) -> "VoyageAIClient":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "VoyageAIClient", batch: list[str]) -> list[list[float]]:
+        if self.config.embedder_model_name == "voyage-multimodal-3":
+            batch = [[text] for text in batch]
+            response = client.multimodal_embed(inputs=batch, model=self.config.embedder_model_name)
+        else:
+            response = client.embed(texts=batch, model=self.config.embedder_model_name)
+        return response.embeddings
+
+
+@dataclass
+class AsyncVoyageAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
+    config: VoyageAIEmbeddingConfig
+
+    def wrap_error(self, e: Exception) -> Exception:
+        return self.config.wrap_error(e=e)
+
+    def get_client(self) -> "AsyncVoyageAIClient":
+        return self.config.get_async_client()
+
+    async def embed_batch(
+        self, client: "AsyncVoyageAIClient", batch: list[str]
+    ) -> list[list[float]]:
+        if self.config.embedder_model_name == "voyage-multimodal-3":
+            batch = [[text] for text in batch]
+            response = await client.multimodal_embed(
+                inputs=batch, model=self.config.embedder_model_name
+            )
+        else:
+            response = await client.embed(texts=batch, model=self.config.embedder_model_name)
+        return response.embeddings