• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

deepset-ai / haystack / 15049844454

15 May 2025 04:07PM UTC coverage: 90.446% (+0.04%) from 90.41%
15049844454

Pull #9345

github

web-flow
Merge 9e4071f83 into 2a64cd4e9
Pull Request #9345: feat: add serialization to `State` / move `State` to utils

10981 of 12141 relevant lines covered (90.45%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.73
haystack/components/generators/chat/azure.py
1
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2
#
3
# SPDX-License-Identifier: Apache-2.0
4

5
import os
1✔
6
from typing import Any, Dict, List, Optional, Union
1✔
7

8
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI
1✔
9

10
from haystack import component, default_from_dict, default_to_dict
1✔
11
from haystack.components.generators.chat import OpenAIChatGenerator
1✔
12
from haystack.dataclasses.streaming_chunk import StreamingCallbackT
1✔
13
from haystack.tools import (
1✔
14
    Tool,
15
    Toolset,
16
    _check_duplicate_tool_names,
17
    deserialize_tools_or_toolset_inplace,
18
    serialize_tools_or_toolset,
19
)
20
from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
1✔
21
from haystack.utils.http_client import init_http_client
1✔
22

23

24
@component
1✔
25
class AzureOpenAIChatGenerator(OpenAIChatGenerator):
1✔
26
    """
27
    Generates text using OpenAI's models on Azure.
28

29
    It works with the gpt-4 - type models and supports streaming responses
30
    from OpenAI API. It uses [ChatMessage](https://docs.haystack.deepset.ai/docs/chatmessage)
31
    format in input and output.
32

33
    You can customize how the text is generated by passing parameters to the
34
    OpenAI API. Use the `**generation_kwargs` argument when you initialize
35
    the component or when you run it. Any parameter that works with
36
    `openai.ChatCompletion.create` will work here too.
37

38
    For details on OpenAI API parameters, see
39
    [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat).
40

41
    ### Usage example
42

43
    ```python
44
    from haystack.components.generators.chat import AzureOpenAIChatGenerator
45
    from haystack.dataclasses import ChatMessage
46
    from haystack.utils import Secret
47

48
    messages = [ChatMessage.from_user("What's Natural Language Processing?")]
49

50
    client = AzureOpenAIChatGenerator(
51
        azure_endpoint="<Your Azure endpoint e.g. `https://your-company.azure.openai.com/>",
52
        api_key=Secret.from_token("<your-api-key>"),
53
        azure_deployment="<this a model name, e.g. gpt-4o-mini>")
54
    response = client.run(messages)
55
    print(response)
56
    ```
57

58
    ```
59
    {'replies':
60
        [ChatMessage(_role=<ChatRole.ASSISTANT: 'assistant'>, _content=[TextContent(text=
61
        "Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on
62
         enabling computers to understand, interpret, and generate human language in a way that is useful.")],
63
         _name=None,
64
         _meta={'model': 'gpt-4o-mini', 'index': 0, 'finish_reason': 'stop',
65
         'usage': {'prompt_tokens': 15, 'completion_tokens': 36, 'total_tokens': 51}})]
66
    }
67
    ```
68
    """
69

70
    # pylint: disable=super-init-not-called
71
    # ruff: noqa: PLR0913
72
    def __init__(  # pylint: disable=too-many-positional-arguments
1✔
73
        self,
74
        azure_endpoint: Optional[str] = None,
75
        api_version: Optional[str] = "2023-05-15",
76
        azure_deployment: Optional[str] = "gpt-4o-mini",
77
        api_key: Optional[Secret] = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False),
78
        azure_ad_token: Optional[Secret] = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False),
79
        organization: Optional[str] = None,
80
        streaming_callback: Optional[StreamingCallbackT] = None,
81
        timeout: Optional[float] = None,
82
        max_retries: Optional[int] = None,
83
        generation_kwargs: Optional[Dict[str, Any]] = None,
84
        default_headers: Optional[Dict[str, str]] = None,
85
        tools: Optional[Union[List[Tool], Toolset]] = None,
