scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge 66ebc722d6 into 1c61c2606e

Browse source
This commit is contained in:
Abhirup Sahoo 2026-09-12 09:47:45 -07:00 • committed by GitHub
commit 9a7d935784
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 357 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/llms/azure/azure.py
View file

@ -46,7 +46,9 @@ from .common_utils import (
AzureOpenAIError,
BaseAzureLLM,
get_azure_ad_token_from_oidc,
get_azure_request_auth_headers,
process_azure_headers,
redact_azure_auth_headers,
select_azure_base_url_or_endpoint,
)
from .image_generation import (
@ -1144,7 +1146,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_key: str,
input: list,
logging_obj: LiteLLMLoggingObj,
headers: dict,
headers: dict[str, str],
client=None,
timeout=None,
model: str | None = None,
@ -1169,7 +1171,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
additional_args={
"complete_input_dict": data,
"api_base": img_gen_api_base,
"headers": headers,
"headers": redact_azure_auth_headers(headers),
},
)
httpx_response: Final[httpx.Response] = await self.make_async_azure_httpx_request(
@ -1228,7 +1230,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
timeout: float,
optional_params: dict,
logging_obj: LiteLLMLoggingObj,
headers: dict,
headers: dict[str, str],
model: str | None = None,
api_key: str | None = None,
api_base: str | None = None,
@ -1263,21 +1265,22 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if not isinstance(max_retries, int):
raise AzureOpenAIError(status_code=422, message="max retries must be an int")
if api_key is None and azure_ad_token_provider is not None:
azure_ad_token = azure_ad_token_provider()
if azure_ad_token:
headers.pop("api-key", None)
headers["Authorization"] = f"Bearer {azure_ad_token}"
# init AzureOpenAI Client
auth_params: Final[dict[str, object]] = {**(litellm_params or {})} # mutable-ok: SDK init takes a dict
if azure_ad_token is not None:
auth_params["azure_ad_token"] = azure_ad_token
if azure_ad_token_provider is not None:
auth_params["azure_ad_token_provider"] = azure_ad_token_provider
azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
litellm_params=auth_params,
api_key=api_key,
model_name=model or "",
api_version=api_version,
api_base=api_base,
is_async=False,
)
request_headers: Final = dict( # mutable-ok: the httpx request helpers take a dict
get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params)
)
if aimg_generation is True:
return self.aimage_generation(
data=data,
@ -1288,7 +1291,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
client=client,
azure_client_params=azure_client_params,
timeout=timeout,
headers=headers,
headers=request_headers,
model=model,
)
@ -1305,7 +1308,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
additional_args={
"complete_input_dict": data,
"api_base": img_gen_api_base,
"headers": headers,
"headers": redact_azure_auth_headers(request_headers),
},
)
httpx_response: Final[httpx.Response] = self.make_sync_azure_httpx_request(
@ -1315,7 +1318,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_version=api_version or "",
api_key=api_key or "",
data=data,
headers=headers,
headers=request_headers,
deployment_name=model,
)
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))

44
litellm/llms/azure/common_utils.py
View file

