From 640b955e12aab32585290a15d5a690b520e7a109 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Mon, 27 May 2024 16:33:37 +0000 Subject: [PATCH 1/6] feat(util.py/azure.py): Add OIDC support when running in Azure Kubernetes Service (AKS). --- litellm/llms/azure.py | 7 ++++--- litellm/utils.py | 8 ++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 02fe4a08f219..4d258cff2913 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -136,9 +136,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", @@ -153,7 +154,7 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): ) 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", diff --git a/litellm/utils.py b/litellm/utils.py index ea0f46c144b2..8c18ab88efb4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10085,6 +10085,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") From 73d9d0d7db78afee938e7293a3245bf7c167a855 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Mon, 27 May 2024 17:35:37 +0000 Subject: [PATCH 2/6] feat(azure.py): Add Azure AD cred caching for OIDC flow. --- litellm/llms/azure.py | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 4d258cff2913..be8220b0ca38 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -18,6 +18,9 @@ from openai import AzureOpenAI, AsyncAzureOpenAI import uuid import os +from litellm.caching import DualCache + +azure_ad_cache = DualCache() class AzureOpenAIError(Exception): @@ -153,6 +156,17 @@ 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"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", data={ @@ -170,14 +184,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 possible_azure_ad_token is None: + if azure_ad_token_access_token is None: raise AzureOpenAIError( - status_code=422, message="Azure AD Token not returned" + status_code=422, message="Azure AD Token access_token not returned" ) - return possible_azure_ad_token + if azure_ad_token_expires_in is None: + raise AzureOpenAIError( + status_code=422, message="Azure AD Token expires_in 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 azure_ad_token_access_token class AzureChatCompletion(BaseLLM): From 7154c2590363ae5dca6fd5d4cc7e9d1262fed762 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 15:51:59 +0000 Subject: [PATCH 3/6] Revert "feat(azure.py): Add Azure AD cred caching for OIDC flow." This reverts commit 73d9d0d7db78afee938e7293a3245bf7c167a855. --- litellm/llms/azure.py | 31 ++++--------------------------- 1 file changed, 4 insertions(+), 27 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index be8220b0ca38..4d258cff2913 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -18,9 +18,6 @@ from openai import AzureOpenAI, AsyncAzureOpenAI import uuid import os -from litellm.caching import DualCache - -azure_ad_cache = DualCache() class AzureOpenAIError(Exception): @@ -156,17 +153,6 @@ 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"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", data={ @@ -184,23 +170,14 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): message=req_token.text, ) - 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) + possible_azure_ad_token = req_token.json().get("access_token", None) - if azure_ad_token_access_token is None: + if possible_azure_ad_token is None: raise AzureOpenAIError( - status_code=422, message="Azure AD Token access_token not returned" + status_code=422, message="Azure AD Token not returned" ) - if azure_ad_token_expires_in is None: - raise AzureOpenAIError( - status_code=422, message="Azure AD Token expires_in 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 azure_ad_token_access_token + return possible_azure_ad_token class AzureChatCompletion(BaseLLM): From 95a4bc46d457e75ecc85cfeec102f995c074eac8 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 15:52:07 +0000 Subject: [PATCH 4/6] Revert "feat(util.py/azure.py): Add OIDC support when running in Azure Kubernetes Service (AKS)." This reverts commit 640b955e12aab32585290a15d5a690b520e7a109. --- litellm/llms/azure.py | 7 +++---- litellm/utils.py | 8 -------- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 4d258cff2913..02fe4a08f219 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -136,10 +136,9 @@ 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_id = os.getenv("AZURE_TENANT_ID", None) - azure_authority_host = os.getenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") + azure_tenant = os.getenv("AZURE_TENANT_ID", None) - if azure_client_id is None or azure_tenant_id is None: + if azure_client_id is None or azure_tenant is None: raise AzureOpenAIError( status_code=422, message="AZURE_CLIENT_ID and AZURE_TENANT_ID must be set", @@ -154,7 +153,7 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): ) req_token = httpx.post( - f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", + f"https://login.microsoftonline.com/{azure_tenant}/oauth2/v2.0/token", data={ "client_id": azure_client_id, "grant_type": "client_credentials", diff --git a/litellm/utils.py b/litellm/utils.py index 8c18ab88efb4..ea0f46c144b2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10085,14 +10085,6 @@ 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") From a31fa5fbc80c29ef3b33ca675e9ff9b3c85b8af3 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Mon, 27 May 2024 16:33:37 +0000 Subject: [PATCH 5/6] feat(util.py/azure.py): Add OIDC support when running in Azure Kubernetes Service (AKS). --- litellm/llms/azure.py | 7 ++++--- litellm/utils.py | 8 ++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 834fcbea96b9..e8bcaff64b63 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -309,9 +309,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", @@ -326,7 +327,7 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str): ) 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", 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") From d9503737689fe7b488d6e609bcbd3d02df03f9e4 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Mon, 27 May 2024 17:35:37 +0000 Subject: [PATCH 6/6] feat(azure.py): Add Azure AD cred caching for OIDC flow. --- litellm/llms/azure.py | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index e8bcaff64b63..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): @@ -326,6 +329,17 @@ 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"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", data={ @@ -343,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):