mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(azure): send the resolved Entra ID token on image generation requests
Azure image generation calls initialize_azure_sdk_client like the chat path does, but then sends the request through httpx with the headers it was given, so a credential resolved from litellm_params or the environment (Entra ID client credentials, managed or workload identity, OIDC, a static azure_ad_token) never reached the wire and Azure answered 401. Only an explicitly passed azure_ad_token_provider was applied Add get_azure_request_auth_headers, which turns the credential in azure_client_params into an Authorization: Bearer header (or api-key, following the SDK's precedence) while keeping any auth header the caller already set, and use it for both the sync and async image requests. The pre-call logging metadata receives a redacted copy of those headers so the token never reaches logging callbacks Fixes #16422
This commit is contained in:
abhirup7
parent
4b3355bdc6
commit
91761d984d
3 changed files with 301 additions and 14 deletions
|
|
@ -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 (
|
||||
|
|
@ -1142,7 +1144,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,
|
||||
|
|
@ -1167,7 +1169,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(
|
||||
|
|
@ -1226,7 +1228,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,
|
||||
|
|
@ -1261,21 +1263,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,
|
||||
|
|
@ -1286,7 +1289,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
client=client,
|
||||
azure_client_params=azure_client_params,
|
||||
timeout=timeout,
|
||||
headers=headers,
|
||||
headers=request_headers,
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
|
@ -1303,7 +1306,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(
|
||||
|
|
@ -1313,7 +1316,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"))
|
||||
|
|
|
|||
|
|
@ -405,6 +405,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(
|
||||
|
|
|
|||
|
|
@ -10,6 +10,11 @@ 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_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 +592,247 @@ 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.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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue