tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 02:25:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14470 from BerriAI/litellm_dev_09_11_2025_p1

Browse Source 
AzureAD Default credentials - select credential type based on environment
This commit is contained in:
Krish Dholakia
2025-10-08 19:03:07 -07:00
committed by GitHub
parent c4022ade49 12a1d081ee
commit adbdf9dae9
5 changed files with 116 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/azure/common_utils.py
+8 -1
View File
@@ -365,6 +365,11 @@ def get_azure_ad_token(
azure_ad_token_provider = get_azure_ad_token_provider(azure_scope=scope)
except ValueError:
verbose_logger.debug("Azure AD Token Provider could not be used.")
except Exception as e:
verbose_logger.error(
f"Error calling Azure AD token provider: {str(e)}. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential"
)
raise e
#########################################################
# If litellm.enable_azure_ad_token_refresh is True and no other token provider is available,
@@ -561,7 +566,9 @@ class BaseAzureLLM(BaseOpenAILLM):
"Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth"
)
try:
azure_ad_token_provider = get_azure_ad_token_provider(azure_scope=scope)
azure_ad_token_provider = get_azure_ad_token_provider(
azure_scope=scope,
)
except ValueError:
verbose_logger.debug("Azure AD Token Provider could not be used.")
if api_version is None:
litellm/secret_managers/get_azure_ad_token_provider.py
+33 -3
View File
@@ -1,11 +1,36 @@
import os
from typing import Any, Callable, Optional, Union
from litellm._logging import verbose_logger
from litellm.types.secret_managers.get_azure_ad_token_provider import (
AzureCredentialType,
)
def infer_credential_type_from_environment() -> AzureCredentialType:
if (
os.environ.get("AZURE_CLIENT_ID")
and os.environ.get("AZURE_CLIENT_SECRET")
and os.environ.get("AZURE_TENANT_ID")
):
return AzureCredentialType.ClientSecretCredential
elif os.environ.get("AZURE_CLIENT_ID"):
return AzureCredentialType.ManagedIdentityCredential
elif (
os.environ.get("AZURE_CLIENT_ID")
and os.environ.get("AZURE_TENANT_ID")
and os.environ.get("AZURE_CERTIFICATE_PATH")
and os.environ.get("AZURE_CERTIFICATE_PASSWORD")
):
return AzureCredentialType.CertificateCredential
elif os.environ.get("AZURE_CERTIFICATE_PASSWORD"):
return AzureCredentialType.CertificateCredential
elif os.environ.get("AZURE_CERTIFICATE_PATH"):
return AzureCredentialType.CertificateCredential
else:
return AzureCredentialType.DefaultAzureCredential
def get_azure_ad_token_provider(
azure_scope: Optional[str] = None,
azure_credential: Optional[AzureCredentialType] = None,
@@ -42,9 +67,14 @@ def get_azure_ad_token_provider(
)
cred: str = (
azure_credential.value if azure_credential else None
or os.environ.get("AZURE_CREDENTIAL", AzureCredentialType.ClientSecretCredential)
or AzureCredentialType.ClientSecretCredential
azure_credential.value
if azure_credential
else None
or os.environ.get("AZURE_CREDENTIAL")
or infer_credential_type_from_environment()
)
verbose_logger.info(
f"For Azure AD Token Provider, choosing credential type: {cred}"
)
credential: Optional[
Union[
tests/llm_translation/test_gemini.py
+34 -20
View File
@@ -267,7 +267,11 @@ def test_gemini_image_generation():
assert len(response.choices[0].message.images) > 0
assert response.choices[0].message.images[0]["image_url"] is not None
assert response.choices[0].message.images[0]["image_url"]["url"] is not None
assert response.choices[0].message.images[0]["image_url"]["url"].startswith("data:image/png;base64,")
assert (
response.choices[0]
.message.images[0]["image_url"]["url"]
.startswith("data:image/png;base64,")
)
def test_gemini_2_5_flash_image_preview():
@@ -772,7 +776,8 @@ def test_system_message_with_no_user_message():
assert response is not None
assert response.choices[0].message.content is not None
def get_current_weather(location, unit="fahrenheit"):
"""Get the current weather in a given location"""
if "tokyo" in location.lower():
@@ -889,9 +894,9 @@ def test_gemini_reasoning_effort_minimal():
# Test with different Gemini models to verify model-specific mapping
test_cases = [
("gemini/gemini-2.5-flash", 1), # Flash: minimum 1 token
("gemini/gemini-2.5-pro", 128), # Pro: minimum 128 tokens
("gemini/gemini-2.5-flash-lite", 512), # Flash-Lite: minimum 512 tokens
("gemini/gemini-2.5-flash", 1), # Flash: minimum 1 token
("gemini/gemini-2.5-pro", 128), # Pro: minimum 128 tokens
("gemini/gemini-2.5-flash-lite", 512), # Flash-Lite: minimum 512 tokens
]
for model, expected_min_budget in test_cases:
@@ -904,24 +909,32 @@ def test_gemini_reasoning_effort_minimal():
"reasoning_effort": "minimal",
},
)
# Verify that the thinking config is set correctly
request_body = raw_request["raw_request_body"]
assert "generationConfig" in request_body, f"Model {model} should have generationConfig"
assert (
"generationConfig" in request_body
), f"Model {model} should have generationConfig"
generation_config = request_body["generationConfig"]
assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig"
assert (
"thinkingConfig" in generation_config
), f"Model {model} should have thinkingConfig"
thinking_config = generation_config["thinkingConfig"]
assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget"
assert (
"thinkingBudget" in thinking_config
), f"Model {model} should have thinkingBudget"
actual_budget = thinking_config["thinkingBudget"]
assert actual_budget == expected_min_budget, \
f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}"
assert (
actual_budget == expected_min_budget
), f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}"
# Verify that includeThoughts is True for minimal reasoning effort
assert thinking_config.get("includeThoughts", True), \
f"Model {model} should have includeThoughts=True for minimal reasoning effort"
assert thinking_config.get(
"includeThoughts", True
), f"Model {model} should have includeThoughts=True for minimal reasoning effort"
# Test with unknown model (should use generic fallback)
try:
@@ -933,13 +946,14 @@ def test_gemini_reasoning_effort_minimal():
"reasoning_effort": "minimal",
},
)
request_body = raw_request["raw_request_body"]
generation_config = request_body["generationConfig"]
thinking_config = generation_config["thinkingConfig"]
# Should use generic fallback (128 tokens)
assert thinking_config["thinkingBudget"] == 128, \
"Unknown model should use generic fallback of 128 tokens"
assert (
thinking_config["thinkingBudget"] == 128
), "Unknown model should use generic fallback of 128 tokens"
except Exception as e:
# If return_raw_request doesn't work for unknown models, that's okay
# The important part is that our known models work correctly
tests/local_testing/test_amazing_vertex_completion.py
+12 -13
View File
@@ -397,7 +397,7 @@ async def test_async_vertexai_response():
| litellm.vertex_text_models
| litellm.vertex_code_text_models
)
test_models = random.sample(list(test_models), 1)
test_models += list(litellm.vertex_language_models) # always test gemini-pro
for model in test_models:
@@ -504,7 +504,6 @@ async def test_async_vertexai_streaming_response():
pytest.fail(f"An exception occurred: {e}")
@pytest.mark.parametrize("load_pdf", [False]) # True,
@pytest.mark.flaky(retries=3, delay=1)
def test_completion_function_plus_pdf(load_pdf):
@@ -547,6 +546,7 @@ def test_completion_function_plus_pdf(load_pdf):
except Exception as e:
pytest.fail("Got={}".format(str(e)))
def encode_image(image_path):
import base64
@@ -910,7 +910,10 @@ async def test_partner_models_httpx(model, region, sync_mode):
[
("vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas", "us-east5"),
("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"),
("vertex_ai/mistral-large-2411", "us-central1"), # critical - we had this issue: https://github.com/BerriAI/litellm/issues/13888
(
"vertex_ai/mistral-large-2411",
"us-central1",
), # critical - we had this issue: https://github.com/BerriAI/litellm/issues/13888
("vertex_ai/openai/gpt-oss-20b-maas", "us-central1"),
],
)
@@ -3773,7 +3776,7 @@ def test_vertex_ai_gemini_audio_ogg():
@pytest.mark.asyncio
async def test_vertex_ai_deepseek():
"""Test that deepseek models use the correct v1 API endpoint instead of v1beta1."""
#load_vertex_ai_credentials()
# load_vertex_ai_credentials()
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@@ -3786,21 +3789,17 @@ async def test_vertex_ai_deepseek():
{
"message": {
"role": "assistant",
"content": "Hello! How can I help you today?"
"content": "Hello! How can I help you today?",
},
"index": 0,
"finish_reason": "stop"
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30
},
"model": "deepseek-ai/deepseek-r1-0528-maas"
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
"model": "deepseek-ai/deepseek-r1-0528-maas",
}
mock_response.status_code = 200
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = await acompletion(
model="vertex_ai/deepseek-ai/deepseek-r1-0528-maas",
tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py
+29
View File
@@ -214,3 +214,32 @@ class TestGetAzureAdTokenProvider:
# Test that the returned callable works
token = result()
assert token == "mock-certificate-token"
@patch.dict(os.environ, {}, clear=True) # Clear all environment variables
@patch("azure.identity.get_bearer_token_provider")
@patch("azure.identity.DefaultAzureCredential")
def test_get_azure_ad_token_provider_defaults_to_default_azure_credential(
self, mock_default_azure_credential, mock_get_bearer_token_provider
):
"""Test get_azure_ad_token_provider defaults to DefaultAzureCredential when no credentials are present."""
# Mock the Azure identity credential instance
mock_credential_instance = MagicMock()
mock_default_azure_credential.return_value = mock_credential_instance
# Mock the bearer token provider
mock_token_provider = MagicMock(return_value="mock-default-token")
mock_get_bearer_token_provider.return_value = mock_token_provider
# Call the function
result = get_azure_ad_token_provider()
# Assertions
assert callable(result)
mock_default_azure_credential.assert_called_once_with()
mock_get_bearer_token_provider.assert_called_once_with(
mock_credential_instance, "https://cognitiveservices.azure.com/.default"
)
# Test that the returned callable works
token = result()
assert token == "mock-default-token"