feat: Enhance error handling in Azure document embedder

haystack/components/embedders/azure_document_embedder.py

@@ -8,9 +8,11 @@
 from openai.lib.azure import AzureOpenAI
 from tqdm import tqdm
 
-from haystack import Document, component, default_from_dict, default_to_dict
+from haystack import Document, component, default_from_dict, default_to_dict, logging
 from haystack.utils import Secret, deserialize_secrets_inplace
 
+logger = logging.getLogger(__name__)
+
 
 @component
 class AzureOpenAIDocumentEmbedder:
@@ -208,24 +210,31 @@ def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> Tuple[List
         meta: Dict[str, Any] = {"model": "", "usage": {"prompt_tokens": 0, "total_tokens": 0}}
         for i in tqdm(range(0, len(texts_to_embed), batch_size), desc="Embedding Texts"):
             batch = texts_to_embed[i : i + batch_size]
-            if self.dimensions is not None:
-                response = self._client.embeddings.create(
-                    model=self.azure_deployment, dimensions=self.dimensions, input=batch
-                )
-            else:
-                response = self._client.embeddings.create(model=self.azure_deployment, input=batch)
-
-            # Append embeddings to the list
-            all_embeddings.extend(el.embedding for el in response.data)
-
-            # Update the meta information only once if it's empty
-            if not meta["model"]:
-                meta["model"] = response.model
-                meta["usage"] = dict(response.usage)
-            else:
-                # Update the usage tokens
-                meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
-                meta["usage"]["total_tokens"] += response.usage.total_tokens
+            try:
+                if self.dimensions is not None:
+                    response = self._client.embeddings.create(
+                        model=self.azure_deployment, dimensions=self.dimensions, input=batch
+                    )
+                else:
+                    response = self._client.embeddings.create(model=self.azure_deployment, input=batch)
+
+                # Append embeddings to the list
+                all_embeddings.extend(el.embedding for el in response.data)
+
+                # Update the meta information only once if it's empty
+                if not meta["model"]:
+                    meta["model"] = response.model
+                    meta["usage"] = dict(response.usage)
+                else:
+                    # Update the usage tokens
+                    meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
+                    meta["usage"]["total_tokens"] += response.usage.total_tokens
+
+            except Exception as e:
+                # Log the error but continue processing
+                batch_range = f"{i} - {i + batch_size}"
+                logger.exception(f"Failed embedding of documents in range: {batch_range} caused by {e}")
+                continue
 
         return all_embeddings, meta
 

releasenotes/notes/add-azure-embedder-exception-handler-c10ea46fb536de3b.yaml

@@ -0,0 +1,6 @@
+---
+enhancements:
+  - |
+    Handle Azure OpenAI embedder exceptions gracefully
+    - Add error handling in _embed_batch inside `AzureOpenAIDocumentEmbedder` to continue processing remaining documents
+    - Add unit tests for graceful error handling

test/components/embedders/test_azure_document_embedder.py

@@ -3,10 +3,14 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 
+from openai import APIError
+
+from haystack.utils.auth import Secret
 import pytest
 
 from haystack import Document
 from haystack.components.embedders import AzureOpenAIDocumentEmbedder
+from unittest.mock import Mock, patch
 
 
 class TestAzureOpenAIDocumentEmbedder:
@@ -83,6 +87,26 @@ def test_from_dict(self, monkeypatch):
         assert component.suffix == ""
         assert component.default_headers == {}
 
+    def test_embed_batch_handles_exceptions_gracefully(self, caplog):
+        embedder = AzureOpenAIDocumentEmbedder(
+            azure_endpoint="https://test.openai.azure.com",
+            api_key=Secret.from_token("fake-api-key"),
+            azure_deployment="text-embedding-ada-002",
+            embedding_separator=" | ",
+        )
+
+        fake_texts_to_embed = ["Cuisine | I love cheese", "ML | A transformer is a deep learning architecture"]
+
+        with patch.object(
+            embedder._client.embeddings,
+            "create",
+            side_effect=APIError(message="Mocked error", request=Mock(), body=None),
+        ):
+            embedder._embed_batch(texts_to_embed=fake_texts_to_embed, batch_size=32)
+
+        assert len(caplog.records) == 1
+        assert "Failed embedding of documents in range: 0 - 32 caused by Mocked error" in caplog.text
+
     @pytest.mark.integration
     @pytest.mark.skipif(
         not os.environ.get("AZURE_OPENAI_API_KEY", None) and not os.environ.get("AZURE_OPENAI_ENDPOINT", None),