diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index bb60680ebc1..5317a9a0ec7 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -125,22 +125,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): def __init__(self) -> None: super().__init__() - def validate_environment(self, api_key, azure_ad_token, azure_ad_token_provider): - headers = { - "content-type": "application/json", - } - if api_key is not None: - headers["api-key"] = api_key - elif azure_ad_token is not None: - if azure_ad_token.startswith("oidc/"): - azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) - headers["Authorization"] = f"Bearer {azure_ad_token}" - elif azure_ad_token_provider is not None: - azure_ad_token = azure_ad_token_provider() - headers["Authorization"] = f"Bearer {azure_ad_token}" - - return headers - def make_sync_azure_openai_chat_completion_request( self, azure_client: AzureOpenAI, @@ -242,6 +226,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): azure_ad_token_provider=azure_ad_token_provider, acompletion=acompletion, client=client, + litellm_params=litellm_params, ) data = {"model": None, "messages": messages, **optional_params} diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 012f47c8517..4ebd54e8fcb 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -158,13 +158,27 @@ def get_azure_ad_token_from_username_password( return token_provider -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) +def get_azure_ad_token_from_oidc( + azure_ad_token: str, + azure_client_id: Optional[str], + azure_tenant_id: Optional[str], +) -> str: + """ + Get Azure AD token from OIDC token + + Args: + azure_ad_token: str + azure_client_id: Optional[str] + azure_tenant_id: Optional[str] + + Returns: + `azure_ad_token_access_token` - str + """ azure_authority_host = os.getenv( "AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com" ) - + azure_client_id = azure_client_id or os.getenv("AZURE_CLIENT_ID") + azure_tenant_id = azure_tenant_id or os.getenv("AZURE_TENANT_ID") if azure_client_id is None or azure_tenant_id is None: raise AzureOpenAIError( status_code=422, @@ -341,7 +355,11 @@ class BaseAzureLLM(BaseOpenAILLM): if azure_ad_token is not None and azure_ad_token.startswith("oidc/"): verbose_logger.debug("Using Azure OIDC Token for Azure Auth") - azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) + azure_ad_token = get_azure_ad_token_from_oidc( + azure_ad_token=azure_ad_token, + azure_client_id=client_id, + azure_tenant_id=tenant_id, + ) elif ( not api_key and azure_ad_token_provider is None @@ -402,6 +420,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_version: str, max_retries: int, timeout: Union[float, httpx.Timeout], + litellm_params: dict, api_key: Optional[str], azure_ad_token: Optional[str], azure_ad_token_provider: Optional[Callable[[], str]], @@ -409,6 +428,8 @@ class BaseAzureLLM(BaseOpenAILLM): client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[AzureOpenAI, AsyncAzureOpenAI]: ## build base url - assume api base includes resource name + tenant_id = litellm_params.get("tenant_id", os.getenv("AZURE_TENANT_ID")) + client_id = litellm_params.get("client_id", os.getenv("AZURE_CLIENT_ID")) if client is None: if not api_base.endswith("/"): api_base += "/" @@ -425,7 +446,11 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params["api_key"] = api_key elif azure_ad_token is not None: if azure_ad_token.startswith("oidc/"): - azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) + azure_ad_token = get_azure_ad_token_from_oidc( + azure_ad_token=azure_ad_token, + azure_client_id=client_id, + azure_tenant_id=tenant_id, + ) azure_client_params["azure_ad_token"] = azure_ad_token if azure_ad_token_provider is not None: diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 8301c4d617d..a44f9045712 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -72,6 +72,7 @@ class AzureTextCompletion(BaseAzureLLM): azure_ad_token=azure_ad_token, azure_ad_token_provider=azure_ad_token_provider, acompletion=acompletion, + litellm_params=litellm_params, ) data = {"model": None, "prompt": prompt, **optional_params} diff --git a/tests/litellm/llms/azure/test_azure_common_utils.py b/tests/litellm/llms/azure/test_azure_common_utils.py index 42b5903ee84..ac5be085776 100644 --- a/tests/litellm/llms/azure/test_azure_common_utils.py +++ b/tests/litellm/llms/azure/test_azure_common_utils.py @@ -159,7 +159,9 @@ def test_initialize_with_username_password(setup_mocks): assert "azure_ad_token_provider" in result -def test_initialize_with_oidc_token(setup_mocks): +def test_initialize_with_oidc_token(setup_mocks, monkeypatch): + monkeypatch.delenv("AZURE_CLIENT_ID", raising=False) + monkeypatch.delenv("AZURE_TENANT_ID", raising=False) # Test with azure_ad_token that starts with "oidc/" result = BaseAzureLLM().initialize_azure_sdk_client( litellm_params={"azure_ad_token": "oidc/test-token"}, @@ -171,7 +173,89 @@ def test_initialize_with_oidc_token(setup_mocks): ) # Verify that get_azure_ad_token_from_oidc was called - setup_mocks["oidc_token"].assert_called_once_with("oidc/test-token") + setup_mocks["oidc_token"].assert_called_once_with( + azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None + ) + + # Verify expected result + assert result["azure_ad_token"] == "mock-oidc-token" + + +def test_initialize_with_oidc_token_and_client_params(setup_mocks): + # Test with azure_ad_token that starts with "oidc/" and explicit client/tenant IDs + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={ + "azure_ad_token": "oidc/test-token", + "client_id": "test-client-id", + "tenant_id": "test-tenant-id", + }, + api_key=None, + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version=None, + is_async=False, + ) + + # Verify that get_azure_ad_token_from_oidc was called with the correct parameters + setup_mocks["oidc_token"].assert_called_once_with( + azure_ad_token="oidc/test-token", + azure_client_id="test-client-id", + azure_tenant_id="test-tenant-id", + ) + + # Verify expected result + assert result["azure_ad_token"] == "mock-oidc-token" + + +def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch): + # Set environment variables + monkeypatch.setenv("AZURE_CLIENT_ID", "env-client-id") + monkeypatch.setenv("AZURE_TENANT_ID", "env-tenant-id") + + # Test with azure_ad_token that starts with "oidc/" but no explicit client/tenant IDs + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={ + "azure_ad_token": "oidc/test-token", + }, + api_key=None, + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version=None, + is_async=False, + ) + + # Verify that get_azure_ad_token_from_oidc was called with environment variables + setup_mocks["oidc_token"].assert_called_once_with( + azure_ad_token="oidc/test-token", + azure_client_id="env-client-id", + azure_tenant_id="env-tenant-id", + ) + + # Verify expected result + assert result["azure_ad_token"] == "mock-oidc-token" + + +def test_initialize_with_oidc_token_no_credentials(setup_mocks, monkeypatch): + # Clear environment variables + monkeypatch.delenv("AZURE_CLIENT_ID", raising=False) + monkeypatch.delenv("AZURE_TENANT_ID", raising=False) + + # Test with azure_ad_token that starts with "oidc/" but no credentials anywhere + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={ + "azure_ad_token": "oidc/test-token", + }, + api_key=None, + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version=None, + is_async=False, + ) + + # Verify that get_azure_ad_token_from_oidc was called with None values + setup_mocks["oidc_token"].assert_called_once_with( + azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None + ) # Verify expected result assert result["azure_ad_token"] == "mock-oidc-token" diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py index fd1adeb9648..cd9230d937e 100644 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ b/tests/litellm_utils_tests/test_secret_manager.py @@ -140,7 +140,11 @@ def test_oidc_circleci_with_azure(): # TODO: Switch to our own Azure account, currently using ai.moda's account os.environ["AZURE_TENANT_ID"] = "17c0a27a-1246-4aa1-a3b6-d294e80e783c" os.environ["AZURE_CLIENT_ID"] = "4faf5422-b2bd-45e8-a6d7-46543a38acd0" - azure_ad_token = get_azure_ad_token_from_oidc("oidc/circleci/") + azure_ad_token = get_azure_ad_token_from_oidc( + azure_ad_token="oidc/circleci/", + azure_client_id=None, + azure_tenant_id=None, + ) print(f"secret_val: {redact_oidc_signature(azure_ad_token)}")