[Feature]: Support azure workload identities to create azure ad tokens on proxy

[krrishdholakia]

The Feature

"we use Azure Active Directory Tokens. Those tokens are defined per deployment and are only valid for a certain amount of time. We cannot add the azure_ad_token flag in the config, as the token would expire after a couple of hours."

Motivation, pitch

User trying to use proxy to call multiple llm's in openai format (in projects in diff languages)

Twitter / LinkedIn details

No response

[krrishdholakia]

Sample code:

from fastapi import FastAPI
from azure.identity import DefaultAzureCredential
from azure.keyvault.secrets import SecretClient
from azure.core.exceptions import HttpResponseError

app = FastAPI()

# Function to get a secret from Azure Key Vault
def get_secret(secret_name):
    credential = DefaultAzureCredential()
    keyvault_name = "<your-keyvault-name>"
    keyvault_uri = f"https://{keyvault_name}.vault.azure.net"
    client = SecretClient(vault_url=keyvault_uri, credential=credential)
    try:
        secret = client.get_secret(secret_name)
        return secret.value
    except HttpResponseError:
        return None

# Example route to demonstrate obtaining an Azure AD token
@app.get("/get_token")
async def get_token():
    tenant_id = get_secret("AzureAd-TenantId-SecretName")
    client_id = get_secret("AzureAd-ClientId-SecretName")
    credential = DefaultAzureCredential()
    token = credential.get_token("https://graph.microsoft.com/.default")
    return {"access_token": token.token}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

[konsti]

Some more inspiration:

from azure.identity import DefaultAzureCredential
import time

class AzureTokenManager:
    def __init__(self, scopes):
        self.scopes = scopes
        self.credential = DefaultAzureCredential()
        self.token = None
        self.expires_on = 0

    def get_token(self, force_refresh=False):
        current_time = time.time()
        # Refresh the token if force_refresh is True or the token is expiring within 720 seconds
        if force_refresh or self.token is None or current_time > self.expires_on - 720:
            result = self.credential.get_token(self.scopes)
            self.token = result.token
            self.expires_on = result.expires_on
        return self.token

# Usage
scopes = "https://cognitiveservices.azure.com/.default"
token_manager = AzureTokenManager(scopes)

token = token_manager.get_token()
token = token_manager.get_token(force_refresh=True)

[olad32]

Even more inspiration, AzureOpenAI python client actually had this feature built in, example here : https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1#scenario-3-authenticating-with-entra-id---service-principal

[RyoYang]

Do we have any updates on this feature? Like support azure_ad_token_provider in litellm proxy?

[krrishdholakia]

@RyoYang we already support azure ad token on litellm proxy?

[Manouchehri]

See #3505, I think that’s what you want?

[RyoYang]

@Manouchehri Thank you for the quick response. From my perspective, this PR indeed supports OIDC authentication for Google, GitHub, and CircleCI but not Azure workload identities. Could you please provide a quick example of how to configure an Azure workload identity case if I misunderstood?

My ideal scenario is to use the Litellm proxy when using the OpenAI SDK, similar to Scenario 4 in this URL

[Manouchehri]

Where/what are you running LiteLLM's proxy in?

[RyoYang]

Oh. I am running LiteLLM's Proxy in AKS as model proxy, and other pods will use openai SDK to chat with endpoints by using litellm proxy service address.

[Manouchehri]

Ah, so all we're missing is fetching an OIDC token I think. We would need to add it to:

litellm/litellm/utils.py

