440 lines
17 KiB
Python
440 lines
17 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
import inspect
|
||
|
|
from typing import Any, Union, Mapping, TypeVar, Callable, Awaitable, overload
|
||
|
|
from typing_extensions import override
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from .._types import NOT_GIVEN, Omit, Timeout, NotGiven
|
||
|
|
from .._utils import is_given, is_mapping
|
||
|
|
from .._client import OpenAI, AsyncOpenAI
|
||
|
|
from .._models import FinalRequestOptions
|
||
|
|
from .._streaming import Stream, AsyncStream
|
||
|
|
from .._exceptions import OpenAIError
|
||
|
|
from .._base_client import DEFAULT_MAX_RETRIES, BaseClient
|
||
|
|
|
||
|
|
_deployments_endpoints = set(
|
||
|
|
[
|
||
|
|
"/completions",
|
||
|
|
"/chat/completions",
|
||
|
|
"/embeddings",
|
||
|
|
"/audio/transcriptions",
|
||
|
|
"/audio/translations",
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
AzureADTokenProvider = Callable[[], str]
|
||
|
|
AsyncAzureADTokenProvider = Callable[[], "str | Awaitable[str]"]
|
||
|
|
_HttpxClientT = TypeVar("_HttpxClientT", bound=Union[httpx.Client, httpx.AsyncClient])
|
||
|
|
_DefaultStreamT = TypeVar("_DefaultStreamT", bound=Union[Stream[Any], AsyncStream[Any]])
|
||
|
|
|
||
|
|
|
||
|
|
# we need to use a sentinel API key value for Azure AD
|
||
|
|
# as we don't want to make the `api_key` in the main client Optional
|
||
|
|
# and Azure AD tokens may be retrieved on a per-request basis
|
||
|
|
API_KEY_SENTINEL = "".join(["<", "missing API key", ">"])
|
||
|
|
|
||
|
|
|
||
|
|
class MutuallyExclusiveAuthError(OpenAIError):
|
||
|
|
def __init__(self) -> None:
|
||
|
|
super().__init__(
|
||
|
|
"The `api_key`, `azure_ad_token` and `azure_ad_token_provider` arguments are mutually exclusive; Only one can be passed at a time"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class BaseAzureClient(BaseClient[_HttpxClientT, _DefaultStreamT]):
|
||
|
|
@override
|
||
|
|
def _build_request(
|
||
|
|
self,
|
||
|
|
options: FinalRequestOptions,
|
||
|
|
) -> httpx.Request:
|
||
|
|
if options.url in _deployments_endpoints and is_mapping(options.json_data):
|
||
|
|
model = options.json_data.get("model")
|
||
|
|
if model is not None and not "/deployments" in str(self.base_url):
|
||
|
|
options.url = f"/deployments/{model}{options.url}"
|
||
|
|
|
||
|
|
return super()._build_request(options)
|
||
|
|
|
||
|
|
|
||
|
|
class AzureOpenAI(BaseAzureClient[httpx.Client, Stream[Any]], OpenAI):
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
azure_endpoint: str,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.Client | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.Client | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
base_url: str,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.Client | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
api_version: str | None = None,
|
||
|
|
azure_endpoint: str | None = None,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
base_url: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.Client | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
"""Construct a new synchronous azure openai client instance.
|
||
|
|
|
||
|
|
This automatically infers the following arguments from their corresponding environment variables if they are not provided:
|
||
|
|
- `api_key` from `AZURE_OPENAI_API_KEY`
|
||
|
|
- `organization` from `OPENAI_ORG_ID`
|
||
|
|
- `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
|
||
|
|
- `api_version` from `OPENAI_API_VERSION`
|
||
|
|
- `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`
|
||
|
|
|
||
|
|
Args:
|
||
|
|
azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`
|
||
|
|
|
||
|
|
azure_ad_token: Your Azure Active Directory token, https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id
|
||
|
|
|
||
|
|
azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on every request.
|
||
|
|
|
||
|
|
azure_deployment: A model deployment, if given sets the base client URL to include `/deployments/{azure_deployment}`.
|
||
|
|
Note: this means you won't be able to use non-deployment endpoints.
|
||
|
|
"""
|
||
|
|
if api_key is None:
|
||
|
|
api_key = os.environ.get("AZURE_OPENAI_API_KEY")
|
||
|
|
|
||
|
|
if azure_ad_token is None:
|
||
|
|
azure_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
|
||
|
|
|
||
|
|
if api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
|
||
|
|
raise OpenAIError(
|
||
|
|
"Missing credentials. Please pass one of `api_key`, `azure_ad_token`, `azure_ad_token_provider`, or the `AZURE_OPENAI_API_KEY` or `AZURE_OPENAI_AD_TOKEN` environment variables."
|
||
|
|
)
|
||
|
|
|
||
|
|
if api_version is None:
|
||
|
|
api_version = os.environ.get("OPENAI_API_VERSION")
|
||
|
|
|
||
|
|
if api_version is None:
|
||
|
|
raise ValueError(
|
||
|
|
"Must provide either the `api_version` argument or the `OPENAI_API_VERSION` environment variable"
|
||
|
|
)
|
||
|
|
|
||
|
|
if default_query is None:
|
||
|
|
default_query = {"api-version": api_version}
|
||
|
|
else:
|
||
|
|
default_query = {"api-version": api_version, **default_query}
|
||
|
|
|
||
|
|
if base_url is None:
|
||
|
|
if azure_endpoint is None:
|
||
|
|
azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||
|
|
|
||
|
|
if azure_endpoint is None:
|
||
|
|
raise ValueError(
|
||
|
|
"Must provide one of the `base_url` or `azure_endpoint` arguments, or the `AZURE_OPENAI_ENDPOINT` environment variable"
|
||
|
|
)
|
||
|
|
|
||
|
|
if azure_deployment is not None:
|
||
|
|
base_url = f"{azure_endpoint}/openai/deployments/{azure_deployment}"
|
||
|
|
else:
|
||
|
|
base_url = f"{azure_endpoint}/openai"
|
||
|
|
else:
|
||
|
|
if azure_endpoint is not None:
|
||
|
|
raise ValueError("base_url and azure_endpoint are mutually exclusive")
|
||
|
|
|
||
|
|
if api_key is None:
|
||
|
|
# define a sentinel value to avoid any typing issues
|
||
|
|
api_key = API_KEY_SENTINEL
|
||
|
|
|
||
|
|
super().__init__(
|
||
|
|
api_key=api_key,
|
||
|
|
organization=organization,
|
||
|
|
base_url=base_url,
|
||
|
|
timeout=timeout,
|
||
|
|
max_retries=max_retries,
|
||
|
|
default_headers=default_headers,
|
||
|
|
default_query=default_query,
|
||
|
|
http_client=http_client,
|
||
|
|
_strict_response_validation=_strict_response_validation,
|
||
|
|
)
|
||
|
|
self._azure_ad_token = azure_ad_token
|
||
|
|
self._azure_ad_token_provider = azure_ad_token_provider
|
||
|
|
|
||
|
|
def _get_azure_ad_token(self) -> str | None:
|
||
|
|
if self._azure_ad_token is not None:
|
||
|
|
return self._azure_ad_token
|
||
|
|
|
||
|
|
provider = self._azure_ad_token_provider
|
||
|
|
if provider is not None:
|
||
|
|
token = provider()
|
||
|
|
if not token or not isinstance(token, str): # pyright: ignore[reportUnnecessaryIsInstance]
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected `azure_ad_token_provider` argument to return a string but it returned {token}",
|
||
|
|
)
|
||
|
|
return token
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
@override
|
||
|
|
def _prepare_options(self, options: FinalRequestOptions) -> None:
|
||
|
|
headers: dict[str, str | Omit] = {**options.headers} if is_given(options.headers) else {}
|
||
|
|
options.headers = headers
|
||
|
|
|
||
|
|
azure_ad_token = self._get_azure_ad_token()
|
||
|
|
if azure_ad_token is not None:
|
||
|
|
if headers.get("Authorization") is None:
|
||
|
|
headers["Authorization"] = f"Bearer {azure_ad_token}"
|
||
|
|
elif self.api_key is not API_KEY_SENTINEL:
|
||
|
|
if headers.get("api-key") is None:
|
||
|
|
headers["api-key"] = self.api_key
|
||
|
|
else:
|
||
|
|
# should never be hit
|
||
|
|
raise ValueError("Unable to handle auth")
|
||
|
|
|
||
|
|
return super()._prepare_options(options)
|
||
|
|
|
||
|
|
|
||
|
|
class AsyncAzureOpenAI(BaseAzureClient[httpx.AsyncClient, AsyncStream[Any]], AsyncOpenAI):
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
azure_endpoint: str,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.AsyncClient | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.AsyncClient | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
base_url: str,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.AsyncClient | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
...
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
azure_endpoint: str | None = None,
|
||
|
|
azure_deployment: str | None = None,
|
||
|
|
api_version: str | None = None,
|
||
|
|
api_key: str | None = None,
|
||
|
|
azure_ad_token: str | None = None,
|
||
|
|
azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||
|
|
organization: str | None = None,
|
||
|
|
base_url: str | None = None,
|
||
|
|
timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
|
||
|
|
max_retries: int = DEFAULT_MAX_RETRIES,
|
||
|
|
default_headers: Mapping[str, str] | None = None,
|
||
|
|
default_query: Mapping[str, object] | None = None,
|
||
|
|
http_client: httpx.AsyncClient | None = None,
|
||
|
|
_strict_response_validation: bool = False,
|
||
|
|
) -> None:
|
||
|
|
"""Construct a new asynchronous azure openai client instance.
|
||
|
|
|
||
|
|
This automatically infers the following arguments from their corresponding environment variables if they are not provided:
|
||
|
|
- `api_key` from `AZURE_OPENAI_API_KEY`
|
||
|
|
- `organization` from `OPENAI_ORG_ID`
|
||
|
|
- `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
|
||
|
|
- `api_version` from `OPENAI_API_VERSION`
|
||
|
|
- `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`
|
||
|
|
|
||
|
|
Args:
|
||
|
|
azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`
|
||
|
|
|
||
|
|
azure_ad_token: Your Azure Active Directory token, https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id
|
||
|
|
|
||
|
|
azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on every request.
|
||
|
|
|
||
|
|
azure_deployment: A model deployment, if given sets the base client URL to include `/deployments/{azure_deployment}`.
|
||
|
|
Note: this means you won't be able to use non-deployment endpoints.
|
||
|
|
"""
|
||
|
|
if api_key is None:
|
||
|
|
api_key = os.environ.get("AZURE_OPENAI_API_KEY")
|
||
|
|
|
||
|
|
if azure_ad_token is None:
|
||
|
|
azure_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
|
||
|
|
|
||
|
|
if api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
|
||
|
|
raise OpenAIError(
|
||
|
|
"Missing credentials. Please pass one of `api_key`, `azure_ad_token`, `azure_ad_token_provider`, or the `AZURE_OPENAI_API_KEY` or `AZURE_OPENAI_AD_TOKEN` environment variables."
|
||
|
|
)
|
||
|
|
|
||
|
|
if api_version is None:
|
||
|
|
api_version = os.environ.get("OPENAI_API_VERSION")
|
||
|
|
|
||
|
|
if api_version is None:
|
||
|
|
raise ValueError(
|
||
|
|
"Must provide either the `api_version` argument or the `OPENAI_API_VERSION` environment variable"
|
||
|
|
)
|
||
|
|
|
||
|
|
if default_query is None:
|
||
|
|
default_query = {"api-version": api_version}
|
||
|
|
else:
|
||
|
|
default_query = {"api-version": api_version, **default_query}
|
||
|
|
|
||
|
|
if base_url is None:
|
||
|
|
if azure_endpoint is None:
|
||
|
|
azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||
|
|
|
||
|
|
if azure_endpoint is None:
|
||
|
|
raise ValueError(
|
||
|
|
"Must provide one of the `base_url` or `azure_endpoint` arguments, or the `AZURE_OPENAI_ENDPOINT` environment variable"
|
||
|
|
)
|
||
|
|
|
||
|
|
if azure_deployment is not None:
|
||
|
|
base_url = f"{azure_endpoint}/openai/deployments/{azure_deployment}"
|
||
|
|
else:
|
||
|
|
base_url = f"{azure_endpoint}/openai"
|
||
|
|
else:
|
||
|
|
if azure_endpoint is not None:
|
||
|
|
raise ValueError("base_url and azure_endpoint are mutually exclusive")
|
||
|
|
|
||
|
|
if api_key is None:
|
||
|
|
# define a sentinel value to avoid any typing issues
|
||
|
|
api_key = API_KEY_SENTINEL
|
||
|
|
|
||
|
|
super().__init__(
|
||
|
|
api_key=api_key,
|
||
|
|
organization=organization,
|
||
|
|
base_url=base_url,
|
||
|
|
timeout=timeout,
|
||
|
|
max_retries=max_retries,
|
||
|
|
default_headers=default_headers,
|
||
|
|
default_query=default_query,
|
||
|
|
http_client=http_client,
|
||
|
|
_strict_response_validation=_strict_response_validation,
|
||
|
|
)
|
||
|
|
self._azure_ad_token = azure_ad_token
|
||
|
|
self._azure_ad_token_provider = azure_ad_token_provider
|
||
|
|
|
||
|
|
async def _get_azure_ad_token(self) -> str | None:
|
||
|
|
if self._azure_ad_token is not None:
|
||
|
|
return self._azure_ad_token
|
||
|
|
|
||
|
|
provider = self._azure_ad_token_provider
|
||
|
|
if provider is not None:
|
||
|
|
token = provider()
|
||
|
|
if inspect.isawaitable(token):
|
||
|
|
token = await token
|
||
|
|
if not token or not isinstance(token, str):
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected `azure_ad_token_provider` argument to return a string but it returned {token}",
|
||
|
|
)
|
||
|
|
return token
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
@override
|
||
|
|
async def _prepare_options(self, options: FinalRequestOptions) -> None:
|
||
|
|
headers: dict[str, str | Omit] = {**options.headers} if is_given(options.headers) else {}
|
||
|
|
options.headers = headers
|
||
|
|
|
||
|
|
azure_ad_token = await self._get_azure_ad_token()
|
||
|
|
if azure_ad_token is not None:
|
||
|
|
if headers.get("Authorization") is None:
|
||
|
|
headers["Authorization"] = f"Bearer {azure_ad_token}"
|
||
|
|
elif self.api_key is not API_KEY_SENTINEL:
|
||
|
|
if headers.get("api-key") is None:
|
||
|
|
headers["api-key"] = self.api_key
|
||
|
|
else:
|
||
|
|
# should never be hit
|
||
|
|
raise ValueError("Unable to handle auth")
|
||
|
|
|
||
|
|
return await super()._prepare_options(options)
|