Repository: deepset-ai/haystack

Issue:
Add run_async for `AzureOpenAIDocumentEmbedder`
We should be able to reuse the implementation once it is made for the `OpenAIDocumentEmbedder`


Discussion:

Patch:
diff --git a/haystack/components/embedders/azure_document_embedder.py b/haystack/components/embedders/azure_document_embedder.py
index 5f9eb647b4..ae08b86781 100644
--- a/haystack/components/embedders/azure_document_embedder.py
+++ b/haystack/components/embedders/azure_document_embedder.py
@@ -3,22 +3,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import httpx
-from more_itertools import batched
-from openai import APIError
-from openai.lib.azure import AzureADTokenProvider, AzureOpenAI
-from tqdm import tqdm
+from openai.lib.azure import AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI
 
-from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.components.embedders import OpenAIDocumentEmbedder
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
 
 logger = logging.getLogger(__name__)
 
 
 @component
-class AzureOpenAIDocumentEmbedder:
+class AzureOpenAIDocumentEmbedder(OpenAIDocumentEmbedder):
     """
     Calculates document embeddings using OpenAI models deployed on Azure.
 
@@ -39,6 +37,7 @@ class AzureOpenAIDocumentEmbedder:
     ```
     """
 
+    # pylint: disable=super-init-not-called
     def __init__(  # noqa: PLR0913 (too-many-arguments) # pylint: disable=too-many-positional-arguments
         self,
         azure_endpoint: Optional[str] = None,
@@ -109,6 +108,9 @@ def __init__(  # noqa: PLR0913 (too-many-arguments) # pylint: disable=too-many-p
             every request.
         :param http_client_kwargs: A dictionary of keyword arguments to configure a custom httpx.Client.
         """
+        # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact
+        # with the API.
+
         # if not provided as a parameter, azure_endpoint is read from the env var AZURE_OPENAI_ENDPOINT
         azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
         if not azure_endpoint:
@@ -117,11 +119,12 @@ def __init__(  # noqa: PLR0913 (too-many-arguments) # pylint: disable=too-many-p
         if api_key is None and azure_ad_token is None:
             raise ValueError("Please provide an API key or an Azure Active Directory token.")
 
-        self.api_key = api_key
+        self.api_key = api_key  # type: ignore[assignment] # mypy does not understand that api_key can be None
         self.azure_ad_token = azure_ad_token
         self.api_version = api_version
         self.azure_endpoint = azure_endpoint
         self.azure_deployment = azure_deployment
+        self.model = azure_deployment
         self.dimensions = dimensions
         self.organization = organization
         self.prefix = prefix
@@ -136,33 +139,31 @@ def __init__(  # noqa: PLR0913 (too-many-arguments) # pylint: disable=too-many-p
         self.azure_ad_token_provider = azure_ad_token_provider
         self.http_client_kwargs = http_client_kwargs
 
-        self._client = AzureOpenAI(
-            api_version=api_version,
-            azure_endpoint=azure_endpoint,
-            azure_deployment=azure_deployment,
-            azure_ad_token_provider=azure_ad_token_provider,
-            api_key=api_key.resolve_value() if api_key is not None else None,
-            azure_ad_token=azure_ad_token.resolve_value() if azure_ad_token is not None else None,
-            organization=organization,
-            timeout=self.timeout,
-            max_retries=self.max_retries,
-            default_headers=self.default_headers,
-            http_client=self._init_http_client(),
-        )
-
-    def _init_http_client(self):
+        client_args: Dict[str, Any] = {
+            "api_version": api_version,
+            "azure_endpoint": azure_endpoint,
+            "azure_deployment": azure_deployment,
+            "azure_ad_token_provider": azure_ad_token_provider,
+            "api_key": api_key.resolve_value() if api_key is not None else None,
+            "azure_ad_token": azure_ad_token.resolve_value() if azure_ad_token is not None else None,
+            "organization": organization,
+            "timeout": self.timeout,
+            "max_retries": self.max_retries,
+            "default_headers": self.default_headers,
+        }
+
+        self.client = AzureOpenAI(http_client=self._init_http_client(async_client=False), **client_args)
+        self.async_client = AsyncAzureOpenAI(http_client=self._init_http_client(async_client=True), **client_args)
+
+    def _init_http_client(self, async_client: bool = False):
         """Internal method to initialize the httpx.Client."""
-        if self.http_client_kwargs:
-            if not isinstance(self.http_client_kwargs, dict):
-                raise TypeError("The parameter 'http_client_kwargs' must be a dictionary.")
-            return httpx.Client(**self.http_client_kwargs)
-        return None
-
-    def _get_telemetry_data(self) -> Dict[str, Any]:
-        """
-        Data that is sent to Posthog for usage analytics.
-        """
-        return {"model": self.azure_deployment}
+        if not self.http_client_kwargs:
+            return None
+        if not isinstance(self.http_client_kwargs, dict):
+            raise TypeError("The parameter 'http_client_kwargs' must be a dictionary.")
+        if async_client:
+            return httpx.AsyncClient(**self.http_client_kwargs)
+        return httpx.Client(**self.http_client_kwargs)
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -213,83 +214,3 @@ def from_dict(cls, data: Dict[str, Any]) -> "AzureOpenAIDocumentEmbedder":
                 serialized_azure_ad_token_provider
             )
         return default_from_dict(cls, data)
-
-    def _prepare_texts_to_embed(self, documents: List[Document]) -> Dict[str, str]:
-        """
-        Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
-        """
-        texts_to_embed = {}
-        for doc in documents:
-            meta_values_to_embed = [
-                str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
-            ]
-
-            text_to_embed = (
-                self.prefix + self.embedding_separator.join(meta_values_to_embed + [doc.content or ""]) + self.suffix
-            ).replace("\n", " ")
-
-            texts_to_embed[doc.id] = text_to_embed
-        return texts_to_embed
-
-    def _embed_batch(self, texts_to_embed: Dict[str, str], batch_size: int) -> Tuple[List[List[float]], Dict[str, Any]]:
-        """
-        Embed a list of texts in batches.
-        """
-
-        all_embeddings: List[List[float]] = []
-        meta: Dict[str, Any] = {"model": "", "usage": {"prompt_tokens": 0, "total_tokens": 0}}
-
-        for batch in tqdm(
-            batched(texts_to_embed.items(), batch_size), disable=not self.progress_bar, desc="Calculating embeddings"
-        ):
-            args: Dict[str, Any] = {"model": self.azure_deployment, "input": [b[1] for b in batch]}
-
-            if self.dimensions is not None:
-                args["dimensions"] = self.dimensions
-
-            try:
-                response = self._client.embeddings.create(**args)
-            except APIError as e:
-                # Log the error but continue processing
-                ids = ", ".join(b[0] for b in batch)
-                logger.exception(f"Failed embedding of documents {ids} caused by {e}")
-                continue
-
-            embeddings = [el.embedding for el in response.data]
-            all_embeddings.extend(embeddings)
-
-            # 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
-
-        return all_embeddings, meta
-
-    @component.output_types(documents=List[Document], meta=Dict[str, Any])
-    def run(self, documents: List[Document]) -> Dict[str, Any]:
-        """
-        Embeds a list of documents.
-
-        :param documents:
-            Documents to embed.
-
-        :returns:
-            A dictionary with the following keys:
-            - `documents`: A list of documents with embeddings.
-            - `meta`: Information about the usage of the model.
-        """
-        if not (isinstance(documents, list) and all(isinstance(doc, Document) for doc in documents)):
-            raise TypeError("Input must be a list of Document instances. For strings, use AzureOpenAITextEmbedder.")
-
-        texts_to_embed = self._prepare_texts_to_embed(documents=documents)
-        embeddings, meta = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
-
-        # Assign the corresponding embeddings to each document
-        for doc, emb in zip(documents, embeddings):
-            doc.embedding = emb
-
-        return {"documents": documents, "meta": meta}
diff --git a/haystack/components/embedders/openai_document_embedder.py b/haystack/components/embedders/openai_document_embedder.py
index 92629e1bad..39952578a2 100644
--- a/haystack/components/embedders/openai_document_embedder.py
+++ b/haystack/components/embedders/openai_document_embedder.py
@@ -175,13 +175,10 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> Dict[str, str]:
                 str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
             ]
 
-            text_to_embed = (
+            texts_to_embed[doc.id] = (
                 self.prefix + self.embedding_separator.join(meta_values_to_embed + [doc.content or ""]) + self.suffix
             )
 
-            # copied from OpenAI embedding_utils (https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py)
-            # replace newlines, which can negatively affect performance.
-            texts_to_embed[doc.id] = text_to_embed.replace("\n", " ")
         return texts_to_embed
 
     def _embed_batch(self, texts_to_embed: Dict[str, str], batch_size: int) -> Tuple[List[List[float]], Dict[str, Any]]:
diff --git a/releasenotes/notes/azure-doc-embedder-inherit-1028a68e99561ab2.yaml b/releasenotes/notes/azure-doc-embedder-inherit-1028a68e99561ab2.yaml
new file mode 100644
index 0000000000..23189aa191
--- /dev/null
+++ b/releasenotes/notes/azure-doc-embedder-inherit-1028a68e99561ab2.yaml
@@ -0,0 +1,10 @@
+---
+
+features:
+  - |
+    The `AzureOpenAIDocumentEmbedder` component now inherits from the `OpenAIDocumentEmbedder` component,
+    enabling asynchronous usage.
+fixes:
+  - |
+    `OpenAIDocumentEmbedder` and `AzureOpenAIDocumentEmbedder` no longer replace newlines with spaces in the text to
+    embed. This was only required for the discontinued v1 embedding models.


Test patch:
diff --git a/test/components/embedders/test_azure_document_embedder.py b/test/components/embedders/test_azure_document_embedder.py
index accc8d254e..52fa72355e 100644
--- a/test/components/embedders/test_azure_document_embedder.py
+++ b/test/components/embedders/test_azure_document_embedder.py
@@ -7,6 +7,7 @@
 
 from haystack.utils.auth import Secret
 import pytest
+import httpx
 
 from haystack import Document
 from haystack.components.embedders import AzureOpenAIDocumentEmbedder
@@ -19,6 +20,7 @@ def test_init_default(self, monkeypatch):
         monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
         embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         assert embedder.azure_deployment == "text-embedding-ada-002"
+        assert embedder.model == "text-embedding-ada-002"
         assert embedder.dimensions is None
         assert embedder.organization is None
         assert embedder.prefix == ""
@@ -38,6 +40,7 @@ def test_init_with_0_max_retries(self, monkeypatch):
             azure_endpoint="https://example-resource.azure.openai.com/", max_retries=0
         )
         assert embedder.azure_deployment == "text-embedding-ada-002"
+        assert embedder.model == "text-embedding-ada-002"
         assert embedder.dimensions is None
         assert embedder.organization is None
         assert embedder.prefix == ""
@@ -203,7 +206,7 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
         fake_texts_to_embed = {"1": "text1", "2": "text2"}
 
         with patch.object(
-            embedder._client.embeddings,
+            embedder.client.embeddings,
             "create",
             side_effect=APIError(message="Mocked error", request=Mock(), body=None),
         ):
@@ -212,6 +215,21 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
         assert len(caplog.records) == 1
         assert "Failed embedding of documents 1, 2 caused by Mocked error" in caplog.text
 
+    def test_init_http_client(self, monkeypatch):
+        monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
+        monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test.openai.azure.com")
+
+        embedder = AzureOpenAIDocumentEmbedder()
+        client = embedder._init_http_client()
+        assert client is None
+
+        embedder.http_client_kwargs = {"proxy": "http://example.com:3128"}
+        client = embedder._init_http_client(async_client=False)
+        assert isinstance(client, httpx.Client)
+
+        client = embedder._init_http_client(async_client=True)
+        assert isinstance(client, httpx.AsyncClient)
+
     def test_http_client_kwargs_type_validation(self, monkeypatch):
         monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
         monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test.openai.azure.com")
@@ -257,4 +275,7 @@ def test_run(self):
             assert isinstance(doc.embedding, list)
             assert len(doc.embedding) == 1536
             assert all(isinstance(x, float) for x in doc.embedding)
-        assert metadata == {"model": "text-embedding-ada-002", "usage": {"prompt_tokens": 15, "total_tokens": 15}}
+
+        assert metadata["usage"]["prompt_tokens"] == 15
+        assert metadata["usage"]["total_tokens"] == 15
+        assert "ada" in metadata["model"]
diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py
index 1e85018732..77c5259efe 100644
--- a/test/components/embedders/test_openai_document_embedder.py
+++ b/test/components/embedders/test_openai_document_embedder.py
@@ -167,13 +167,12 @@ def test_prepare_texts_to_embed_w_metadata(self):
 
         prepared_texts = embedder._prepare_texts_to_embed(documents)
 
-        # note that newline is replaced by space
         assert prepared_texts == {
-            "0": "meta_value 0 | document number 0: content",
-            "1": "meta_value 1 | document number 1: content",
-            "2": "meta_value 2 | document number 2: content",
-            "3": "meta_value 3 | document number 3: content",
-            "4": "meta_value 4 | document number 4: content",
+            "0": "meta_value 0 | document number 0:\ncontent",
+            "1": "meta_value 1 | document number 1:\ncontent",
+            "2": "meta_value 2 | document number 2:\ncontent",
+            "3": "meta_value 3 | document number 3:\ncontent",
+            "4": "meta_value 4 | document number 4:\ncontent",
         }
 
     def test_prepare_texts_to_embed_w_suffix(self):