86
        tools_strict: bool = False,
87
        *,
88
        azure_ad_token_provider: Optional[Union[AzureADTokenProvider, AsyncAzureADTokenProvider]] = None,
89
        http_client_kwargs: Optional[Dict[str, Any]] = None,
90
    ):
91
        """
92
        Initialize the Azure OpenAI Chat Generator component.
93

94
        :param azure_endpoint: The endpoint of the deployed model, for example `"https://example-resource.azure.openai.com/"`.
95
        :param api_version: The version of the API to use. Defaults to 2023-05-15.
96
        :param azure_deployment: The deployment of the model, usually the model name.
97
        :param api_key: The API key to use for authentication.
98
        :param azure_ad_token: [Azure Active Directory token](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id).
99
        :param organization: Your organization ID, defaults to `None`. For help, see
100
        [Setting up your organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization).
101
        :param streaming_callback: A callback function called when a new token is received from the stream.
102
            It accepts [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk)
103
            as an argument.
104
        :param timeout: Timeout for OpenAI client calls. If not set, it defaults to either the
105
            `OPENAI_TIMEOUT` environment variable, or 30 seconds.
106
        :param max_retries: Maximum number of retries to contact OpenAI after an internal error.
107
            If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5.
108
        :param generation_kwargs: Other parameters to use for the model. These parameters are sent directly to
109
            the OpenAI endpoint. For details, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat).
110
            Some of the supported parameters:
111
            - `max_tokens`: The maximum number of tokens the output text can have.
112
            - `temperature`: The sampling temperature to use. Higher values mean the model takes more risks.
113
                Try 0.9 for more creative applications and 0 (argmax sampling) for ones with a well-defined answer.
114
            - `top_p`: Nucleus sampling is an alternative to sampling with temperature, where the model considers
115
                tokens with a top_p probability mass. For example, 0.1 means only the tokens comprising
116
                the top 10% probability mass are considered.
117
            - `n`: The number of completions to generate for each prompt. For example, with 3 prompts and n=2,
118
                the LLM will generate two completions per prompt, resulting in 6 completions total.
119
            - `stop`: One or more sequences after which the LLM should stop generating tokens.
120
            - `presence_penalty`: The penalty applied if a token is already present.
121
                Higher values make the model less likely to repeat the token.
122
            - `frequency_penalty`: Penalty applied if a token has already been generated.
123
                Higher values make the model less likely to repeat the token.
124
            - `logit_bias`: Adds a logit bias to specific tokens. The keys of the dictionary are tokens, and the
125
                values are the bias to add to that token.
126
        :param default_headers: Default headers to use for the AzureOpenAI client.
127
        :param tools:
128
            A list of tools or a Toolset for which the model can prepare calls. This parameter can accept either a
129
            list of `Tool` objects or a `Toolset` instance.
130
        :param tools_strict:
131
            Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
132
            the schema provided in the `parameters` field of the tool definition, but this may increase latency.
133
        :param azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on
134
            every request.
135
        :param http_client_kwargs:
136
            A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
137
            For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
138
        """
139
        # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact
140
        # with the API.
141

142
        # Why is this here?
143
        # AzureOpenAI init is forcing us to use an init method that takes either base_url or azure_endpoint as not
144
        # None init parameters. This way we accommodate the use case where env var AZURE_OPENAI_ENDPOINT is set instead
145
        # of passing it as a parameter.
146
        azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
1✔
147
        if not azure_endpoint:
1✔
148
            raise ValueError("Please provide an Azure endpoint or set the environment variable AZURE_OPENAI_ENDPOINT.")
×
149

150
        if api_key is None and azure_ad_token is None:
1✔
151
            raise ValueError("Please provide an API key or an Azure Active Directory token.")
×
152

153
        # The check above makes mypy incorrectly infer that api_key is never None,
154
        # which propagates the incorrect type.
155
        self.api_key = api_key  # type: ignore
1✔
156
        self.azure_ad_token = azure_ad_token
1✔
157
        self.generation_kwargs = generation_kwargs or {}
