duckdb-nsql/manifest/manifest/clients/azureopenai_chat.py

"""Azure client."""
import logging
import os
from typing import Any, Dict, Optional

from manifest.clients.openai_chat import OPENAICHAT_ENGINES, OpenAIChatClient
from manifest.request import LMRequest

logger = logging.getLogger(__name__)

# Azure deployment name can only use letters and numbers, no spaces. Hyphens ("-") and
# underscores ("_") may be used, except as ending characters. We create this mapping to
# handle difference between Azure and OpenAI
AZURE_DEPLOYMENT_NAME_MAPPING = {
    "gpt-3.5-turbo": "gpt-35-turbo",
    "gpt-3.5-turbo-0301": "gpt-35-turbo-0301",
}
OPENAI_DEPLOYMENT_NAME_MAPPING = {
    "gpt-35-turbo": "gpt-3.5-turbo",
    "gpt-35-turbo-0301": "gpt-3.5-turbo-0301",
}


class AzureChatClient(OpenAIChatClient):
    """Azure chat client."""

    # User param -> (client param, default value)
    PARAMS = OpenAIChatClient.PARAMS
    REQUEST_CLS = LMRequest
    NAME = "azureopenaichat"
    IS_CHAT = True

    def connect(
        self,
        connection_str: Optional[str] = None,
        client_args: Dict[str, Any] = {},
    ) -> None:
        """
        Connect to the AzureOpenAI server.

        connection_str is passed as default AZURE_OPENAI_KEY if variable not set.

        Args:
            connection_str: connection string.
            client_args: client arguments.
        """
        self.api_key, self.host = None, None
        if connection_str:
            connection_parts = connection_str.split("::")
            if len(connection_parts) == 1:
                self.api_key = connection_parts[0]
            elif len(connection_parts) == 2:
                self.api_key, self.host = connection_parts
            else:
                raise ValueError(
                    "Invalid connection string. "
                    "Must be either AZURE_OPENAI_KEY or "
                    "AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
                )
        self.api_key = self.api_key or os.environ.get("AZURE_OPENAI_KEY")
        if self.api_key is None:
            raise ValueError(
                "AzureOpenAI API key not set. Set AZURE_OPENAI_KEY environment "
                "variable or pass through `client_connection`."
            )
        self.host = self.host or os.environ.get("AZURE_OPENAI_ENDPOINT")
        if self.host is None:
            raise ValueError(
                "Azure Service URL not set "
                "(e.g. https://openai-azure-service.openai.azure.com/)."
                " Set AZURE_OPENAI_ENDPOINT or pass through `client_connection`."
                " as AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
            )
        self.host = self.host.rstrip("/")
        for key in self.PARAMS:
            setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
        if getattr(self, "engine") not in OPENAICHAT_ENGINES:
            raise ValueError(
                f"Invalid engine {getattr(self, 'engine')}. "
                f"Must be {OPENAICHAT_ENGINES}."
            )

    def get_generation_url(self) -> str:
        """Get generation URL."""
        engine = getattr(self, "engine")
        deployment_name = AZURE_DEPLOYMENT_NAME_MAPPING.get(engine, engine)
        return (
            self.host
            + "/openai/deployments/"
            + deployment_name
            + "/chat/completions?api-version=2023-05-15"
        )

    def get_generation_header(self) -> Dict[str, str]:
        """
        Get generation header.

        Returns:
            header.
        """
        return {"api-key": f"{self.api_key}"}

    def get_model_params(self) -> Dict:
        """
        Get model params.

        By getting model params from the server, we can add to request
        and make sure cache keys are unique to model.

        Returns:
            model params.
        """
        # IMPORTANT!!!
        # Azure models are the same as openai models. So we want to unify their
        # cached. Make sure we retrun the OpenAI name here.
        return {"model_name": OpenAIChatClient.NAME, "engine": getattr(self, "engine")}