migrate to new mixedbread python sdk (#492)

* migrate to new mixedbread python sdk

* fix lint

* update changelog

* update mixedbread unit test

* update dependency in pyproject

* fix fixedbread unit test

Author: bryan-unstructured

CHANGELOG.md

@@ -1,13 +1,14 @@
+## 1.0.22-dev1
+
+* **Migrate to new Mixedbread Python SDK**
+
 ## 1.0.22
 
 * **Fix Notion connector missing database properties fields**
 
 ## 1.0.21
 
 * **Fix Jira connector cloud option not working issue**
-
-## 1.0.20
-
 * **Fix Weaviate connector issue with names being wrongly transformed to match collections naming conventions**
 
 ## 1.0.19

requirements/embed/mixedbreadai.txt

@@ -1,3 +1,3 @@
 -c ../common/constraints.txt
 
-mixedbread-ai
+mixedbread

test/integration/embedders/test_mixedbread.py

@@ -27,7 +27,6 @@ def get_api_key() -> str:
 
 
 @requires_env(API_KEY)
-@pytest.mark.skip(reason="need to migrate to new mixedbread sdk.")
 def test_mixedbread_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="mixedbread-ai", embedding_api_key=api_key)
@@ -41,7 +40,6 @@ def test_mixedbread_embedder(embedder_file: Path):
 
 
 @requires_env(API_KEY)
-@pytest.mark.skip(reason="need to migrate to new mixedbread sdk.")
 def test_raw_mixedbread_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder = MixedbreadAIEmbeddingEncoder(
@@ -60,7 +58,6 @@ def test_raw_mixedbread_embedder(embedder_file: Path):
 
 @requires_env(API_KEY)
 @pytest.mark.asyncio
-@pytest.mark.skip(reason="need to migrate to new mixedbread sdk.")
 async def test_raw_async_mixedbread_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder = AsyncMixedbreadAIEmbeddingEncoder(

test/unit/embed/test_mixedbreadai.py

@@ -7,23 +7,23 @@
 def test_embed_documents_does_not_break_element_to_dict(mocker):
     mock_client = mocker.MagicMock()
 
-    def mock_embeddings(
+    def mock_embed(
         model,
+        input,
         normalized,
         encoding_format,
-        truncation_strategy,
-        request_options,
-        input,
+        extra_headers,
+        timeout,
     ):
         mock_response = mocker.MagicMock()
         mock_response.data = [mocker.MagicMock(embedding=[i, i + 1]) for i in range(len(input))]
         return mock_response
 
-    mock_client.embeddings.side_effect = mock_embeddings
+    mock_client.embed.side_effect = mock_embed
 
     # Mock get_client to return our mock_client
     mocker.patch.object(MixedbreadAIEmbeddingConfig, "get_client", return_value=mock_client)
-    mocker.patch.object(MixedbreadAIEmbeddingEncoder, "get_request_options", return_value={})
+    mocker.patch.object(MixedbreadAIEmbeddingEncoder, "get_client", return_value=mock_client)
 
     encoder = MixedbreadAIEmbeddingEncoder(
         config=MixedbreadAIEmbeddingConfig(

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.22"  # pragma: no cover
+__version__ = "1.0.22-dev1"  # pragma: no cover

unstructured_ingest/embed/mixedbreadai.py

@@ -19,8 +19,7 @@
 
 
 if TYPE_CHECKING:
-    from mixedbread_ai.client import AsyncMixedbreadAI, MixedbreadAI
-    from mixedbread_ai.core import RequestOptions
+    from mixedbread import AsyncMixedbread, Mixedbread
 
 
 class MixedbreadAIEmbeddingConfig(EmbeddingConfig):
@@ -44,31 +43,33 @@ class MixedbreadAIEmbeddingConfig(EmbeddingConfig):
     )
 
     @requires_dependencies(
-        ["mixedbread_ai"],
-        extras="mixedbreadai",
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
     )
-    def get_client(self) -> "MixedbreadAI":
+    def get_client(self) -> "Mixedbread":
         """
         Create the Mixedbread AI client.
 
         Returns:
-            MixedbreadAI: Initialized client.
+            Mixedbread: Initialized client.
         """
-        from mixedbread_ai.client import MixedbreadAI
+        from mixedbread import Mixedbread
 
-        return MixedbreadAI(
+        return Mixedbread(
             api_key=self.api_key.get_secret_value(),
+            max_retries=MAX_RETRIES,
         )
 
     @requires_dependencies(
-        ["mixedbread_ai"],
-        extras="mixedbreadai",
+        ["mixedbread"],
+        extras="embed-mixedbreadai",
     )
-    def get_async_client(self) -> "AsyncMixedbreadAI":
-        from mixedbread_ai.client import AsyncMixedbreadAI
+    def get_async_client(self) -> "AsyncMixedbread":
+        from mixedbread import AsyncMixedbread
 
-        return AsyncMixedbreadAI(
+        return AsyncMixedbread(
             api_key=self.api_key.get_secret_value(),
+            max_retries=MAX_RETRIES,
         )
 
 
@@ -88,29 +89,20 @@ def get_exemplary_embedding(self) -> list[float]:
         return self.embed_query(query="Q")
 
     @requires_dependencies(
-        ["mixedbread_ai"],
+        ["mixedbread"],
         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":
+    def get_client(self) -> "Mixedbread":
         return self.config.get_client()
 
-    def embed_batch(self, client: "MixedbreadAI", batch: list[str]) -> list[list[float]]:
-        response = client.embeddings(
+    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,
-            truncation_strategy=TRUNCATION_STRATEGY,
-            request_options=self.get_request_options(),
-            input=batch,
+            extra_headers={"User-Agent": USER_AGENT},
+            timeout=TIMEOUT,
         )
         return [datum.embedding for datum in response.data]
 
@@ -124,28 +116,19 @@ async def get_exemplary_embedding(self) -> list[float]:
         return await self.embed_query(query="Q")
 
     @requires_dependencies(
-        ["mixedbread_ai"],
+        ["mixedbread"],
         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":
+    def get_client(self) -> "AsyncMixedbread":
         return self.config.get_async_client()
 
-    async def embed_batch(self, client: "AsyncMixedbreadAI", batch: list[str]) -> list[list[float]]:
-        response = await client.embeddings(
+    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,
-            truncation_strategy=TRUNCATION_STRATEGY,
-            request_options=self.get_request_options(),
-            input=batch,
+            extra_headers={"User-Agent": USER_AGENT},
+            timeout=TIMEOUT,
         )
         return [datum.embedding for datum in response.data]