Coverage for haystack/components/generators/chat/azure.py: 96%
85 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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, ClassVar
8from openai.lib._pydantic import to_strict_json_schema
9from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI
10from pydantic import BaseModel
12from haystack import component, default_from_dict, default_to_dict
13from haystack.components.generators.chat import OpenAIChatGenerator
14from haystack.dataclasses.streaming_chunk import StreamingCallbackT
15from haystack.tools import (
16 ToolsType,
17 _check_duplicate_tool_names,
18 deserialize_tools_or_toolset_inplace,
19 flatten_tools_or_toolsets,
20 serialize_tools_or_toolset,
21)
22from haystack.utils import Secret, deserialize_callable, serialize_callable
23from haystack.utils.http_client import init_http_client
26@component
27class AzureOpenAIChatGenerator(OpenAIChatGenerator):
28 """
29 Generates text using OpenAI's models on Azure.
31 It works with OpenAI's GPT models deployed on Azure and supports streaming responses
32 from OpenAI API. It uses [ChatMessage](https://docs.haystack.deepset.ai/docs/chatmessage)
33 format in input and output.
35 You can customize how the text is generated by passing parameters to the
36 OpenAI API. Use the `**generation_kwargs` argument when you initialize
37 the component or when you run it. Any parameter that works with
38 `openai.ChatCompletion.create` will work here too.
40 For details on OpenAI API parameters, see
41 [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat).
43 ### Usage example
44 <!-- test-ignore -->
45 ```python
46 from haystack.components.generators.chat import AzureOpenAIChatGenerator
47 from haystack.dataclasses import ChatMessage
48 from haystack.utils import Secret
50 messages = [ChatMessage.from_user("What's Natural Language Processing?")]
52 client = AzureOpenAIChatGenerator(
53 azure_endpoint="<Your Azure endpoint e.g. `https://your-company.azure.openai.com/>",
54 api_key=Secret.from_token("<your-api-key>"),
55 azure_deployment="<this is a model name, e.g. gpt-5.6-luna>")
56 response = client.run(messages)
57 print(response)
58 ```
60 ```
61 {'replies':
62 [ChatMessage(_role=<ChatRole.ASSISTANT: 'assistant'>, _content=[TextContent(text=
63 "Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on
64 enabling computers to understand, interpret, and generate human language in a way that is useful.")],
65 _name=None,
66 _meta={'model': 'gpt-5.6-luna', 'index': 0, 'finish_reason': 'stop',
67 'usage': {'prompt_tokens': 15, 'completion_tokens': 36, 'total_tokens': 51}})]
68 }
69 ```
70 """
72 SUPPORTED_MODELS: ClassVar[list[str]] = [
73 "gpt-5.4",
74 "gpt-5.4-pro",
75 "gpt-5.3-codex",
76 "gpt-5.2",
77 "gpt-5.2-codex",
78 "gpt-5.2-chat",
79 "gpt-5.1",
80 "gpt-5.1-chat",
81 "gpt-5.1-codex",
82 "gpt-5.1-codex-mini",
83 "gpt-5",
84 "gpt-5-mini",
85 "gpt-5-nano",
86 "gpt-5-chat",
87 "gpt-4.1",
88 "gpt-4.1-mini",
89 "gpt-4.1-nano",
90 "gpt-4o",
91 "gpt-4o-mini",
92 "gpt-4o-audio-preview",
93 "gpt-realtime-1.5",
94 "gpt-audio-1.5",
95 "o1",
96 "o1-mini",
97 "o3",
98 "o3-mini",
99 "o4-mini",
100 "codex-mini",
101 "gpt-4",
102 "gpt-35-turbo",
103 "gpt-oss-120b",
104 "computer-use-preview",
105 ]
106 """A non-exhaustive list of chat models supported by this component.
107 See https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure
108 for the full list."""
110 # ruff: noqa: PLR0913
111 def __init__(
112 self,
113 azure_endpoint: str | Secret | None = None,
114 api_version: str | Secret | None = "2024-12-01-preview",
115 azure_deployment: str | None = "gpt-4.1-mini",
116 api_key: Secret | None = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False),
117 azure_ad_token: Secret | None = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False),
118 organization: str | None = None,
119 streaming_callback: StreamingCallbackT | None = None,
120 timeout: float | None = None,
121 max_retries: int | None = None,
122 generation_kwargs: dict[str, Any] | None = None,
123 default_headers: dict[str, str] | None = None,
124 tools: ToolsType | None = None,
125 tools_strict: bool = False,
126 *,
127 azure_ad_token_provider: AzureADTokenProvider | AsyncAzureADTokenProvider | None = None,
128 http_client_kwargs: dict[str, Any] | None = None,
129 ) -> None:
130 """
131 Initialize the Azure OpenAI Chat Generator component.
133 :param azure_endpoint: The endpoint of the deployed model, for example `"https://example-resource.azure.openai.com/"`.
134 Can also be a [Secret](https://docs.haystack.deepset.ai/docs/secret-management), for example
135 `Secret.from_env_var("AZURE_OPENAI_ENDPOINT")`, to resolve the value from an environment variable at
136 runtime. This is useful to switch endpoints between environments (e.g. dev and prod) without changing the
137 serialized pipeline.
138 :param api_version: The version of the API to use. Defaults to 2024-12-01-preview.
139 Can also be a [Secret](https://docs.haystack.deepset.ai/docs/secret-management), for example
140 `Secret.from_env_var("AZURE_OPENAI_API_VERSION")`, to resolve the value from an environment variable at
141 runtime.
142 :param azure_deployment: The deployment of the model, usually the model name.
143 :param api_key: The API key to use for authentication.
144 :param azure_ad_token: [Azure Active Directory token](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id).
145 :param organization: Your organization ID, defaults to `None`. For help, see
146 [Setting up your organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization).
147 :param streaming_callback: A callback function called when a new token is received from the stream.
148 It accepts [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk)
149 as an argument.
150 :param timeout: Timeout for OpenAI client calls. If not set, it defaults to either the
151 `OPENAI_TIMEOUT` environment variable, or 30 seconds.
152 :param max_retries: Maximum number of retries to contact OpenAI after an internal error.
153 If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5.
154 :param generation_kwargs: Other parameters to use for the model. These parameters are sent directly to
155 the OpenAI endpoint. For details, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat).
156 Some of the supported parameters:
157 - `max_completion_tokens`: An upper bound for the number of tokens that can be generated for a completion,
158 including visible output tokens and reasoning tokens.
159 - `temperature`: The sampling temperature to use. Higher values mean the model takes more risks.
160 Try 0.9 for more creative applications and 0 (argmax sampling) for ones with a well-defined answer.
161 - `top_p`: Nucleus sampling is an alternative to sampling with temperature, where the model considers
162 tokens with a top_p probability mass. For example, 0.1 means only the tokens comprising
163 the top 10% probability mass are considered.
164 - `n`: The number of completions to generate for each prompt. For example, with 3 prompts and n=2,
165 the LLM will generate two completions per prompt, resulting in 6 completions total.
166 - `stop`: One or more sequences after which the LLM should stop generating tokens.
167 - `presence_penalty`: The penalty applied if a token is already present.
168 Higher values make the model less likely to repeat the token.
169 - `frequency_penalty`: Penalty applied if a token has already been generated.
170 Higher values make the model less likely to repeat the token.
171 - `logit_bias`: Adds a logit bias to specific tokens. The keys of the dictionary are tokens, and the
172 values are the bias to add to that token.
173 - `response_format`: A JSON schema or a Pydantic model that enforces the structure of the model's response.
174 If provided, the output will always be validated against this
175 format (unless the model returns a tool call).
176 For details, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs).
177 Notes:
178 - This parameter accepts Pydantic models and JSON schemas for latest models starting from GPT-4o.
179 Older models only support basic version of structured outputs through `{"type": "json_object"}`.
180 For detailed information on JSON mode, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs#json-mode).
181 - For structured outputs with streaming,
182 the `response_format` must be a JSON schema and not a Pydantic model.
183 :param default_headers: Default headers to use for the AzureOpenAI client.
184 :param tools:
185 A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
186 :param tools_strict:
187 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
188 the schema provided in the `parameters` field of the tool definition, but this may increase latency.
189 :param azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on
190 every request.
191 :param http_client_kwargs:
192 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
193 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
194 """
195 # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact
196 # with the API.
198 # Why is this here?
199 # AzureOpenAI init is forcing us to use an init method that takes either base_url or azure_endpoint as not
200 # None init parameters. This way we accommodate the use case where env var AZURE_OPENAI_ENDPOINT is set instead
201 # of passing it as a parameter.
202 azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
203 # `azure_endpoint` accepts either a plain string or a `Secret`. We keep the original value on the instance for
204 # serialization and resolve it to a string only to validate that an endpoint was provided.
205 resolved_azure_endpoint = (
206 azure_endpoint.resolve_value() if isinstance(azure_endpoint, Secret) else azure_endpoint
207 )
208 if not resolved_azure_endpoint:
209 raise ValueError("Please provide an Azure endpoint or set the environment variable AZURE_OPENAI_ENDPOINT.")
211 if api_key is None and azure_ad_token is None:
212 raise ValueError("Please provide an API key or an Azure Active Directory token.")
214 # The check above makes mypy incorrectly infer that api_key is never None,
215 # which propagates the incorrect type.
216 self.api_key = api_key # type: ignore
217 self.azure_ad_token = azure_ad_token
218 self.generation_kwargs = generation_kwargs or {}
219 self.streaming_callback = streaming_callback
220 self.api_version = api_version
221 self.azure_endpoint = azure_endpoint
222 self.azure_deployment = azure_deployment
223 self.organization = organization
224 self.model = azure_deployment or "gpt-4.1-mini"
225 self.timeout = timeout
226 self.max_retries = max_retries
227 self.default_headers = default_headers or {}
228 self.azure_ad_token_provider = azure_ad_token_provider
229 self.http_client_kwargs = http_client_kwargs
230 _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
231 self.tools = tools
232 self.tools_strict = tools_strict
234 self.client: AzureOpenAI | None = None
235 self.async_client: AsyncAzureOpenAI | None = None
237 def _client_kwargs(self) -> dict[str, Any]:
238 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
239 max_retries = (
240 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
241 )
242 resolved_azure_endpoint = (
243 self.azure_endpoint.resolve_value() if isinstance(self.azure_endpoint, Secret) else self.azure_endpoint
244 )
245 resolved_api_version = (
246 self.api_version.resolve_value() if isinstance(self.api_version, Secret) else self.api_version
247 )
248 return {
249 "api_version": resolved_api_version,
250 "azure_endpoint": resolved_azure_endpoint,
251 "azure_deployment": self.azure_deployment,
252 "api_key": self.api_key.resolve_value() if self.api_key is not None else None,
253 "azure_ad_token": self.azure_ad_token.resolve_value() if self.azure_ad_token is not None else None,
254 "organization": self.organization,
255 "timeout": timeout,
256 "max_retries": max_retries,
257 "default_headers": self.default_headers,
258 "azure_ad_token_provider": self.azure_ad_token_provider,
259 }
261 def warm_up(self) -> None:
262 """
263 Initialize the synchronous Azure OpenAI client.
264 """
265 if self.client is None:
266 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
267 # https://github.com/openai/openai-python/blob/main/httpx2.md
268 http_client = init_http_client(self.http_client_kwargs, async_client=False)
269 self.client = AzureOpenAI(
270 http_client=http_client, # type: ignore[arg-type]
271 **self._client_kwargs(),
272 )
274 async def warm_up_async(self) -> None: # noqa: RUF029
275 """
276 Initialize the asynchronous Azure OpenAI client on the serving event loop.
277 """
278 if self.async_client is None:
279 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
280 # https://github.com/openai/openai-python/blob/main/httpx2.md
281 http_client = init_http_client(self.http_client_kwargs, async_client=True)
282 self.async_client = AsyncAzureOpenAI(
283 http_client=http_client, # type: ignore[arg-type]
284 **self._client_kwargs(),
285 )
287 def close(self) -> None:
288 """
289 Releases the synchronous Azure OpenAI client.
290 """
291 if self.client is not None:
292 self.client.close()
293 self.client = None
295 async def close_async(self) -> None:
296 """
297 Releases the asynchronous Azure OpenAI client.
298 """
299 if self.async_client is not None:
300 await self.async_client.close()
301 self.async_client = None
303 def to_dict(self) -> dict[str, Any]:
304 """
305 Serialize this component to a dictionary.
307 :returns:
308 The serialized component as a dictionary.
309 """
310 callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
311 azure_ad_token_provider_name = None
312 if self.azure_ad_token_provider:
313 azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider)
314 # If the response format is a Pydantic model, it's converted to openai's json schema format
315 # If it's already a json schema, it's left as is
316 generation_kwargs = self.generation_kwargs.copy()
317 response_format = generation_kwargs.get("response_format")
318 if response_format and isinstance(response_format, type) and issubclass(response_format, BaseModel):
319 json_schema = {
320 "type": "json_schema",
321 "json_schema": {
322 "name": response_format.__name__,
323 "strict": True,
324 "schema": to_strict_json_schema(response_format),
325 },
326 }
327 generation_kwargs["response_format"] = json_schema
328 return default_to_dict(
329 self,
330 azure_endpoint=self.azure_endpoint.to_dict()
331 if isinstance(self.azure_endpoint, Secret)
332 else self.azure_endpoint,
333 azure_deployment=self.azure_deployment,
334 organization=self.organization,
335 api_version=self.api_version.to_dict() if isinstance(self.api_version, Secret) else self.api_version,
336 streaming_callback=callback_name,
337 generation_kwargs=generation_kwargs,
338 timeout=self.timeout,
339 max_retries=self.max_retries,
340 api_key=self.api_key,
341 azure_ad_token=self.azure_ad_token,
342 default_headers=self.default_headers,
343 tools=serialize_tools_or_toolset(self.tools),
344 tools_strict=self.tools_strict,
345 azure_ad_token_provider=azure_ad_token_provider_name,
346 http_client_kwargs=self.http_client_kwargs,
347 )
349 @classmethod
350 def from_dict(cls, data: dict[str, Any]) -> "AzureOpenAIChatGenerator":
351 """
352 Deserialize this component from a dictionary.
354 :param data: The dictionary representation of this component.
355 :returns:
356 The deserialized component instance.
357 """
358 deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools")
359 init_params = data.get("init_parameters", {})
360 serialized_callback_handler = init_params.get("streaming_callback")
361 if serialized_callback_handler:
362 data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
363 serialized_azure_ad_token_provider = init_params.get("azure_ad_token_provider")
364 if serialized_azure_ad_token_provider:
365 data["init_parameters"]["azure_ad_token_provider"] = deserialize_callable(
366 serialized_azure_ad_token_provider
367 )
368 return default_from_dict(cls, data)