# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import os from typing import Any, ClassVar from openai.lib._pydantic import to_strict_json_schema from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI from pydantic import BaseModel from haystack import component, default_from_dict, default_to_dict from haystack.components.generators.chat import OpenAIChatGenerator from haystack.dataclasses.streaming_chunk import StreamingCallbackT from haystack.tools import ( ToolsType, _check_duplicate_tool_names, deserialize_tools_or_toolset_inplace, flatten_tools_or_toolsets, serialize_tools_or_toolset, warm_up_tools, ) from haystack.utils import Secret, deserialize_callable, serialize_callable from haystack.utils.http_client import init_http_client @component class AzureOpenAIChatGenerator(OpenAIChatGenerator): """ Generates text using OpenAI's models on Azure. It works with the gpt-4 - type models and supports streaming responses from OpenAI API. It uses [ChatMessage](https://docs.haystack.deepset.ai/docs/chatmessage) format in input and output. You can customize how the text is generated by passing parameters to the OpenAI API. Use the `**generation_kwargs` argument when you initialize the component or when you run it. Any parameter that works with `openai.ChatCompletion.create` will work here too. For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat). ### Usage example ```python from haystack.components.generators.chat import AzureOpenAIChatGenerator from haystack.dataclasses import ChatMessage from haystack.utils import Secret messages = [ChatMessage.from_user("What's Natural Language Processing?")] client = AzureOpenAIChatGenerator( azure_endpoint="", api_key=Secret.from_token(""), azure_deployment="") response = client.run(messages) print(response) ``` ``` {'replies': [ChatMessage(_role=, _content=[TextContent(text= "Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on enabling computers to understand, interpret, and generate human language in a way that is useful.")], _name=None, _meta={'model': 'gpt-4.1-mini', 'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 15, 'completion_tokens': 36, 'total_tokens': 51}})] } ``` """ SUPPORTED_MODELS: ClassVar[list[str]] = [ "gpt-5.4", "gpt-5.4-pro", "gpt-5.3-codex", "gpt-5.2", "gpt-5.2-codex", "gpt-5.2-chat", "gpt-5.1", "gpt-5.1-chat", "gpt-5.1-codex", "gpt-5.1-codex-mini", "gpt-5", "gpt-5-mini", "gpt-5-nano", "gpt-5-chat", "gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini", "gpt-4o-audio-preview", "gpt-realtime-1.5", "gpt-audio-1.5", "o1", "o1-mini", "o3", "o3-mini", "o4-mini", "codex-mini", "gpt-4", "gpt-35-turbo", "gpt-oss-120b", "computer-use-preview", ] """A non-exhaustive list of chat models supported by this component. See https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure for the full list.""" # ruff: noqa: PLR0913 def __init__( self, azure_endpoint: str | Secret | None = None, api_version: str | Secret | None = "2024-12-01-preview", azure_deployment: str | None = "gpt-4.1-mini", api_key: Secret | None = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False), azure_ad_token: Secret | None = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False), organization: str | None = None, streaming_callback: StreamingCallbackT | None = None, timeout: float | None = None, max_retries: int | None = None, generation_kwargs: dict[str, Any] | None = None, default_headers: dict[str, str] | None = None, tools: ToolsType | None = None, tools_strict: bool = False, *, azure_ad_token_provider: AzureADTokenProvider | AsyncAzureADTokenProvider | None = None, http_client_kwargs: dict[str, Any] | None = None, ) -> None: """ Initialize the Azure OpenAI Chat Generator component. :param azure_endpoint: The endpoint of the deployed model, for example `"https://example-resource.azure.openai.com/"`. Can also be a [Secret](https://docs.haystack.deepset.ai/docs/secret-management), for example `Secret.from_env_var("AZURE_OPENAI_ENDPOINT")`, to resolve the value from an environment variable at runtime. This is useful to switch endpoints between environments (e.g. dev and prod) without changing the serialized pipeline. :param api_version: The version of the API to use. Defaults to 2024-12-01-preview. Can also be a [Secret](https://docs.haystack.deepset.ai/docs/secret-management), for example `Secret.from_env_var("AZURE_OPENAI_API_VERSION")`, to resolve the value from an environment variable at runtime. :param azure_deployment: The deployment of the model, usually the model name. :param api_key: The API key to use for authentication. :param azure_ad_token: [Azure Active Directory token](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id). :param organization: Your organization ID, defaults to `None`. For help, see [Setting up your organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization). :param streaming_callback: A callback function called when a new token is received from the stream. It accepts [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk) as an argument. :param timeout: Timeout for OpenAI client calls. If not set, it defaults to either the `OPENAI_TIMEOUT` environment variable, or 30 seconds. :param max_retries: Maximum number of retries to contact OpenAI after an internal error. If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5. :param generation_kwargs: Other parameters to use for the model. These parameters are sent directly to the OpenAI endpoint. For details, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat). Some of the supported parameters: - `max_completion_tokens`: An upper bound for the number of tokens that can be generated for a completion, including visible output tokens and reasoning tokens. - `temperature`: The sampling temperature to use. Higher values mean the model takes more risks. Try 0.9 for more creative applications and 0 (argmax sampling) for ones with a well-defined answer. - `top_p`: Nucleus sampling is an alternative to sampling with temperature, where the model considers tokens with a top_p probability mass. For example, 0.1 means only the tokens comprising the top 10% probability mass are considered. - `n`: The number of completions to generate for each prompt. For example, with 3 prompts and n=2, the LLM will generate two completions per prompt, resulting in 6 completions total. - `stop`: One or more sequences after which the LLM should stop generating tokens. - `presence_penalty`: The penalty applied if a token is already present. Higher values make the model less likely to repeat the token. - `frequency_penalty`: Penalty applied if a token has already been generated. Higher values make the model less likely to repeat the token. - `logit_bias`: Adds a logit bias to specific tokens. The keys of the dictionary are tokens, and the values are the bias to add to that token. - `response_format`: A JSON schema or a Pydantic model that enforces the structure of the model's response. If provided, the output will always be validated against this format (unless the model returns a tool call). For details, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs). Notes: - This parameter accepts Pydantic models and JSON schemas for latest models starting from GPT-4o. Older models only support basic version of structured outputs through `{"type": "json_object"}`. For detailed information on JSON mode, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs#json-mode). - For structured outputs with streaming, the `response_format` must be a JSON schema and not a Pydantic model. :param default_headers: Default headers to use for the AzureOpenAI client. :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls. :param tools_strict: Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly the schema provided in the `parameters` field of the tool definition, but this may increase latency. :param azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on every request. :param http_client_kwargs: A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`. For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client). """ # We intentionally do not call super().__init__ here because we only need to instantiate the client to interact # with the API. # Why is this here? # AzureOpenAI init is forcing us to use an init method that takes either base_url or azure_endpoint as not # None init parameters. This way we accommodate the use case where env var AZURE_OPENAI_ENDPOINT is set instead # of passing it as a parameter. azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT") # `azure_endpoint` accepts either a plain string or a `Secret`. We keep the original value on the instance for # serialization and resolve it to a string only to validate that an endpoint was provided. resolved_azure_endpoint = ( azure_endpoint.resolve_value() if isinstance(azure_endpoint, Secret) else azure_endpoint ) if not resolved_azure_endpoint: raise ValueError("Please provide an Azure endpoint or set the environment variable AZURE_OPENAI_ENDPOINT.") if api_key is None and azure_ad_token is None: raise ValueError("Please provide an API key or an Azure Active Directory token.") # The check above makes mypy incorrectly infer that api_key is never None, # which propagates the incorrect type. self.api_key = api_key # type: ignore self.azure_ad_token = azure_ad_token self.generation_kwargs = generation_kwargs or {} self.streaming_callback = streaming_callback self.api_version = api_version self.azure_endpoint = azure_endpoint self.azure_deployment = azure_deployment self.organization = organization self.model = azure_deployment or "gpt-4.1-mini" self.timeout = timeout self.max_retries = max_retries self.default_headers = default_headers or {} self.azure_ad_token_provider = azure_ad_token_provider self.http_client_kwargs = http_client_kwargs _check_duplicate_tool_names(flatten_tools_or_toolsets(tools)) self.tools = tools self.tools_strict = tools_strict self.client: AzureOpenAI | None = None self.async_client: AsyncAzureOpenAI | None = None self._tools_warmed_up = False def _client_kwargs(self) -> dict[str, Any]: timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0")) max_retries = ( self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5")) ) resolved_azure_endpoint = ( self.azure_endpoint.resolve_value() if isinstance(self.azure_endpoint, Secret) else self.azure_endpoint ) resolved_api_version = ( self.api_version.resolve_value() if isinstance(self.api_version, Secret) else self.api_version ) return { "api_version": resolved_api_version, "azure_endpoint": resolved_azure_endpoint, "azure_deployment": self.azure_deployment, "api_key": self.api_key.resolve_value() if self.api_key is not None else None, "azure_ad_token": self.azure_ad_token.resolve_value() if self.azure_ad_token is not None else None, "organization": self.organization, "timeout": timeout, "max_retries": max_retries, "default_headers": self.default_headers, "azure_ad_token_provider": self.azure_ad_token_provider, } def _warm_up_tools(self) -> None: if not self._tools_warmed_up: warm_up_tools(self.tools) self._tools_warmed_up = True def warm_up(self) -> None: """ Warm up the tools and initialize the synchronous Azure OpenAI client. """ self._warm_up_tools() if self.client is None: self.client = AzureOpenAI( http_client=init_http_client(self.http_client_kwargs, async_client=False), **self._client_kwargs() ) async def warm_up_async(self) -> None: # noqa: RUF029 """ Warm up the tools and initialize the asynchronous Azure OpenAI client on the serving event loop. """ self._warm_up_tools() if self.async_client is None: self.async_client = AsyncAzureOpenAI( http_client=init_http_client(self.http_client_kwargs, async_client=True), **self._client_kwargs() ) def close(self) -> None: """ Releases the synchronous Azure OpenAI client. """ if self.client is not None: self.client.close() self.client = None async def close_async(self) -> None: """ Releases the asynchronous Azure OpenAI client. """ if self.async_client is not None: await self.async_client.close() self.async_client = None def to_dict(self) -> dict[str, Any]: """ Serialize this component to a dictionary. :returns: The serialized component as a dictionary. """ callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None azure_ad_token_provider_name = None if self.azure_ad_token_provider: azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider) # If the response format is a Pydantic model, it's converted to openai's json schema format # If it's already a json schema, it's left as is generation_kwargs = self.generation_kwargs.copy() response_format = generation_kwargs.get("response_format") if response_format and issubclass(response_format, BaseModel): json_schema = { "type": "json_schema", "json_schema": { "name": response_format.__name__, "strict": True, "schema": to_strict_json_schema(response_format), }, } generation_kwargs["response_format"] = json_schema return default_to_dict( self, azure_endpoint=self.azure_endpoint.to_dict() if isinstance(self.azure_endpoint, Secret) else self.azure_endpoint, azure_deployment=self.azure_deployment, organization=self.organization, api_version=self.api_version.to_dict() if isinstance(self.api_version, Secret) else self.api_version, streaming_callback=callback_name, generation_kwargs=generation_kwargs, timeout=self.timeout, max_retries=self.max_retries, api_key=self.api_key, azure_ad_token=self.azure_ad_token, default_headers=self.default_headers, tools=serialize_tools_or_toolset(self.tools), tools_strict=self.tools_strict, azure_ad_token_provider=azure_ad_token_provider_name, http_client_kwargs=self.http_client_kwargs, ) @classmethod def from_dict(cls, data: dict[str, Any]) -> "AzureOpenAIChatGenerator": """ Deserialize this component from a dictionary. :param data: The dictionary representation of this component. :returns: The deserialized component instance. """ deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools") init_params = data.get("init_parameters", {}) serialized_callback_handler = init_params.get("streaming_callback") if serialized_callback_handler: data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler) serialized_azure_ad_token_provider = init_params.get("azure_ad_token_provider") if serialized_azure_ad_token_provider: data["init_parameters"]["azure_ad_token_provider"] = deserialize_callable( serialized_azure_ad_token_provider ) return default_from_dict(cls, data)