unstructured-ingest 0.5.2__tar.gz → 0.5.3__tar.gz

{unstructured_ingest-0.5.2/unstructured_ingest.egg-info → unstructured_ingest-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: unstructured-ingest
-Version: 0.5.2
+Version: 0.5.3
 Summary: A library that prepares raw documents for downstream ML tasks.
 Home-page: https://github.com/Unstructured-IO/unstructured-ingest
 Author: Unstructured Technologies
@@ -22,13 +22,13 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9.0,<3.14
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
-Requires-Dist: pydantic>=2.7
-Requires-Dist: click
-Requires-Dist: pandas
+Requires-Dist: python-dateutil
 Requires-Dist: opentelemetry-sdk
-Requires-Dist: dataclasses_json
 Requires-Dist: tqdm
-Requires-Dist: python-dateutil
+Requires-Dist: pandas
+Requires-Dist: click
+Requires-Dist: pydantic>=2.7
+Requires-Dist: dataclasses_json
 Provides-Extra: remote
 Requires-Dist: unstructured-client>=0.26.1; extra == "remote"
 Provides-Extra: csv
@@ -71,39 +71,39 @@ Requires-Dist: fsspec; extra == "azure"
 Provides-Extra: azure-ai-search
 Requires-Dist: azure-search-documents; extra == "azure-ai-search"
 Provides-Extra: biomed
-Requires-Dist: bs4; extra == "biomed"
 Requires-Dist: requests; extra == "biomed"
+Requires-Dist: bs4; extra == "biomed"
 Provides-Extra: box
-Requires-Dist: boxfs; extra == "box"
 Requires-Dist: fsspec; extra == "box"
+Requires-Dist: boxfs; extra == "box"
 Provides-Extra: chroma
 Requires-Dist: chromadb; extra == "chroma"
 Provides-Extra: clarifai
 Requires-Dist: clarifai; extra == "clarifai"
 Provides-Extra: confluence
-Requires-Dist: requests; extra == "confluence"
 Requires-Dist: atlassian-python-api; extra == "confluence"
+Requires-Dist: requests; extra == "confluence"
 Provides-Extra: couchbase
 Requires-Dist: couchbase; extra == "couchbase"
 Provides-Extra: delta-table
-Requires-Dist: boto3; extra == "delta-table"
 Requires-Dist: deltalake; extra == "delta-table"
+Requires-Dist: boto3; extra == "delta-table"
 Provides-Extra: discord
 Requires-Dist: discord.py; extra == "discord"
 Provides-Extra: dropbox
-Requires-Dist: dropboxdrivefs; extra == "dropbox"
 Requires-Dist: fsspec; extra == "dropbox"
+Requires-Dist: dropboxdrivefs; extra == "dropbox"
 Provides-Extra: duckdb
 Requires-Dist: duckdb; extra == "duckdb"
 Provides-Extra: elasticsearch
 Requires-Dist: elasticsearch[async]; extra == "elasticsearch"
 Provides-Extra: gcs
-Requires-Dist: bs4; extra == "gcs"
 Requires-Dist: gcsfs; extra == "gcs"
 Requires-Dist: fsspec; extra == "gcs"
+Requires-Dist: bs4; extra == "gcs"
 Provides-Extra: github
-Requires-Dist: requests; extra == "github"
 Requires-Dist: pygithub>1.58.0; extra == "github"
+Requires-Dist: requests; extra == "github"
 Provides-Extra: gitlab
 Requires-Dist: python-gitlab; extra == "gitlab"
 Provides-Extra: google-drive
@@ -124,14 +124,14 @@ Requires-Dist: pymilvus; extra == "milvus"
 Provides-Extra: mongodb
 Requires-Dist: pymongo; extra == "mongodb"
 Provides-Extra: neo4j
-Requires-Dist: neo4j; extra == "neo4j"
 Requires-Dist: cymple; extra == "neo4j"
+Requires-Dist: neo4j-rust-ext; extra == "neo4j"
 Requires-Dist: networkx; extra == "neo4j"
 Provides-Extra: notion
+Requires-Dist: htmlBuilder; extra == "notion"
 Requires-Dist: notion-client; extra == "notion"
 Requires-Dist: backoff; extra == "notion"
 Requires-Dist: httpx; extra == "notion"
-Requires-Dist: htmlBuilder; extra == "notion"
 Provides-Extra: onedrive
 Requires-Dist: Office365-REST-Python-Client; extra == "onedrive"
 Requires-Dist: msal; extra == "onedrive"
@@ -152,16 +152,16 @@ Requires-Dist: praw; extra == "reddit"
 Provides-Extra: redis
 Requires-Dist: redis; extra == "redis"
 Provides-Extra: s3
-Requires-Dist: s3fs; extra == "s3"
 Requires-Dist: fsspec; extra == "s3"
+Requires-Dist: s3fs; extra == "s3"
 Provides-Extra: sharepoint
 Requires-Dist: Office365-REST-Python-Client; extra == "sharepoint"
 Requires-Dist: msal; extra == "sharepoint"
 Provides-Extra: salesforce
 Requires-Dist: simple-salesforce; extra == "salesforce"
 Provides-Extra: sftp
-Requires-Dist: paramiko; extra == "sftp"
 Requires-Dist: fsspec; extra == "sftp"
+Requires-Dist: paramiko; extra == "sftp"
 Provides-Extra: slack
 Requires-Dist: slack_sdk[optional]; extra == "slack"
 Provides-Extra: snowflake
@@ -178,13 +178,13 @@ Requires-Dist: databricks-sql-connector; extra == "databricks-delta-tables"
 Provides-Extra: singlestore
 Requires-Dist: singlestoredb; extra == "singlestore"
 Provides-Extra: vectara
-Requires-Dist: requests; extra == "vectara"
 Requires-Dist: aiofiles; extra == "vectara"
+Requires-Dist: requests; extra == "vectara"
 Requires-Dist: httpx; extra == "vectara"
 Provides-Extra: vastdb
+Requires-Dist: vastdb; extra == "vastdb"
 Requires-Dist: pyarrow; extra == "vastdb"
 Requires-Dist: ibis; extra == "vastdb"
-Requires-Dist: vastdb; extra == "vastdb"
 Provides-Extra: embed-huggingface
 Requires-Dist: sentence-transformers; extra == "embed-huggingface"
 Provides-Extra: embed-octoai
@@ -200,10 +200,22 @@ Provides-Extra: openai
 Requires-Dist: openai; extra == "openai"
 Requires-Dist: tiktoken; extra == "openai"
 Provides-Extra: bedrock
-Requires-Dist: boto3; extra == "bedrock"
 Requires-Dist: aioboto3; extra == "bedrock"
+Requires-Dist: boto3; extra == "bedrock"
 Provides-Extra: togetherai
 Requires-Dist: together; extra == "togetherai"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # Unstructured Ingest  
 

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/test/unit/v2/embedders/test_bedrock.py

@@ -15,7 +15,7 @@ def generate_embedder_config_params() -> dict:
         "region_name": fake.city(),
     }
     if random.random() < 0.5:
