Coverage for haystack/components/embedders/azure_document_embedder.py: 97%
69 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import os
6from typing import Any
8from openai.lib.azure import AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI
10from haystack import component, default_from_dict, default_to_dict, logging
11from haystack.components.embedders import OpenAIDocumentEmbedder
12from haystack.utils import Secret, deserialize_callable, serialize_callable
13from haystack.utils.http_client import init_http_client
15logger = logging.getLogger(__name__)
18@component
19class AzureOpenAIDocumentEmbedder(OpenAIDocumentEmbedder):
20 """
21 Calculates document embeddings using OpenAI models deployed on Azure.
23 ### Usage example
24 <!-- test-ignore -->
25 ```python
26 from haystack import Document
27 from haystack.components.embedders import AzureOpenAIDocumentEmbedder
29 doc = Document(content="I love pizza!")
30 document_embedder = AzureOpenAIDocumentEmbedder()
32 result = document_embedder.run([doc])
33 print(result['documents'][0].embedding)
35 # [0.017020374536514282, -0.023255806416273117, ...]
36 ```
37 """
39 def __init__( # noqa: PLR0913, PLR0917 (too-many-arguments, too-many-positional-arguments)
40 self,
41 azure_endpoint: str | None = None,
42 api_version: str | None = "2023-05-15",
43 azure_deployment: str = "text-embedding-ada-002",
44 dimensions: int | None = None,
45 api_key: Secret | None = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False),
46 azure_ad_token: Secret | None = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False),
47 organization: str | None = None,
48 prefix: str = "",
49 suffix: str = "",
50 batch_size: int = 32,
51 progress_bar: bool = True,
52 meta_fields_to_embed: list[str] | None = None,
53 embedding_separator: str = "\n",
54 timeout: float | None = None,
55 max_retries: int | None = None,
56 *,
57 default_headers: dict[str, str] | None = None,
58 azure_ad_token_provider: AzureADTokenProvider | None = None,
59 http_client_kwargs: dict[str, Any] | None = None,
60 raise_on_failure: bool = False,
61 ) -> None:
62 """
63 Creates an AzureOpenAIDocumentEmbedder component.
65 :param azure_endpoint:
66 The endpoint of the model deployed on Azure.
67 :param api_version:
68 The version of the API to use.
69 :param azure_deployment:
70 The name of the model deployed on Azure. The default model is text-embedding-ada-002.
71 :param dimensions:
72 The number of dimensions of the resulting embeddings. Only supported in text-embedding-3
73 and later models.
74 :param api_key:
75 The Azure OpenAI API key.
76 You can set it with an environment variable `AZURE_OPENAI_API_KEY`, or pass with this
77 parameter during initialization.
78 :param azure_ad_token:
79 Microsoft Entra ID token, see Microsoft's
80 [Entra ID](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id)
81 documentation for more information. You can set it with an environment variable
82 `AZURE_OPENAI_AD_TOKEN`, or pass with this parameter during initialization.
83 Previously called Azure Active Directory.
84 :param organization:
85 Your organization ID. See OpenAI's
86 [Setting Up Your Organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
87 for more information.
88 :param prefix:
89 A string to add at the beginning of each text.
90 :param suffix:
91 A string to add at the end of each text.
92 :param batch_size:
93 Number of documents to embed at once.
94 :param progress_bar:
95 If `True`, shows a progress bar when running.
96 :param meta_fields_to_embed:
97 List of metadata fields to embed along with the document text.
98 :param embedding_separator:
99 Separator used to concatenate the metadata fields to the document text.
100 :param timeout: The timeout for `AzureOpenAI` client calls, in seconds.
101 If not set, defaults to either the
102 `OPENAI_TIMEOUT` environment variable, or 30 seconds.
103 :param max_retries: Maximum number of retries to contact AzureOpenAI after an internal error.
104 If not set, defaults to either the `OPENAI_MAX_RETRIES` environment variable or to 5 retries.
105 :param default_headers: Default headers to send to the AzureOpenAI client.
106 :param azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on
107 every request.
108 :param http_client_kwargs:
109 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
110 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
111 :param raise_on_failure:
112 Whether to raise an exception if the embedding request fails. If `False`, the component will log the error
113 and continue processing the remaining documents. If `True`, it will raise an exception on failure.
114 """
115 # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact
116 # with the API.
118 # if not provided as a parameter, azure_endpoint is read from the env var AZURE_OPENAI_ENDPOINT
119 azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
120 if not azure_endpoint:
121 raise ValueError("Please provide an Azure endpoint or set the environment variable AZURE_OPENAI_ENDPOINT.")
123 if api_key is None and azure_ad_token is None:
124 raise ValueError("Please provide an API key or an Azure Active Directory token.")
126 self.api_key = api_key # type: ignore[assignment] # mypy does not understand that api_key can be None
127 self.azure_ad_token = azure_ad_token
128 self.api_version = api_version
129 self.azure_endpoint = azure_endpoint
130 self.azure_deployment = azure_deployment
131 self.model = azure_deployment
132 self.dimensions = dimensions
133 self.organization = organization
134 self.prefix = prefix
135 self.suffix = suffix
136 self.batch_size = batch_size
137 self.progress_bar = progress_bar
138 self.meta_fields_to_embed = meta_fields_to_embed or []
139 self.embedding_separator = embedding_separator
140 self.timeout = timeout
141 self.max_retries = max_retries
142 self.default_headers = default_headers or {}
143 self.azure_ad_token_provider = azure_ad_token_provider
144 self.http_client_kwargs = http_client_kwargs
145 self.raise_on_failure = raise_on_failure
147 self.client: AzureOpenAI | None = None
148 self.async_client: AsyncAzureOpenAI | None = None
150 def _client_kwargs(self) -> dict[str, Any]:
151 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
152 max_retries = (
153 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
154 )
155 return {
156 "api_version": self.api_version,
157 "azure_endpoint": self.azure_endpoint,
158 "azure_deployment": self.azure_deployment,
159 "azure_ad_token_provider": self.azure_ad_token_provider,
160 "api_key": self.api_key.resolve_value() if self.api_key is not None else None,
161 "azure_ad_token": self.azure_ad_token.resolve_value() if self.azure_ad_token is not None else None,
162 "organization": self.organization,
163 "timeout": timeout,
164 "max_retries": max_retries,
165 "default_headers": self.default_headers,
166 }
168 def warm_up(self) -> None:
169 """
170 Initializes the synchronous AzureOpenAI client.
171 """
172 if self.client is None:
173 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
174 # https://github.com/openai/openai-python/blob/main/httpx2.md
175 http_client = init_http_client(self.http_client_kwargs, async_client=False)
176 self.client = AzureOpenAI(
177 http_client=http_client, # type: ignore[arg-type]
178 **self._client_kwargs(),
179 )
181 async def warm_up_async(self) -> None: # noqa: RUF029
182 """
183 Initializes the asynchronous AzureOpenAI client on the serving event loop.
184 """
185 if self.async_client is None:
186 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
187 # https://github.com/openai/openai-python/blob/main/httpx2.md
188 http_client = init_http_client(self.http_client_kwargs, async_client=True)
189 self.async_client = AsyncAzureOpenAI(
190 http_client=http_client, # type: ignore[arg-type]
191 **self._client_kwargs(),
192 )
194 def close(self) -> None:
195 """
196 Releases the synchronous AzureOpenAI client.
197 """
198 if self.client is not None:
199 self.client.close()
200 self.client = None
202 async def close_async(self) -> None:
203 """
204 Releases the asynchronous AzureOpenAI client.
205 """
206 if self.async_client is not None:
207 await self.async_client.close()
208 self.async_client = None
210 def to_dict(self) -> dict[str, Any]:
211 """
212 Serializes the component to a dictionary.
214 :returns:
215 Dictionary with serialized data.
216 """
217 azure_ad_token_provider_name = None
218 if self.azure_ad_token_provider:
219 azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider)
220 return default_to_dict(
221 self,
222 azure_endpoint=self.azure_endpoint,
223 azure_deployment=self.azure_deployment,
224 dimensions=self.dimensions,
225 organization=self.organization,
226 api_version=self.api_version,
227 prefix=self.prefix,
228 suffix=self.suffix,
229 batch_size=self.batch_size,
230 progress_bar=self.progress_bar,
231 meta_fields_to_embed=self.meta_fields_to_embed,
232 embedding_separator=self.embedding_separator,
233 api_key=self.api_key,
234 azure_ad_token=self.azure_ad_token,
235 timeout=self.timeout,
236 max_retries=self.max_retries,
237 default_headers=self.default_headers,
238 azure_ad_token_provider=azure_ad_token_provider_name,
239 http_client_kwargs=self.http_client_kwargs,
240 raise_on_failure=self.raise_on_failure,
241 )
243 @classmethod
244 def from_dict(cls, data: dict[str, Any]) -> "AzureOpenAIDocumentEmbedder":
245 """
246 Deserializes the component from a dictionary.
248 :param data:
249 Dictionary to deserialize from.
250 :returns:
251 Deserialized component.
252 """
253 serialized_azure_ad_token_provider = data["init_parameters"].get("azure_ad_token_provider")
254 if serialized_azure_ad_token_provider:
255 data["init_parameters"]["azure_ad_token_provider"] = deserialize_callable(
256 serialized_azure_ad_token_provider
257 )
258 return default_from_dict(cls, data)