@ -96,6 +96,11 @@ def _cached_entra_id_token_provider(
return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope)
@lru_cache(maxsize=128)
def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]:
return get_azure_ad_token_provider(azure_scope=scope)
def get_azure_ad_token_from_entra_id(
tenant_id: str,
client_id: str,
@ -406,6 +411,41 @@ def get_azure_ad_token(
return azure_ad_token
_AZURE_AUTH_HEADER_NAMES: Final = frozenset(("api-key", "authorization"))
_REDACTED_AZURE_HEADER_VALUE: Final = "***REDACTED***"
def _resolve_azure_ad_token(azure_client_params: Mapping[str, object]) -> str | None:
azure_ad_token: Final = azure_client_params.get("azure_ad_token")
if isinstance(azure_ad_token, str) and azure_ad_token:
return azure_ad_token
token_provider: Final = azure_client_params.get("azure_ad_token_provider")
provided_token: Final = token_provider() if callable(token_provider) else None
return provided_token if isinstance(provided_token, str) and provided_token else None
def get_azure_request_auth_headers(
headers: Mapping[str, str],
azure_client_params: Mapping[str, object],
) -> Mapping[str, str]:
if any(name.lower() in _AZURE_AUTH_HEADER_NAMES for name in headers):
return headers
azure_ad_token: Final = _resolve_azure_ad_token(azure_client_params)
if azure_ad_token is not None:
return MappingProxyType({**headers, "Authorization": f"Bearer {azure_ad_token}"})
api_key: Final = azure_client_params.get("api_key")
if isinstance(api_key, str) and api_key:
return MappingProxyType({**headers, "api-key": api_key})
return headers
def redact_azure_auth_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
return { # mutable-ok: logging callbacks JSON-serialize this copy
name: (_REDACTED_AZURE_HEADER_VALUE if name.lower() in _AZURE_AUTH_HEADER_NAMES else value)
for name, value in headers.items()
}
class BaseAzureLLM(BaseOpenAILLM):
@staticmethod
def _try_get_default_azure_credential_provider(
@ -616,9 +656,7 @@ 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 = _cached_azure_ad_token_refresh_provider(scope)
except ValueError:
verbose_logger.debug("Azure AD Token Provider could not be used.")
if api_version is None:

296
tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
View file

@ -10,6 +10,12 @@ import respx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.common_utils import (
_cached_azure_ad_token_refresh_provider,
_cached_entra_id_token_provider,
get_azure_request_auth_headers,
redact_azure_auth_headers,
)
from litellm.llms.azure.image_generation.http_utils import (
azure_deployment_image_generation_json_body,
)
@ -587,3 +593,293 @@ def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_moc
sent_body = json.loads(request.content)
assert sent_body["model"] == model
assert sent_body["prompt"] == prompt
@pytest.fixture
def fake_entra_id(monkeypatch: pytest.MonkeyPatch):
built_credentials = []
class FakeClientSecretCredential:
def __init__(self, tenant_id: str, client_id: str, client_secret: str) -> None:
built_credentials.append((tenant_id, client_id, client_secret))
monkeypatch.setattr("azure.identity.ClientSecretCredential", FakeClientSecretCredential)
monkeypatch.setattr("azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "entra-id-token")
_cached_entra_id_token_provider.cache_clear()
yield built_credentials
_cached_entra_id_token_provider.cache_clear()
def _mock_image_generation_route(respx_mock: respx.MockRouter, api_base: str, model: str) -> respx.Route:
return respx_mock.post(f"{api_base}/openai/deployments/{model}/images/generations").mock(
return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]})
)
@pytest.mark.parametrize("credentials_in_litellm_params", [False, True])
def test_azure_image_generation_keyless_entra_id_sends_bearer_token(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
fake_entra_id: list,
credentials_in_litellm_params: bool,
):
for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"):
monkeypatch.delenv(name, raising=False)
api_base = "https://my-resource.openai.azure.com"
api_version = "2025-04-01-preview"
litellm_params = {"api_base": api_base, "api_version": api_version}
if credentials_in_litellm_params:
litellm_params.update(
tenant_id="tenant-from-params", client_id="client-from-params", client_secret="secret-from-params"
)
expected_credential = ("tenant-from-params", "client-from-params", "secret-from-params")
else:
monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env")
monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env")
monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env")
expected_credential = ("tenant-from-env", "client-from-env", "secret-from-env")
route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
logging_obj = MagicMock()
response = AzureChatCompletion().image_generation(
prompt="a cat",
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=logging_obj,
headers={"Content-Type": "application/json"},
model="gpt-image-1",
api_key=None,
api_base=api_base,
api_version=api_version,
litellm_params=litellm_params,
)
request = route.calls.last.request
assert request.headers["Authorization"] == "Bearer entra-id-token"
assert "api-key" not in request.headers
assert fake_entra_id == [expected_credential]
assert response.data[0].b64_json == "aaaa"
logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]
assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"}
assert "entra-id-token" not in str(logging_obj.pre_call.call_args)
@pytest.mark.asyncio
async def test_azure_aimage_generation_keyless_entra_id_sends_bearer_token(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
api_base = "https://my-resource.openai.azure.com"
api_version = "2025-04-01-preview"
route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
logging_obj = MagicMock()
response = await AzureChatCompletion().image_generation(
prompt="a cat",
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=logging_obj,
headers={"Content-Type": "application/json"},
model="gpt-image-1",
api_key=None,
api_base=api_base,
api_version=api_version,
aimg_generation=True,
litellm_params={
"api_base": api_base,
"api_version": api_version,
"tenant_id": "tenant-from-params",
"client_id": "client-from-params",
"client_secret": "secret-from-params",
},
)
request = route.calls.last.request
assert request.headers["Authorization"] == "Bearer entra-id-token"
assert "api-key" not in request.headers
assert fake_entra_id == [("tenant-from-params", "client-from-params", "secret-from-params")]
assert response.data[0].b64_json == "aaaa"
logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]
assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"}
assert "entra-id-token" not in str(logging_obj.pre_call.call_args)
@pytest.mark.parametrize(
"credential_kwargs, expected_authorization",
[
({"azure_ad_token": "static-ad-token"}, "Bearer static-ad-token"),
({"azure_ad_token_provider": lambda: "provider-token"}, "Bearer provider-token"),
],
)
def test_azure_image_generation_explicit_azure_ad_credential_sends_bearer_token(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
credential_kwargs: dict,
expected_authorization: str,
):
for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"):
monkeypatch.delenv(name, raising=False)
api_base = "https://my-resource.openai.azure.com"
api_version = "2025-04-01-preview"
route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
response = AzureChatCompletion().image_generation(
prompt="a cat",
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=MagicMock(),
headers={"Content-Type": "application/json"},
model="gpt-image-1",
api_key=None,
api_base=api_base,
api_version=api_version,
litellm_params={"api_base": api_base, "api_version": api_version},
**credential_kwargs,
)
request = route.calls.last.request
assert request.headers["Authorization"] == expected_authorization
assert "api-key" not in request.headers
assert response.data[0].b64_json == "aaaa"
def test_azure_image_generation_with_api_key_keeps_api_key_header(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list
):
monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env")
monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env")
monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env")
api_base = "https://my-resource.openai.azure.com"
api_version = "2025-04-01-preview"
route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
logging_obj = MagicMock()
response = AzureChatCompletion().image_generation(
prompt="a cat",
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=logging_obj,
headers={"Content-Type": "application/json", "api-key": "sk-test"},
model="gpt-image-1",
api_key="sk-test",
api_base=api_base,
api_version=api_version,
litellm_params={"api_base": api_base, "api_version": api_version},
)
request = route.calls.last.request
assert request.headers["api-key"] == "sk-test"
assert "Authorization" not in request.headers
assert fake_entra_id == []
assert response.data[0].b64_json == "aaaa"
assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***"
@pytest.fixture
def fake_default_azure_credential(monkeypatch: pytest.MonkeyPatch):
built_credentials = []
class FakeDefaultAzureCredential:
def __init__(self) -> None:
built_credentials.append(self)
for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", "AZURE_AD_TOKEN"):
monkeypatch.delenv(name, raising=False)
monkeypatch.setattr("azure.identity.DefaultAzureCredential", FakeDefaultAzureCredential)
monkeypatch.setattr(
"azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "default-credential-token"
)
monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True)
_cached_azure_ad_token_refresh_provider.cache_clear()
yield built_credentials
_cached_azure_ad_token_refresh_provider.cache_clear()
def test_azure_image_generation_token_refresh_reuses_credential_across_requests(
respx_mock: respx.MockRouter, fake_default_azure_credential: list
):
api_base = "https://my-resource.openai.azure.com"
api_version = "2025-04-01-preview"
route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
for _ in range(3):
AzureChatCompletion().image_generation(
prompt="a cat",
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=MagicMock(),
headers={"Content-Type": "application/json"},
model="gpt-image-1",
api_key=None,
api_base=api_base,
api_version=api_version,
litellm_params={"api_base": api_base, "api_version": api_version},
)
assert route.call_count == 3
assert all(call.request.headers["Authorization"] == "Bearer default-credential-token" for call in route.calls)
assert len(fake_default_azure_credential) == 1
@pytest.mark.parametrize(
"caller_auth_header",
[{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}],
)
def test_get_azure_request_auth_headers_keeps_caller_auth_header(caller_auth_header: dict):
headers = {"Content-Type": "application/json", **caller_auth_header}
azure_client_params = {
"api_key": "sk-resolved",
"azure_ad_token": "resolved-token",
"azure_ad_token_provider": lambda: "provider-token",
}
assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers
def test_get_azure_request_auth_headers_prefers_azure_ad_token_over_provider_and_api_key():
headers = {"Content-Type": "application/json"}
azure_client_params = {
"api_key": "sk-resolved",
"azure_ad_token": "static-token",
"azure_ad_token_provider": lambda: "provider-token",
}
out = get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params)
assert dict(out) == {"Content-Type": "application/json", "Authorization": "Bearer static-token"}
assert headers == {"Content-Type": "application/json"}
def test_get_azure_request_auth_headers_uses_token_provider_over_api_key():
azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": lambda: "pt"}
out = get_azure_request_auth_headers(headers={}, azure_client_params=azure_client_params)
assert dict(out) == {"Authorization": "Bearer pt"}
def test_get_azure_request_auth_headers_falls_back_to_api_key():
azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": None}
out = get_azure_request_auth_headers(headers={"Content-Type": "application/json"}, azure_client_params=azure_client_params)
assert dict(out) == {"Content-Type": "application/json", "api-key": "sk-resolved"}
@pytest.mark.parametrize(
"azure_client_params",
[
{},
{"api_key": "", "azure_ad_token": "", "azure_ad_token_provider": None},
{"azure_ad_token_provider": lambda: None},
{"azure_ad_token_provider": lambda: ""},
],
)
def test_get_azure_request_auth_headers_without_credential_leaves_headers_unchanged(azure_client_params: dict):
headers = {"Content-Type": "application/json"}
assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers
def test_redact_azure_auth_headers_masks_only_credential_values():
headers = {"Content-Type": "application/json", "api-key": "sk-secret", "authorization": "Bearer secret"}
assert redact_azure_auth_headers(headers) == {
"Content-Type": "application/json",
"api-key": "***REDACTED***",
"authorization": "***REDACTED***",
}
assert headers["api-key"] == "sk-secret"
assert headers["authorization"] == "Bearer secret"

3
tests/test_litellm/llms/azure/test_azure_common_utils.py
View file

@ -9,6 +9,7 @@ import pytest
import litellm
from litellm.llms.azure.common_utils import (
BaseAzureLLM,
_cached_azure_ad_token_refresh_provider,
_cached_entra_id_token_provider,
get_azure_ad_token,
get_azure_ad_token_from_entra_id,
@ -34,6 +35,7 @@ def setup_mocks(monkeypatch):
monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
monkeypatch.delenv("AZURE_SCOPE", raising=False)
monkeypatch.delenv("AZURE_AD_TOKEN", raising=False)
_cached_azure_ad_token_refresh_provider.cache_clear()
with (
patch(
@ -78,6 +80,7 @@ def setup_mocks(monkeypatch):
"logger": mock_logger,
"select_url": mock_select_url,
}
_cached_azure_ad_token_refresh_provider.cache_clear()
def test_initialize_with_api_key(setup_mocks):