diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 834fcbea96b9..46ab62a8d3ef 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -36,6 +36,9 @@ AsyncAssistantStreamManager, AssistantStreamManager, ) +from litellm.caching import DualCache + +azure_ad_cache = DualCache() class AzureOpenAIError(Exception): @@ -309,9 +312,10 @@ def select_azure_base_url_or_endpoint(azure_client_params: dict): def get_azure_ad_token_from_oidc(azure_ad_token: str): azure_client_id = os.getenv("AZURE_CLIENT_ID", None) - azure_tenant = os.getenv("AZURE_TENANT_ID", None) + azure_tenant_id = os.getenv("AZURE_TENANT_ID", None) + azure_authority_host = os.getenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") - if azure_client_id is None or azure_tenant is None: + if azure_client_id is None or azure_tenant_id is None: raise AzureOpenAIError( status_code=422, message="AZURE_CLIENT_ID and AZURE_TENANT_ID must be set", @@ -325,8 +329,19 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): message="OIDC token could not be retrieved from secret manager.", ) + azure_ad_token_cache_key = json.dumps({ + "azure_client_id": azure_client_id, + "azure_tenant_id": azure_tenant_id, + "azure_authority_host": azure_authority_host, + "oidc_token": oidc_token, + }) + + azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key) + if azure_ad_token_access_token is not None: + return azure_ad_token_access_token + req_token = httpx.post( - f"https://login.microsoftonline.com/{azure_tenant}/oauth2/v2.0/token", + f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", data={ "client_id": azure_client_id, "grant_type": "client_credentials", @@ -342,12 +357,23 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): message=req_token.text, ) - possible_azure_ad_token = req_token.json().get("access_token", None) + azure_ad_token_json = req_token.json() + azure_ad_token_access_token = azure_ad_token_json.get("access_token", None) + azure_ad_token_expires_in = azure_ad_token_json.get("expires_in", None) + + if azure_ad_token_access_token is None: + raise AzureOpenAIError( + status_code=422, message="Azure AD Token access_token not returned" + ) + + if azure_ad_token_expires_in is None: + raise AzureOpenAIError( + status_code=422, message="Azure AD Token expires_in not returned" + ) - if possible_azure_ad_token is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token not returned") + azure_ad_cache.set_cache(key=azure_ad_token_cache_key, value=azure_ad_token_access_token, ttl=azure_ad_token_expires_in) - return possible_azure_ad_token + return azure_ad_token_access_token class AzureChatCompletion(BaseLLM): diff --git a/litellm/utils.py b/litellm/utils.py index 5e85419dcd0f..b872687b55e1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10050,6 +10050,14 @@ def get_secret( return oidc_token else: raise ValueError("Github OIDC provider failed") + elif oidc_provider == "azure": + # https://azure.github.io/azure-workload-identity/docs/quick-start.html + azure_federated_token_file = os.getenv("AZURE_FEDERATED_TOKEN_FILE") + if azure_federated_token_file is None: + raise ValueError("AZURE_FEDERATED_TOKEN_FILE not found in environment") + with open(azure_federated_token_file, "r") as f: + oidc_token = f.read() + return oidc_token else: raise ValueError("Unsupported OIDC provider")