deepset-ai--haystack/haystack/components/embedders/azure_text_embedder.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import os
from typing import Any

from openai.lib.azure import AsyncAzureOpenAI, AzureADTokenProvider, AzureOpenAI

from haystack import component, default_from_dict, default_to_dict
from haystack.components.embedders import OpenAITextEmbedder
from haystack.utils import Secret, deserialize_callable, serialize_callable
from haystack.utils.http_client import init_http_client


@component
class AzureOpenAITextEmbedder(OpenAITextEmbedder):
    """
    Embeds strings using OpenAI models deployed on Azure.

    ### Usage example
    <!-- test-ignore -->
    ```python
    from haystack.components.embedders import AzureOpenAITextEmbedder

    text_to_embed = "I love pizza!"
    text_embedder = AzureOpenAITextEmbedder()

    print(text_embedder.run(text_to_embed))

    # {'embedding': [0.017020374536514282, -0.023255806416273117, ...],
    # 'meta': {'model': 'text-embedding-ada-002-v2',
    #          'usage': {'prompt_tokens': 4, 'total_tokens': 4}}}
    ```
    """

    def __init__(  # noqa: PLR0913
        self,
        azure_endpoint: str | None = None,
        api_version: str | None = "2023-05-15",
        azure_deployment: str = "text-embedding-ada-002",
        dimensions: int | None = None,
        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,
        timeout: float | None = None,
        max_retries: int | None = None,
        prefix: str = "",
        suffix: str = "",
        *,
        default_headers: dict[str, str] | None = None,
        azure_ad_token_provider: AzureADTokenProvider | None = None,
        http_client_kwargs: dict[str, Any] | None = None,
    ) -> None:
        """
        Creates an AzureOpenAITextEmbedder component.

        :param azure_endpoint:
            The endpoint of the model deployed on Azure.
        :param api_version:
            The version of the API to use.
        :param azure_deployment:
            The name of the model deployed on Azure. The default model is text-embedding-ada-002.
        :param dimensions:
            The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3
            and later models.
        :param api_key:
            The Azure OpenAI API key.
            You can set it with an environment variable `AZURE_OPENAI_API_KEY`, or pass with this
            parameter during initialization.
        :param azure_ad_token:
            Microsoft Entra ID token, see Microsoft's
            [Entra ID](https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id)
            documentation for more information. You can set it with an environment variable
            `AZURE_OPENAI_AD_TOKEN`, or pass with this parameter during initialization.
            Previously called Azure Active Directory.
        :param organization:
            Your organization ID. See OpenAI's
            [Setting Up Your Organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
            for more information.
        :param timeout: The timeout for `AzureOpenAI` client calls, in seconds.
            If not set, defaults to either the
            `OPENAI_TIMEOUT` environment variable, or 30 seconds.
        :param max_retries: Maximum number of retries to contact AzureOpenAI after an internal error.
            If not set, defaults to either the `OPENAI_MAX_RETRIES` environment variable, or to 5 retries.
        :param prefix:
            A string to add at the beginning of each text.
        :param suffix:
            A string to add at the end of each text.
        :param default_headers: Default headers to send to the AzureOpenAI client.
        :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")
        if not 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.")

        self.api_key = api_key  # type: ignore[assignment] # mypy does not understand that api_key can be None
        self.azure_ad_token = azure_ad_token
        self.api_version = api_version
        self.azure_endpoint = azure_endpoint
        self.azure_deployment = azure_deployment
        self.model = azure_deployment
        self.dimensions = dimensions
        self.organization = organization
        self.timeout = timeout
        self.max_retries = max_retries
        self.prefix = prefix
        self.suffix = suffix
        self.default_headers = default_headers or {}
        self.azure_ad_token_provider = azure_ad_token_provider
        self.http_client_kwargs = http_client_kwargs

        self.client: AzureOpenAI | None = None
        self.async_client: AsyncAzureOpenAI | None = None

    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"))
        )
        return {
            "api_version": self.api_version,
            "azure_endpoint": self.azure_endpoint,
            "azure_deployment": self.azure_deployment,
            "azure_ad_token_provider": self.azure_ad_token_provider,
            "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,
        }

    def warm_up(self) -> None:
        """
        Initializes the synchronous Azure OpenAI client.
        """
        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
        """
        Initializes the asynchronous Azure OpenAI client on the serving event loop.
        """
        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]:
        """
        Serializes the component to a dictionary.

        :returns:
            Dictionary with serialized data.
        """
        azure_ad_token_provider_name = None
        if self.azure_ad_token_provider:
            azure_ad_token_provider_name = serialize_callable(self.azure_ad_token_provider)
        return default_to_dict(
            self,
            azure_endpoint=self.azure_endpoint,
            azure_deployment=self.azure_deployment,
            dimensions=self.dimensions,
            organization=self.organization,
            api_version=self.api_version,
            prefix=self.prefix,
            suffix=self.suffix,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            timeout=self.timeout,
            max_retries=self.max_retries,
            default_headers=self.default_headers,
            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]) -> "AzureOpenAITextEmbedder":
        """
        Deserializes the component from a dictionary.

        :param data:
            Dictionary to deserialize from.
        :returns:
            Deserialized component.
        """
        serialized_azure_ad_token_provider = data["init_parameters"].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)