unstructured-ingest 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl

test/integration/connectors/weaviate/test_local.py

@@ -25,7 +25,7 @@ def wait_for_container(timeout: int = 10, interval: int = 1) -> None:
     start_time = time.time()
     while time.time() - start_time < timeout:
         try:
-            requests.get("http://localhost:8080/v1/.well-known/read")
+            requests.get("http://localhost:8080/v1/.well-known/read", timeout=1)
             return
         except Exception as e:
             print(f"Failed to validate container healthy, sleeping for {interval} seconds: {e}")
@@ -34,15 +34,20 @@ def wait_for_container(timeout: int = 10, interval: int = 1) -> None:
 
 
 @pytest.fixture
-def collection(collections_schema_config: dict) -> str:
+def weaviate_instance():
     with container_context(
         image="semitechnologies/weaviate:1.27.3",
         ports={8080: 8080, 50051: 50051},
-    ):
+    ) as ctx:
         wait_for_container()
-        with weaviate.connect_to_local() as weaviate_client:
-            weaviate_client.collections.create_from_dict(config=collections_schema_config)
-        yield COLLECTION_NAME
+        yield ctx
+
+
+@pytest.fixture
+def collection(weaviate_instance, collections_schema_config: dict) -> str:
+    with weaviate.connect_to_local() as weaviate_client:
+        weaviate_client.collections.create_from_dict(config=collections_schema_config)
+    return COLLECTION_NAME
 
 
 def get_count(client: WeaviateClient) -> int:
@@ -129,3 +134,19 @@ def test_weaviate_local_destination(upload_file: Path, collection: str, tmp_path
         file_data=file_data,
         expected_count=expected_count,
     )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
+def test_weaviate_local_create_destination(weaviate_instance):
+    uploader = LocalWeaviateUploader(
+        upload_config=LocalWeaviateUploaderConfig(),
+        connection_config=LocalWeaviateConnectionConfig(),
+    )
+    collection_name = "system_created"
+    created = uploader.create_destination(destination_name=collection_name)
+    assert created
+    with uploader.connection_config.get_client() as weaviate_client:
+        assert weaviate_client.collections.exists(name=collection_name)
+
+    created = uploader.create_destination(destination_name=collection_name)
+    assert not created

test/integration/embedders/test_azure_openai.py

@@ -54,6 +54,4 @@ def test_raw_azure_openai_embedder(embedder_file: Path):
             azure_endpoint=azure_data.endpoint,
         )
     )
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1536,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1536)

test/integration/embedders/test_bedrock.py

@@ -55,7 +55,7 @@ def test_raw_bedrock_embedder(embedder_file: Path):
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(1536,),
+        expected_dimension=1536,
         expected_is_unit_vector=False,
     )
 
@@ -98,6 +98,6 @@ async def test_raw_async_bedrock_embedder(embedder_file: Path):
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(1536,),
+        expected_dimension=1536,
         expected_is_unit_vector=False,
     )

test/integration/embedders/test_huggingface.py

@@ -21,6 +21,4 @@ def test_huggingface_embedder(embedder_file: Path):
 
 def test_raw_hugginface_embedder(embedder_file: Path):
     embedder = HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig())
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(384,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=384)

test/integration/embedders/test_mixedbread.py

@@ -49,7 +49,7 @@ def test_raw_mixedbread_embedder(embedder_file: Path):
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(1024,),
+        expected_dimension=1024,
         expected_is_unit_vector=False,
     )
 
@@ -66,6 +66,6 @@ async def test_raw_async_mixedbread_embedder(embedder_file: Path):
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(1024,),
+        expected_dimension=1024,
         expected_is_unit_vector=False,
     )

test/integration/embedders/test_octoai.py

@@ -47,9 +47,7 @@ def test_raw_octoai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1024)
 
 
 @pytest.mark.skip(reason="Unexpected connection error at the moment")
@@ -73,5 +71,5 @@ async def test_raw_async_octoai_embedder(embedder_file: Path):
         )
     )
     await validate_raw_embedder_async(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
+        embedder=embedder, embedder_file=embedder_file, expected_dimension=1024
     )