-        params["embed_model_name"] = fake.word()
+        params["embedder_model_name"] = fake.word()
     return params
 
 

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/test/unit/v2/embedders/test_huggingface.py

@@ -16,7 +16,7 @@ fake = faker.Faker()
 def generate_embedder_config_params() -> dict:
     params = {}
     if random.random() < 0.5:
-        params["embed_model_name"] = fake.word() if random.random() < 0.5 else None
+        params["embedder_model_name"] = fake.word() if random.random() < 0.5 else None
         params["embedder_model_kwargs"] = (
             generate_random_dictionary(key_type=str, value_type=Any)
             if random.random() < 0.5

unstructured_ingest-0.5.3/unstructured_ingest/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.5.3"  # pragma: no cover

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/unstructured_ingest/embed/azure_openai.py

@@ -44,7 +44,13 @@ class AzureOpenAIEmbeddingConfig(OpenAIEmbeddingConfig):
 class AzureOpenAIEmbeddingEncoder(OpenAIEmbeddingEncoder):
     config: AzureOpenAIEmbeddingConfig
 
+    def get_client(self) -> "AzureOpenAI":
+        return self.config.get_client()
+
 
 @dataclass
 class AsyncAzureOpenAIEmbeddingEncoder(AsyncOpenAIEmbeddingEncoder):
     config: AzureOpenAIEmbeddingConfig
+
+    def get_client(self) -> "AsyncAzureOpenAI":
+        return self.config.get_async_client()

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/unstructured_ingest/embed/bedrock.py

@@ -15,7 +15,13 @@ from unstructured_ingest.embed.interfaces import (
 )
 from unstructured_ingest.logger import logger
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.errors import ProviderError, RateLimitError, UserAuthError, UserError
+from unstructured_ingest.v2.errors import (
+    ProviderError,
+    RateLimitError,
+    UserAuthError,
+    UserError,
+    is_internal_error,
+)
 
 if TYPE_CHECKING:
     from botocore.client import BaseClient
@@ -51,9 +57,11 @@ class BedrockEmbeddingConfig(EmbeddingConfig):
     aws_access_key_id: SecretStr
     aws_secret_access_key: SecretStr
     region_name: str = "us-west-2"
-    embed_model_name: str = Field(default="amazon.titan-embed-text-v1", alias="model_name")
+    embedder_model_name: str = Field(default="amazon.titan-embed-text-v1", alias="model_name")
 
     def wrap_error(self, e: Exception) -> Exception:
+        if is_internal_error(e=e):
+            return e
         from botocore.exceptions import ClientError
 
         if isinstance(e, ClientError):
@@ -122,7 +130,7 @@ class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
 
     def embed_query(self, query: str) -> list[float]:
         """Call out to Bedrock embedding endpoint."""
-        provider = self.config.embed_model_name.split(".")[0]
+        provider = self.config.embedder_model_name.split(".")[0]
         body = conform_query(query=query, provider=provider)
 
         bedrock_client = self.config.get_client()
@@ -130,7 +138,7 @@ class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
         try:
             response = bedrock_client.invoke_model(
                 body=json.dumps(body),
-                modelId=self.config.embed_model_name,
+                modelId=self.config.embedder_model_name,
                 accept="application/json",
                 contentType="application/json",
             )
@@ -148,6 +156,8 @@ class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
     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
@@ -163,7 +173,7 @@ class AsyncBedrockEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
 
     async def embed_query(self, query: str) -> list[float]:
         """Call out to Bedrock embedding endpoint."""
-        provider = self.config.embed_model_name.split(".")[0]
+        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:
@@ -171,7 +181,7 @@ class AsyncBedrockEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
                 try:
                     response = await bedrock_client.invoke_model(
                         body=json.dumps(body),
-                        modelId=self.config.embed_model_name,
+                        modelId=self.config.embedder_model_name,
                         accept="application/json",
                         contentType="application/json",
                     )

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/unstructured_ingest/embed/huggingface.py

@@ -47,7 +47,7 @@ class HuggingFaceEmbeddingConfig(EmbeddingConfig):
 class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
     config: HuggingFaceEmbeddingConfig
 
-    def embed_query(self, query: str) -> list[float]:
+    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]]:
@@ -58,6 +58,8 @@ class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
     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

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/unstructured_ingest/embed/interfaces.py

@@ -1,11 +1,12 @@
-import asyncio
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 from pydantic import BaseModel, Field
 
+from unstructured_ingest.utils.data_prep import batch_generator
+
 EMBEDDINGS_KEY = "embeddings"
 
 
@@ -50,21 +51,37 @@ class BaseEmbeddingEncoder(BaseEncoder, ABC):
         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
+    def get_client(self):
+        raise NotImplementedError
 
-    @abstractmethod
-    def embed_query(self, query: str) -> list[float]:
-        pass
+    def embed_batch(self, client: Any, batch: list[str]) -> list[list[float]]:
+        raise NotImplementedError
 
-    def _embed_documents(self, elements: list[str]) -> list[list[float]]:
-        results = []
-        for text in elements:
-            response = self.embed_query(query=text)
-            results.append(response)
+    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]
+        embeddings = []
+        try:
+            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
+                embeddings = self.embed_batch(client=client, batch=batch)
+                embeddings.extend(embeddings)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        for element, embedding in zip(elements_with_text, embeddings):
+            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]
 
-        return results
+    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
@@ -88,14 +105,35 @@ class AsyncBaseEmbeddingEncoder(BaseEncoder, ABC):
         exemplary_embedding = await self.get_exemplary_embedding()
         return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
 
-    @abstractmethod
+    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]:
-        pass
+        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]
+        embeddings = []
+        try:
+            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
+                embeddings = await self.embed_batch(client=client, batch=batch)
+                embeddings.extend(embeddings)
+        except Exception as e:
+            raise self.wrap_error(e=e)
+        for element, embedding in zip(elements_with_text, embeddings):
+            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]
 
