partners/openai + community: Async Azure AD token provider support for Azure OpenAI (#27488)

This PR introduces a new `azure_ad_async_token_provider` attribute to
the `AzureOpenAI` and `AzureChatOpenAI` classes in `partners/openai` and
`community` packages, given it's currently supported on `openai` package
as
[AsyncAzureADTokenProvider](https://github.com/openai/openai-python/blob/main/src/openai/lib/azure.py#L33)
type.

The reason for creating a new attribute is to avoid breaking changes.
Let's say you have an existing code that uses a `AzureOpenAI` or
`AzureChatOpenAI` instance to perform both sync and async operations.
The `azure_ad_token_provider` will work exactly as it is today, while
`azure_ad_async_token_provider` will override it for async requests.


If no one reviews your PR within a few days, please @-mention one of
baskaryan, efriis, eyurtsev, ccurme, vbarda, hwchase17.
This commit is contained in:
Fernando de Oliveira
2024-10-22 17:43:06 -04:00
committed by GitHub
parent 34684423bf
commit ab205e7389
6 changed files with 96 additions and 11 deletions

View File

@@ -8,6 +8,7 @@ from typing import (
AbstractSet,
Any,
AsyncIterator,
Awaitable,
Callable,
Collection,
Dict,
@@ -804,7 +805,13 @@ class AzureOpenAI(BaseOpenAI):
azure_ad_token_provider: Union[Callable[[], str], None] = None
"""A function that returns an Azure Active Directory token.
Will be invoked on every request.
Will be invoked on every sync request. For async requests,
will be invoked if `azure_ad_async_token_provider` is not provided.
"""
azure_ad_async_token_provider: Union[Callable[[], Awaitable[str]], None] = None
"""A function that returns an Azure Active Directory token.
Will be invoked on every async request.
"""
openai_api_type: str = ""
"""Legacy, for openai<1.0.0 support."""
@@ -922,6 +929,12 @@ class AzureOpenAI(BaseOpenAI):
"http_client": values["http_client"],
}
values["client"] = openai.AzureOpenAI(**client_params).completions
azure_ad_async_token_provider = values["azure_ad_async_token_provider"]
if azure_ad_async_token_provider:
client_params["azure_ad_token_provider"] = azure_ad_async_token_provider
values["async_client"] = openai.AsyncAzureOpenAI(
**client_params
).completions