test/integration/embedders/test_openai.py

@@ -47,9 +47,7 @@ def test_raw_openai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1536,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1536)
 
 
 def test_raw_openai_embedder_invalid_credentials():
@@ -72,5 +70,5 @@ async def test_raw_async_openai_embedder(embedder_file: Path):
         )
     )
     await validate_raw_embedder_async(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1536,)
+        embedder=embedder, embedder_file=embedder_file, expected_dimension=1536
     )

test/integration/embedders/test_togetherai.py

@@ -46,7 +46,7 @@ def test_raw_togetherai_embedder(embedder_file: Path):
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(768,),
+        expected_dimension=768,
         expected_is_unit_vector=False,
     )
 
@@ -66,6 +66,6 @@ async def test_raw_async_togetherai_embedder(embedder_file: Path):
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,
-        expected_dimensions=(768,),
+        expected_dimension=768,
         expected_is_unit_vector=False,
     )

test/integration/embedders/test_vertexai.py

@@ -46,9 +46,7 @@ def test_raw_vertexai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(768,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=768)
 
 
 @requires_env(API_KEY)
@@ -61,5 +59,5 @@ async def test_raw_async_vertexai_embedder(embedder_file: Path):
         )
     )
     await validate_raw_embedder_async(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(768,)
+        embedder=embedder, embedder_file=embedder_file, expected_dimension=768
     )

test/integration/embedders/test_voyageai.py

@@ -46,9 +46,7 @@ def test_raw_voyageai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
-    validate_raw_embedder(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
-    )
+    validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1024)
 
 
 @requires_env(API_KEY)
@@ -61,5 +59,5 @@ async def test_raw_async_voyageai_embedder(embedder_file: Path):
         )
     )
     await validate_raw_embedder_async(
-        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1024,)
+        embedder=embedder, embedder_file=embedder_file, expected_dimension=1024
     )

test/integration/embedders/utils.py

@@ -23,23 +23,22 @@ def validate_embedding_output(original_elements: list[dict], output_elements: li
 def validate_raw_embedder(
     embedder: BaseEmbeddingEncoder,
     embedder_file: Path,
-    expected_dimensions: Optional[tuple[int, ...]] = None,
+    expected_dimension: Optional[int] = None,
     expected_is_unit_vector: bool = True,
 ):
     with open(embedder_file) as f:
         elements = json.load(f)
     all_text = [element["text"] for element in elements]
     single_text = all_text[0]
-    num_of_dimensions = embedder.num_of_dimensions
-    if expected_dimensions:
+    dimension = embedder.dimension
+    if expected_dimension:
         assert (
-            num_of_dimensions == expected_dimensions
-        ), f"number of dimensions {num_of_dimensions} didn't match expected: {expected_dimensions}"
+            dimension == expected_dimension
+        ), f"dimensions {dimension} didn't match expected: {expected_dimension}"
     is_unit_vector = embedder.is_unit_vector
     assert is_unit_vector == expected_is_unit_vector
     single_embedding = embedder.embed_query(query=single_text)
-    expected_length = num_of_dimensions[0]
-    assert len(single_embedding) == expected_length
+    assert len(single_embedding) == dimension
     embedded_elements = embedder.embed_documents(elements=elements)
     validate_embedding_output(original_elements=elements, output_elements=embedded_elements)
 
@@ -47,22 +46,21 @@ def validate_raw_embedder(
 async def validate_raw_embedder_async(
     embedder: AsyncBaseEmbeddingEncoder,
     embedder_file: Path,
-    expected_dimensions: Optional[tuple[int, ...]] = None,
+    expected_dimension: Optional[int] = None,
     expected_is_unit_vector: bool = True,
 ):
     with open(embedder_file) as f:
         elements = json.load(f)
     all_text = [element["text"] for element in elements]
     single_text = all_text[0]
-    num_of_dimensions = await embedder.num_of_dimensions
-    if expected_dimensions:
+    dimension = await embedder.dimension
+    if expected_dimension:
         assert (
-            num_of_dimensions == expected_dimensions
-        ), f"number of dimensions {num_of_dimensions} didn't match expected: {expected_dimensions}"
+            dimension == expected_dimension
+        ), f"dimension {dimension} didn't match expected: {expected_dimension}"
     is_unit_vector = await embedder.is_unit_vector
     assert is_unit_vector == expected_is_unit_vector
     single_embedding = await embedder.embed_query(query=single_text)
-    expected_length = num_of_dimensions[0]
-    assert len(single_embedding) == expected_length
+    assert len(single_embedding) == dimension
     embedded_elements = await embedder.embed_documents(elements=elements)
     validate_embedding_output(original_elements=elements, output_elements=embedded_elements)

test/unit/embed/test_openai.py

@@ -3,18 +3,26 @@ from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbedd
 
 def test_embed_documents_does_not_break_element_to_dict(mocker):
     # Mocked client with the desired behavior for embed_documents
+    raw_elements = [{"text": f"This is sentence {i + 1}"} for i in range(4)]
+    mock_response = mocker.MagicMock()
+    mock_response_data = []
+    for i in range(2):
+        mock_response_d = mocker.MagicMock()
+        mock_response_d.embedding = [1, 2]
+        mock_response_data.append(mock_response_d)
+    mock_response.data = mock_response_data
     mock_client = mocker.MagicMock()
-    mock_client.embed_documents.return_value = [1, 2]
+    mock_client.embeddings.create.return_value = mock_response
 
     # Mock get_client to return our mock_client
     mocker.patch.object(OpenAIEmbeddingConfig, "get_client", return_value=mock_client)
 
-    encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key"))
-    raw_elements = [{"text": f"This is sentence {i + 1}"} for i in range(2)]
+    encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key", batch_size=2))
 
     elements = encoder.embed_documents(
         elements=raw_elements,
     )
-    assert len(elements) == 2
+    assert len(elements) == 4
     assert elements[0]["text"] == "This is sentence 1"
     assert elements[1]["text"] == "This is sentence 2"
+    assert mock_client.embeddings.create.call_count == 2

test/unit/test_html.py