Lines 10030 to 10097 in a6a84e5

	# Example: oidc/google/https://bedrock-runtime.us-east-1.amazonaws.com/model/stability.stable-diffusion-xl-v1/invoke
	if secret_name.startswith("oidc/"):
	secret_name_split = secret_name.replace("oidc/", "")
	oidc_provider, oidc_aud = secret_name_split.split("/", 1)
	# TODO: Add caching for HTTP requests
	if oidc_provider == "google":
	oidc_token = oidc_cache.get_cache(key=secret_name)
	if oidc_token is not None:
	return oidc_token
	oidc_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0))
	# https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature
	response = oidc_client.get(
	"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity",
	params={"audience": oidc_aud},
	headers={"Metadata-Flavor": "Google"},
	)
	if response.status_code == 200:
	oidc_token = response.text
	oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=3600 - 60)
	return oidc_token
	else:
	raise ValueError("Google OIDC provider failed")
	elif oidc_provider == "circleci":
	# https://circleci.com/docs/openid-connect-tokens/
	env_secret = os.getenv("CIRCLE_OIDC_TOKEN")
	if env_secret is None:
	raise ValueError("CIRCLE_OIDC_TOKEN not found in environment")
	return env_secret
	elif oidc_provider == "circleci_v2":
	# https://circleci.com/docs/openid-connect-tokens/
	env_secret = os.getenv("CIRCLE_OIDC_TOKEN_V2")
	if env_secret is None:
	raise ValueError("CIRCLE_OIDC_TOKEN_V2 not found in environment")
	return env_secret
	elif oidc_provider == "github":
	# https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers#using-custom-actions
	actions_id_token_request_url = os.getenv("ACTIONS_ID_TOKEN_REQUEST_URL")
	actions_id_token_request_token = os.getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
	if (
	actions_id_token_request_url is None
	or actions_id_token_request_token is None
	):
	raise ValueError(
	"ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment"
	)
	oidc_token = oidc_cache.get_cache(key=secret_name)
	if oidc_token is not None:
	return oidc_token
	oidc_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0))
	response = oidc_client.get(
	actions_id_token_request_url,
	params={"audience": oidc_aud},
	headers={
	"Authorization": f"Bearer {actions_id_token_request_token}",
	"Accept": "application/json; api-version=2.0",
	},
	)
	if response.status_code == 200:
	oidc_token = response.text["value"]
	oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5)
	return oidc_token
	else:
	raise ValueError("Github OIDC provider failed")
	else:
	raise ValueError("Unsupported OIDC provider")

Right now, the config for using OIDC with Google Cloud Run is like:

model_list:
  - model_name: gpt-4-0125-preview
    litellm_params:
      model: azure/gpt-4-0125-preview
      api_version: "2024-05-01-preview"
      azure_ad_token: "oidc/google/https://litellm.example.com"
      api_base: "https://example.openai.azure.com"
    model_info:
      base_model: azure/gpt-4-0125-preview

Instead of google, we should add support for like aks.

[olad32]

Hi, i suggest to use official AzureOpenAI python client as it actually had this feature built in (token retrieve and use, via the azure_ad_token_provider parameter), example use here : https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1#scenario-3-authenticating-with-entra-id---service-principal and https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1?tab=readme-ov-file#scenario-4-authenticating-with-entra-id---managed-identity

[Manouchehri]

Eh, using the official python client like that will make us lose out on being able to authenticate AKS -> Amazon Bedrock I think. It'd also lose out on being able to cache creds.

[RyoYang]

Hi, i suggest to use official AzureOpenAI python client as it actually had this feature built in (token retrieve and use, via the azure_ad_token_provider parameter), example use here : https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1#scenario-3-authenticating-with-entra-id---service-principal and https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1?tab=readme-ov-file#scenario-4-authenticating-with-entra-id---managed-identity嗨，我建议使用官方的 AzureOpenAI Python 客户端，因为它实际上已经内置了这个功能（通过 azure_ad_token_provider 参数检索和使用令牌），示例用法在这里：https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1#scenario-3-authenticating-with-entra-id---service-principal 和 https://github.com/LazaUK/AOAI-EntraIDAuth-SDKv1?tab=readme-ov-file#scenario-4-authenticating-with-entra-id---managed-identity

I have tried to use openai sdk by using azure_ad_token_provider, but it seems like the proxy will still try to use api_key or ad_token to do auth... Can you please show me how to config litellm proxy in the example above?

[RyoYang]

Ah, so all we're missing is fetching an OIDC token I think. We would need to add it to:

litellm/litellm/utils.py

Lines 10030 to 10097 in a6a84e5

	# Example: oidc/google/https://bedrock-runtime.us-east-1.amazonaws.com/model/stability.stable-diffusion-xl-v1/invoke
	if secret_name.startswith("oidc/"):
	secret_name_split = secret_name.replace("oidc/", "")
	oidc_provider, oidc_aud = secret_name_split.split("/", 1)
	# TODO: Add caching for HTTP requests
	if oidc_provider == "google":
	oidc_token = oidc_cache.get_cache(key=secret_name)
	if oidc_token is not None:
	return oidc_token
	oidc_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0))
	# https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature
	response = oidc_client.get(
	"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity",
	params={"audience": oidc_aud},
	headers={"Metadata-Flavor": "Google"},
	)
	if response.status_code == 200:
	oidc_token = response.text
	oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=3600 - 60)
	return oidc_token
	else:
	raise ValueError("Google OIDC provider failed")
	elif oidc_provider == "circleci":
	# https://circleci.com/docs/openid-connect-tokens/
	env_secret = os.getenv("CIRCLE_OIDC_TOKEN")
	if env_secret is None:
	raise ValueError("CIRCLE_OIDC_TOKEN not found in environment")
	return env_secret
	elif oidc_provider == "circleci_v2":
	# https://circleci.com/docs/openid-connect-tokens/
	env_secret = os.getenv("CIRCLE_OIDC_TOKEN_V2")
	if env_secret is None:
	raise ValueError("CIRCLE_OIDC_TOKEN_V2 not found in environment")
	return env_secret
	elif oidc_provider == "github":
	# https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers#using-custom-actions
	actions_id_token_request_url = os.getenv("ACTIONS_ID_TOKEN_REQUEST_URL")
	actions_id_token_request_token = os.getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
	if (
	actions_id_token_request_url is None
	or actions_id_token_request_token is None
	):
	raise ValueError(
	"ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment"
	)
	oidc_token = oidc_cache.get_cache(key=secret_name)
	if oidc_token is not None:
	return oidc_token
	oidc_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0))
	response = oidc_client.get(
	actions_id_token_request_url,
	params={"audience": oidc_aud},
	headers={
	"Authorization": f"Bearer {actions_id_token_request_token}",
	"Accept": "application/json; api-version=2.0",
	},
	)
	if response.status_code == 200:
	oidc_token = response.text["value"]
	oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5)
	return oidc_token
	else:
	raise ValueError("Github OIDC provider failed")
	else:
	raise ValueError("Unsupported OIDC provider")

Right now, the config for using OIDC with Google Cloud Run is like:

model_list:
  - model_name: gpt-4-0125-preview
    litellm_params:
      model: azure/gpt-4-0125-preview
      api_version: "2024-05-01-preview"
      azure_ad_token: "oidc/google/https://litellm.example.com"
      api_base: "https://example.openai.azure.com"
    model_info:
      base_model: azure/gpt-4-0125-preview

Instead of google, we should add support for like aks.

Yea! Exactly! Does the oidc URL will be the oidc issuer from (az aks show --name myAKScluster --resource-group myResourceGroup --query "oidcIssuerProfile.issuerUrl" -o tsv)?

[Manouchehri]

I don't need the issuer URL atm, right now I'm trying to figure out how we can get the OIDC token without resorting to the Python SDK if possible.

[Manouchehri]

Ah, looks like it might be AZURE_FEDERATED_TOKEN_FILE.

[Manouchehri]

@RyoYang or @olad32: Could you share AZURE_AUTHORITY_HOST and the contents of AZURE_FEDERATED_TOKEN_FILE with me?

For AZURE_FEDERATED_TOKEN_FILE, make sure to remove the part after the last dot so you don't give me access to your real infrastructure.

e.g. if you have:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c

Please only shared with me:

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SIGNATURE_REMOVED

[Manouchehri]

@RyoYang and @olad32, can you help me over at #3861?