1✔
158
        self.streaming_callback = streaming_callback
1✔
159
        self.api_version = api_version
1✔
160
        self.azure_endpoint = azure_endpoint
1✔
161
        self.azure_deployment = azure_deployment
1✔
162
        self.organization = organization
1✔
163
        self.model = azure_deployment or "gpt-4o-mini"
1✔
164
        self.timeout = timeout if timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
1✔
165
        self.max_retries = max_retries if max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
1✔
166
        self.default_headers = default_headers or {}
1✔
167
        self.azure_ad_token_provider = azure_ad_token_provider
1✔
168
        self.http_client_kwargs = http_client_kwargs
1✔
169
        _check_duplicate_tool_names(list(tools or []))
1✔
170
        self.tools = tools
1✔
171
        self.tools_strict = tools_strict
1✔
172

173
        client_args: Dict[str, Any] = {
1✔
174
            "api_version": api_version,
175
            "azure_endpoint": azure_endpoint,
176
            "azure_deployment": azure_deployment,
177
            "api_key": api_key.resolve_value() if api_key is not None else None,
178
            "azure_ad_token": azure_ad_token.resolve_value() if azure_ad_token is not None else None,
179
            "organization": organization,
180
            "timeout": self.timeout,
181
            "max_retries": self.max_retries,
182
            "default_headers": self.default_headers,
183
            "azure_ad_token_provider": azure_ad_token_provider,
184
        }
185

186
        self.client = AzureOpenAI(
1✔
187
            http_client=init_http_client(self.http_client_kwargs, async_client=False), **client_args
188
        )
189
        self.async_client = AsyncAzureOpenAI(
1✔
190
            http_client=init_http_client(self.http_client_kwargs, async_client=True), **client_args
191
        )
192

193
    def to_dict(self) -> Dict[str, Any]:
1✔
194
        """
195
        Serialize this component to a dictionary.
196

197
        :returns:
198
            The serialized component as a dictionary.
199
        """
200
        callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
1✔
201
        azure_ad_token_provider_name = None
1✔
202
        if self.azure_ad_token_provider:
1✔
203
            azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider)
1✔
204
        return default_to_dict(
1✔
205
            self,
206
            azure_endpoint=self.azure_endpoint,
207
            azure_deployment=self.azure_deployment,
208
            organization=self.organization,
209
            api_version=self.api_version,
210
            streaming_callback=callback_name,
211
            generation_kwargs=self.generation_kwargs,
212
            timeout=self.timeout,
213
            max_retries=self.max_retries,
214
            api_key=self.api_key.to_dict() if self.api_key is not None else None,
215
            azure_ad_token=self.azure_ad_token.to_dict() if self.azure_ad_token is not None else None,
216
            default_headers=self.default_headers,
217
            tools=serialize_tools_or_toolset(self.tools),
218
            tools_strict=self.tools_strict,
219
            azure_ad_token_provider=azure_ad_token_provider_name,
220
            http_client_kwargs=self.http_client_kwargs,
221
        )
222

223
    @classmethod
1✔
224
    def from_dict(cls, data: Dict[str, Any]) -> "AzureOpenAIChatGenerator":
1✔
225
        """
226
        Deserialize this component from a dictionary.
227

228
        :param data: The dictionary representation of this component.
229
        :returns:
230
            The deserialized component instance.
231
        """
232
        deserialize_secrets_inplace(data["init_parameters"], keys=["api_key", "azure_ad_token"])
1✔
233
        deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools")
1✔
234
        init_params = data.get("init_parameters", {})
1✔
235
        serialized_callback_handler = init_params.get("streaming_callback")
1✔
236
        if serialized_callback_handler:
1✔
237
            data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
×
238
        serialized_azure_ad_token_provider = init_params.get("azure_ad_token_provider")
1✔
239
        if serialized_azure_ad_token_provider:
1✔
240
            data["init_parameters"]["azure_ad_token_provider"] = deserialize_callable(
×
241
                serialized_azure_ad_token_provider
242
            )
243
        return default_from_dict(cls, data)
1✔
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc