Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: prioritize api_key over tenant_id for more Azure AD token provider #8318

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+18 −29
Open

feat: prioritize api_key over tenant_id for more Azure AD token provider #8318

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
16f158f
fix: prioritize api_key over tenant_id for Azure AD token provider
you-n-g Feb 6, 2025
96e9f57
Merge branch 'main' into main
you-n-g Feb 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions litellm/router_utils/client_initalization_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,8 @@ def set_client( # noqa: PLR0915
organization = get_secret_str(organization_env_name)
litellm_params["organization"] = organization
azure_ad_token_provider: Optional[Callable[[], str]] = None
if litellm_params.get("tenant_id"):
# If we have api_key, then we have higher priority
if not api_key and litellm_params.get("tenant_id"):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this logic is getting quite complicated, can we please refactor into a smaller function and add testing for this

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which section are you referring to (from which line to which line)?
Hi ,@krrishdholakia

this logic is getting quite complicated, can we please refactor into a smaller function and add testing for this

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What kind of test is expected? Do we any example that I can refer to?

verbose_router_logger.debug(
"Using Azure AD Token Provider for Azure Auth"
)
Expand Down Expand Up @@ -232,7 +233,7 @@ def set_client( # noqa: PLR0915
if azure_ad_token.startswith("oidc/"):
azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token)
elif (
azure_ad_token_provider is None
not api_key and azure_ad_token_provider is None
and litellm.enable_azure_ad_token_refresh is True
):
try:
Expand Down
42 changes: 15 additions & 27 deletions litellm/secret_managers/get_azure_ad_token_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import os
from typing import Callable, Union

from typing import Callable
from litellm._logging import verbose_logger


Expand All @@ -16,31 +15,20 @@ def get_azure_ad_token_provider() -> Callable[[], str]:
Returns:
Callable that returns a temporary authentication token.
"""
from azure.identity import (
ClientSecretCredential,
DefaultAzureCredential,
get_bearer_token_provider,
)
from azure.identity import get_bearer_token_provider
import azure.identity as identity
azure_scope = os.environ.get("AZURE_SCOPE", "https://cognitiveservices.azure.com/.default")
cred = os.environ.get("AZURE_CREDENTIAL", "ClientSecretCredential")

try:
credential: Union[ClientSecretCredential, DefaultAzureCredential] = (
ClientSecretCredential(
client_id=os.environ["AZURE_CLIENT_ID"],
client_secret=os.environ["AZURE_CLIENT_SECRET"],
tenant_id=os.environ["AZURE_TENANT_ID"],
)
)
except KeyError as e:
verbose_logger.exception(
"Missing environment variable required by Azure AD workflow. "
"DefaultAzureCredential will be used"
" {}".format(str(e))
cred_cls = getattr(identity, cred)
# ClientSecretCredential, DefaultAzureCredential, AzureCliCredential
if cred == "ClientSecretCredential":
credential = cred_cls(
client_id=os.environ["AZURE_CLIENT_ID"],
client_secret=os.environ["AZURE_CLIENT_SECRET"],
tenant_id=os.environ["AZURE_TENANT_ID"],
)
credential = DefaultAzureCredential()
except Exception:
raise
else:
credential = cred_cls()

return get_bearer_token_provider(
credential,
"https://cognitiveservices.azure.com/.default",
)
return get_bearer_token_provider(credential, azure_scope)