@@ -0,0 +1,112 @@
+import base64
+from pathlib import Path
+
+from bs4 import BeautifulSoup
+from pytest_mock import MockerFixture
+
+from unstructured_ingest.utils.html import HtmlMixin
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+
+
+def test_extract_images(mocker: MockerFixture):
+    mixin = HtmlMixin(extract_images=True)
+    mock_download_response = b"DOWNLOADED"
+    expected_image_src = base64.b64encode(mock_download_response).decode()
+    mocked_download_response = mocker.patch(
+        "unstructured_ingest.utils.html.HtmlMixin.download_content",
+        return_value=mock_download_response,
+    )
+    url = "http://mywebsite.com/path/to/page"
+    html = """
+    <img src="http://mywebsite.com/img1.jpg"/>
+    <img src="http://notmywebsite.com/img2.jpg"/>
+    <img src="img3.jpg"/>
+    <img src="data:image/png;base64,24689654..."/>
+    """
+    expected_html = f"""
+    <img src="data:image/png;base64,{expected_image_src}"/>
+    <img src="http://notmywebsite.com/img2.jpg"/>
+    <img src="data:image/png;base64,{expected_image_src}"/>
+    <img src="data:image/png;base64,24689654..."/>
+    """
+    expected_soup = BeautifulSoup(expected_html, "html.parser")
+    result = mixin.extract_html_images(url=url, html=html)
+    result_soup = BeautifulSoup(result, "html.parser")
+    assert expected_soup == result_soup
+    assert mocked_download_response.call_count == 2
+    urls_to_download = [
+        call_args_list.kwargs["url"] for call_args_list in mocked_download_response.call_args_list
+    ]
+    assert urls_to_download == ["http://mywebsite.com/img1.jpg", "http://mywebsite.com/img3.jpg"]
+
+
+def test_extract_images_allow_list(mocker: MockerFixture):
+    mixin = HtmlMixin(
+        extract_images=True, allow_list=["http://allowedwebsite1.com", "http://allowedwebsite2.com"]
+    )
+    mock_download_response = b"DOWNLOADED"
+    expected_image_src = base64.b64encode(mock_download_response).decode()
+    mocked_download_response = mocker.patch(
+        "unstructured_ingest.utils.html.HtmlMixin.download_content",
+        return_value=mock_download_response,
+    )
+    url = "http://mywebsite.com/path/to/page"
+    html = """
+    <img src="http://mywebsite.com/img1.jpg"/>
+    <img src="http://notmywebsite.com/img2.jpg"/>
+    <img src="http://allowedwebsite1.com/img2.jpg"/>
+    <img src="http://allowedwebsite2.com/img2.jpg"/>
+    """
+
+    expected_html = f"""
+    <img src="http://mywebsite.com/img1.jpg"/>
+    <img src="http://notmywebsite.com/img2.jpg"/>
+    <img src="data:image/png;base64,{expected_image_src}"/>
+    <img src="data:image/png;base64,{expected_image_src}"/>
+    """
+    expected_soup = BeautifulSoup(expected_html, "html.parser")
+    result = mixin.extract_html_images(url=url, html=html)
+    result_soup = BeautifulSoup(result, "html.parser")
+    assert expected_soup == result_soup
+    assert mocked_download_response.call_count == 2
+    urls_to_download = [
+        call_args_list.kwargs["url"] for call_args_list in mocked_download_response.call_args_list
+    ]
+    assert urls_to_download == [
+        "http://allowedwebsite1.com/img2.jpg",
+        "http://allowedwebsite2.com/img2.jpg",
+    ]
+
+
+def test_extract_embedded_docs(mocker: MockerFixture):
+    mixin = HtmlMixin(extract_files=True)
+    mock_download_response = b"DOWNLOADED"
+    mocked_download_response = mocker.patch(
+        "unstructured_ingest.utils.html.HtmlMixin.download_content",
+        return_value=mock_download_response,
+    )
+    mocked_write_content = mocker.patch("unstructured_ingest.utils.html.HtmlMixin.write_content")
+    url = "http://mywebsite.com/path/to/page"
+    html = """
+    <a href="http://mywebsite.com/file.pdf"/>
+    <a href="http://notmywebsite.com/file.pdf"/>
+    <a href="http://mywebsite.com/another/link"/>
+    <a href="another/link/2"/>
+    <a href="file.doc"/>
+    """
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath="file.txt",
+            filename="file.txt",
+        ),
+        connector_type="my_connector",
+        identifier="mock_file_data",
+    )
+    results = mixin.extract_embedded_files(
+        url=url, html=html, download_dir=Path("/tmp/download/location"), original_filedata=file_data
+    )
+    assert len(results) == 2
+    downloaded_urls = [r["file_data"].metadata.url for r in results]
+    assert downloaded_urls == ["http://mywebsite.com/file.pdf", "http://mywebsite.com/file.doc"]
+    assert mocked_download_response.call_count == 2
+    assert mocked_write_content.call_count == 2

test/unit/v2/embedders/test_voyageai.py

@@ -14,7 +14,7 @@ def generate_embedder_config_params() -> dict:
     }
     if random.random() < 0.5:
         params["embedder_model_name"] = fake.word()
-        params["batch_size"] = fake.random_int()
+        params["batch_size"] = fake.random_int(max=100)
         params["truncation"] = fake.boolean()
         params["max_retries"] = fake.random_int()
         params["timeout_in_seconds"] = fake.random_int()

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.2"  # pragma: no cover
+__version__ = "0.4.3"  # pragma: no cover

unstructured_ingest/embed/huggingface.py

@@ -33,6 +33,11 @@ class HuggingFaceEmbeddingConfig(EmbeddingConfig):
             **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):
@@ -43,7 +48,7 @@ class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
 
     def _embed_documents(self, texts: list[str]) -> list[list[float]]:
         client = self.config.get_client()
-        embeddings = client.encode(texts, **self.config.encode_kwargs)
+        embeddings = client.encode(texts, **self.config.get_encoder_kwargs())
         return embeddings.tolist()
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:

