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