-    @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
+        try:
+            return await self._embed_query(query=query)
+        except Exception as e:
+            raise self.wrap_error(e=e)

unstructured_ingest-0.5.3/unstructured_ingest/embed/mixedbreadai.py

@@ -0,0 +1,149 @@
+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_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_query(query="Q")
+
+    @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 get_client(self) -> "MixedbreadAI":
+        return self.config.get_client()
+
+    def embed_batch(self, client: "MixedbreadAI", batch: list[str]) -> list[list[float]]:
+        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,
+        )
+        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_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 get_client(self) -> "AsyncMixedbreadAI":
+        return self.config.get_async_client()
+
+    async def embed_batch(self, client: "AsyncMixedbreadAI", batch: list[str]) -> list[list[float]]:
+        response = await 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,
+        )
+        return [datum.embedding for datum in response.data]

{unstructured_ingest-0.5.2 → unstructured_ingest-0.5.3}/unstructured_ingest/embed/octoai.py

@@ -4,13 +4,11 @@ from typing import TYPE_CHECKING
 from pydantic import Field, SecretStr
 
 from unstructured_ingest.embed.interfaces import (
-    EMBEDDINGS_KEY,
     AsyncBaseEmbeddingEncoder,
     BaseEmbeddingEncoder,
     EmbeddingConfig,
 )
 from unstructured_ingest.logger import logger
-from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.errors import (
     ProviderError,
@@ -18,6 +16,7 @@ from unstructured_ingest.v2.errors import (
     RateLimitError,
     UserAuthError,
     UserError,
+    is_internal_error,
 )
 
 if TYPE_CHECKING:
@@ -30,6 +29,8 @@ class OctoAiEmbeddingConfig(EmbeddingConfig):
     base_url: str = Field(default="https://text.octoai.run/v1")
 
     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
 
@@ -81,31 +82,17 @@ class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
     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)
+    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 embed_documents(self, elements: list[dict]) -> list[dict]:
-        elements = elements.copy()
-        elements_with_text = [e for e in elements if e.get("text")]
-        texts = [e["text"] for e in elements_with_text]
-        embeddings = []
-        client = self.config.get_client()
-        try:
-            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
-                response = client.embeddings.create(
-                    input=batch, model=self.config.embedder_model_name
-                )
-                embeddings.extend([data.embedding for data in response.data])
-        except Exception as e:
-            raise self.wrap_error(e=e)
-        for element, embedding in zip(elements_with_text, embeddings):
-            element[EMBEDDINGS_KEY] = embedding
-        return elements
+    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
@@ -115,30 +102,11 @@ class AsyncOctoAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
     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
+    def get_client(self) -> "AsyncOpenAI":
+        return self.config.get_async_client()
 
-    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")]
-        texts = [e["text"] for e in elements_with_text]
-        client = self.config.get_async_client()
-        embeddings = []
-        try:
-            for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
-                response = await client.embeddings.create(
-                    input=batch, model=self.config.embedder_model_name
-                )
-                embeddings.extend([data.embedding for data in response.data])
-        except Exception as e:
-            raise self.wrap_error(e=e)
-        for element, embedding in zip(elements_with_text, embeddings):
-            element[EMBEDDINGS_KEY] = embedding
-        return elements
+    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]