unstructured_ingest/embed/interfaces.py

@@ -1,13 +1,16 @@
 import asyncio
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import Optional
 
 import numpy as np
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class EmbeddingConfig(BaseModel):
-    pass
+    batch_size: Optional[int] = Field(
+        default=32, description="Optional batch size for embedding requests."
+    )
 
 
 @dataclass
@@ -53,9 +56,9 @@ class BaseEmbeddingEncoder(BaseEncoder, ABC):
         is properly configured: e.g., embed a single a element"""
 
     @property
-    def num_of_dimensions(self) -> tuple[int, ...]:
+    def dimension(self):
         exemplary_embedding = self.get_exemplary_embedding()
-        return np.shape(exemplary_embedding)
+        return len(exemplary_embedding)
 
     def get_exemplary_embedding(self) -> list[float]:
         return self.embed_query(query="Q")
@@ -91,9 +94,9 @@ class AsyncBaseEmbeddingEncoder(BaseEncoder, ABC):
         is properly configured: e.g., embed a single a element"""
 
     @property
-    async def num_of_dimensions(self) -> tuple[int, ...]:
+    async def dimension(self):
         exemplary_embedding = await self.get_exemplary_embedding()
-        return np.shape(exemplary_embedding)
+        return len(exemplary_embedding)
 
     async def get_exemplary_embedding(self) -> list[float]:
         return await self.embed_query(query="Q")

unstructured_ingest/embed/mixedbreadai.py

@@ -10,10 +10,10 @@ from unstructured_ingest.embed.interfaces import (
     BaseEmbeddingEncoder,
     EmbeddingConfig,
 )
+from unstructured_ingest.utils.data_prep import batch_generator
 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"
@@ -109,13 +109,10 @@ class MixedbreadAIEmbeddingEncoder(BaseEmbeddingEncoder):
         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]
+        for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
             response = client.embeddings(
                 model=self.config.embedder_model_name,
                 normalized=True,
@@ -186,13 +183,9 @@ class AsyncMixedbreadAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
         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]
+        for batch in batch_generator(texts, batch_size=self.config.batch_size or len(texts)):
             tasks.append(
                 client.embeddings(
                     model=self.config.embedder_model_name,

unstructured_ingest/embed/octoai.py

@@ -9,6 +9,7 @@ from unstructured_ingest.embed.interfaces import (
     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,
@@ -89,12 +90,16 @@ class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
 
     def embed_documents(self, elements: list[dict]) -> list[dict]:
         texts = [e.get("text", "") for e in elements]
+        embeddings = []
+        client = self.config.get_client()
         try:
-            client = self.config.get_client()
-            response = client.embeddings.create(input=texts, model=self.config.embedder_model_name)
+            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)
-        embeddings = [data.embedding for data in response.data]
         elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
         return elements_with_embeddings
 
@@ -119,12 +124,14 @@ class AsyncOctoAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
     async def embed_documents(self, elements: list[dict]) -> list[dict]:
         texts = [e.get("text", "") for e in elements]
         client = self.config.get_async_client()
+        embeddings = []
         try:
-            response = await client.embeddings.create(
-                input=texts, model=self.config.embedder_model_name
-            )
+            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)
-        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

@@ -9,6 +9,7 @@ from unstructured_ingest.embed.interfaces import (
     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,
@@ -80,7 +81,17 @@ class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
         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])
+        client = self.config.get_client()
+        texts = [e.get("text", "") for e in elements]
+        embeddings = []
+        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)
         elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
         return elements_with_embeddings
 
@@ -105,12 +116,14 @@ class AsyncOpenAIEmbeddingEncoder(AsyncBaseEmbeddingEncoder):
     async def embed_documents(self, elements: list[dict]) -> list[dict]:
         client = self.config.get_async_client()
         texts = [e.get("text", "") for e in elements]
+        embeddings = []
         try:
-            response = await client.embeddings.create(
-                input=texts, model=self.config.embedder_model_name
-            )
+            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)
-        embeddings = [data.embedding for data in response.data]
         elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
         return elements_with_embeddings