Coverage for haystack/components/embedders/azure_text_embedder.py: 97%
63 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
11from haystack.components.embedders import OpenAITextEmbedder
12from haystack.utils import Secret, deserialize_callable, serialize_callable
13from haystack.utils.http_client import init_http_client
16@component
17class AzureOpenAITextEmbedder(OpenAITextEmbedder):
18 """
19 Embeds strings using OpenAI models deployed on Azure.
21 ### Usage example
22 <!-- test-ignore -->
23 ```python
24 from haystack.components.embedders import AzureOpenAITextEmbedder
26 text_to_embed = "I love pizza!"
27 text_embedder = AzureOpenAITextEmbedder()
29 print(text_embedder.run(text_to_embed))
31 # {'embedding': [0.017020374536514282, -0.023255806416273117, ...],
32 # 'meta': {'model': 'text-embedding-ada-002-v2',
33 # 'usage': {'prompt_tokens': 4, 'total_tokens': 4}}}
34 ```
35 """
37 def __init__( # noqa: PLR0913
38 self,
39 azure_endpoint: str | None = None,
40 api_version: str | None = "2023-05-15",
41 azure_deployment: str = "text-embedding-ada-002",
42 dimensions: int | None = None,
43 api_key: Secret | None = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False),
44 azure_ad_token: Secret | None = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False),
45 organization: str | None = None,
46 timeout: float | None = None,
47 max_retries: int | None = None,
48 prefix: str = "",
49 suffix: str = "",
50 *,
51 default_headers: dict[str, str] | None = None,
52 azure_ad_token_provider: AzureADTokenProvider | None = None,
53 http_client_kwargs: dict[str, Any] | None = None,
54 ) -> None:
55 """
56 Creates an AzureOpenAITextEmbedder component.
58 :param azure_endpoint:
59 The endpoint of the model deployed on Azure.
60 :param api_version:
61 The version of the API to use.
62 :param azure_deployment:
63 The name of the model deployed on Azure. The default model is text-embedding-ada-002.
64 :param dimensions:
65 The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3
66 and later models.
67 :param api_key:
68 The Azure OpenAI API key.
69 You can set it with an environment variable `AZURE_OPENAI_API_KEY`, or pass with this
70 parameter during initialization.
71 :param azure_ad_token:
72 Microsoft Entra ID token, see Microsoft's
73 [Entra ID](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id)
74 documentation for more information. You can set it with an environment variable
75 `AZURE_OPENAI_AD_TOKEN`, or pass with this parameter during initialization.
76 Previously called Azure Active Directory.
77 :param organization:
78 Your organization ID. See OpenAI's
79 [Setting Up Your Organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
80 for more information.
81 :param timeout: The timeout for `AzureOpenAI` client calls, in seconds.
82 If not set, defaults to either the
83 `OPENAI_TIMEOUT` environment variable, or 30 seconds.
84 :param max_retries: Maximum number of retries to contact AzureOpenAI after an internal error.
85 If not set, defaults to either the `OPENAI_MAX_RETRIES` environment variable, or to 5 retries.
86 :param prefix:
87 A string to add at the beginning of each text.
88 :param suffix:
89 A string to add at the end of each text.
90 :param default_headers: Default headers to send to the AzureOpenAI client.
91 :param azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on
92 every request.
93 :param http_client_kwargs:
94 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
95 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
97 """
98 # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact
99 # with the API.
101 # Why is this here?
102 # AzureOpenAI init is forcing us to use an init method that takes either base_url or azure_endpoint as not
103 # None init parameters. This way we accommodate the use case where env var AZURE_OPENAI_ENDPOINT is set instead
104 # of passing it as a parameter.
105 azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
106 if not azure_endpoint:
107 raise ValueError("Please provide an Azure endpoint or set the environment variable AZURE_OPENAI_ENDPOINT.")
109 if api_key is None and azure_ad_token is None:
110 raise ValueError("Please provide an API key or an Azure Active Directory token.")
112 self.api_key = api_key # type: ignore[assignment] # mypy does not understand that api_key can be None
113 self.azure_ad_token = azure_ad_token
114 self.api_version = api_version
115 self.azure_endpoint = azure_endpoint
116 self.azure_deployment = azure_deployment
117 self.model = azure_deployment
118 self.dimensions = dimensions
119 self.organization = organization
120 self.timeout = timeout
121 self.max_retries = max_retries
122 self.prefix = prefix
123 self.suffix = suffix
124 self.default_headers = default_headers or {}
125 self.azure_ad_token_provider = azure_ad_token_provider
126 self.http_client_kwargs = http_client_kwargs
128 self.client: AzureOpenAI | None = None
129 self.async_client: AsyncAzureOpenAI | None = None
131 def _client_kwargs(self) -> dict[str, Any]:
132 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
133 max_retries = (
134 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
135 )
136 return {
137 "api_version": self.api_version,
138 "azure_endpoint": self.azure_endpoint,
139 "azure_deployment": self.azure_deployment,
140 "azure_ad_token_provider": self.azure_ad_token_provider,
141 "api_key": self.api_key.resolve_value() if self.api_key is not None else None,
142 "azure_ad_token": self.azure_ad_token.resolve_value() if self.azure_ad_token is not None else None,
143 "organization": self.organization,
144 "timeout": timeout,
145 "max_retries": max_retries,
146 "default_headers": self.default_headers,
147 }
149 def warm_up(self) -> None:
150 """
151 Initializes the synchronous Azure OpenAI client.
152 """
153 if self.client is None:
154 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
155 # https://github.com/openai/openai-python/blob/main/httpx2.md
156 http_client = init_http_client(self.http_client_kwargs, async_client=False)
157 self.client = AzureOpenAI(
158 http_client=http_client, # type: ignore[arg-type]
159 **self._client_kwargs(),
160 )
162 async def warm_up_async(self) -> None: # noqa: RUF029
163 """
164 Initializes the asynchronous Azure OpenAI client on the serving event loop.
165 """
166 if self.async_client is None:
167 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
168 # https://github.com/openai/openai-python/blob/main/httpx2.md
169 http_client = init_http_client(self.http_client_kwargs, async_client=True)
170 self.async_client = AsyncAzureOpenAI(
171 http_client=http_client, # type: ignore[arg-type]
172 **self._client_kwargs(),
173 )
175 def close(self) -> None:
176 """
177 Releases the synchronous Azure OpenAI client.
178 """
179 if self.client is not None:
180 self.client.close()
181 self.client = None
183 async def close_async(self) -> None:
184 """
185 Releases the asynchronous Azure OpenAI client.
186 """
187 if self.async_client is not None:
188 await self.async_client.close()
189 self.async_client = None
191 def to_dict(self) -> dict[str, Any]:
192 """
193 Serializes the component to a dictionary.
195 :returns:
196 Dictionary with serialized data.
197 """
198 azure_ad_token_provider_name = None
199 if self.azure_ad_token_provider:
200 azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider)
201 return default_to_dict(
202 self,
203 azure_endpoint=self.azure_endpoint,
204 azure_deployment=self.azure_deployment,
205 dimensions=self.dimensions,
206 organization=self.organization,
207 api_version=self.api_version,
208 prefix=self.prefix,
209 suffix=self.suffix,
210 api_key=self.api_key,
211 azure_ad_token=self.azure_ad_token,
212 timeout=self.timeout,
213 max_retries=self.max_retries,
214 default_headers=self.default_headers,
215 azure_ad_token_provider=azure_ad_token_provider_name,
216 http_client_kwargs=self.http_client_kwargs,
217 )
219 @classmethod
220 def from_dict(cls, data: dict[str, Any]) -> "AzureOpenAITextEmbedder":
221 """
222 Deserializes the component from a dictionary.
224 :param data:
225 Dictionary to deserialize from.
226 :returns:
227 Deserialized component.
228 """
229 serialized_azure_ad_token_provider = data["init_parameters"].get("azure_ad_token_provider")
230 if serialized_azure_ad_token_provider:
231 data["init_parameters"]["azure_ad_token_provider"] = deserialize_callable(
232 serialized_azure_ad_token_provider
233 )
234 return default_from_dict(cls, data)