fix(azure): add keyless Decisions authentication (#48386) (#48529)

* fix(azure): support Entra auth for Decisions APIs (#48386)

* fix(azure): simplify Decisions provider registration (#48386)

* refactor(azure): reuse base Azure auth for Decisions (#48386)

litellm/__init__.py

@@ -1741,6 +1741,9 @@ if TYPE_CHECKING:
     from .llms.openai.decisions.transformation import (
         OpenAIDecisionsConfig as OpenAIDecisionsConfig,
     )
+    from .llms.azure.decisions.transformation import (
+        AzureDecisionsConfig as AzureDecisionsConfig,
+    )
     from .llms.azure_ai.decisions.transformation import (
         AzureAIDecisionsConfig as AzureAIDecisionsConfig,
     )

litellm/_lazy_imports_registry.py

@@ -174,6 +174,7 @@ LLM_CONFIG_NAMES: Final = (
     "LayaDecisionsConfig",
     "BespokeDecisionsConfig",
     "OpenAIDecisionsConfig",
+    "AzureDecisionsConfig",
     "AzureAIDecisionsConfig",
     "NvidiaNimRerankConfig",
     "NvidiaNimRankingConfig",
@@ -751,6 +752,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
     "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"),
     "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"),
     "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
+    "AzureDecisionsConfig": (".llms.azure.decisions.transformation", "AzureDecisionsConfig"),
     "AzureAIDecisionsConfig": (".llms.azure_ai.decisions.transformation", "AzureAIDecisionsConfig"),
     "NvidiaNimRerankConfig": (
         ".llms.nvidia_nim.rerank.transformation",

litellm/llms/azure/decisions/transformation.py

@@ -0,0 +1,77 @@
+import asyncio
+from collections.abc import Mapping
+from typing import Final
+
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.azure.common_utils import (
+    BaseAzureLLM,
+    cached_azure_ad_token_refresh_provider,
+)
+from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
+
+
+class AzureDecisionsConfig(OpenAIDecisionsConfig):
+    """Azure OpenAI Decisions, with the Azure deployment name in the request body."""
+
+    api_key_required = False
+
+    def resolve_api_base(self, api_base: str | None) -> str | None:
+        return BaseAzureLLM.resolve_azure_api_base(api_base)
+
+    def resolve_api_key(self, api_key: str | None) -> str | None:
+        return api_key
+
+    def resolve_credential(
+        self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object]
+    ) -> str | None:
+        return api_key
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        api_key: str | None,
+        litellm_params: Mapping[str, object],
+    ) -> Mapping[str, str]:
+        try:
+            params: Final = GenericLiteLLMParams.model_validate({**litellm_params, "api_key": api_key})
+        except ValidationError as error:
+            raise litellm.BadRequestError(
+                message="Invalid Azure Decisions credential parameters",
+                model=model,
+                llm_provider="azure",
+            ) from error
+        clean_headers: Final = super().validate_environment(
+            {name: value for name, value in headers.items() if name.lower() != "api-key"},
+            model,
+            None,
+            litellm_params,
+        )
+        azure_headers: Final = BaseAzureLLM.base_validate_azure_environment(
+            headers=dict(clean_headers), litellm_params=params
+        )
+        if "api-key" in azure_headers:
+            return azure_headers
+        if "Authorization" in azure_headers:
+            return azure_headers
+        scope: Final = (
+            params.azure_scope or get_secret_str("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default"
+        )
+        azure_headers["Authorization"] = f"Bearer {cached_azure_ad_token_refresh_provider(scope)()}"
+        return azure_headers
+
+    async def avalidate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        api_key: str | None,
+        litellm_params: Mapping[str, object],
+    ) -> Mapping[str, str]:
+        return await asyncio.to_thread(self.validate_environment, headers, model, api_key, litellm_params)
+
+    def get_complete_url(self, api_base: str, model: str) -> str:
+        return BaseAzureLLM.get_base_azure_url(api_base, None, route="/openai/v1/decisions")

litellm/utils.py

@@ -9616,8 +9616,12 @@ class ProviderConfigManager:
             return litellm.LayaDecisionsConfig()
         if provider == LlmProviders.BESPOKE:
             return litellm.BespokeDecisionsConfig()
-        if provider == LlmProviders.OPENAI:
-            return litellm.OpenAIDecisionsConfig()
+        if provider in (LlmProviders.OPENAI, LlmProviders.AZURE):
+            configs: Final = {
+                LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig,
+                LlmProviders.AZURE: litellm.AzureDecisionsConfig,
+            }
+            return configs[provider]()
         if provider == LlmProviders.AZURE_AI:
             return litellm.AzureAIDecisionsConfig()
         return None

tests/unit/decisions/test_main.py

@@ -1532,6 +1532,7 @@ async def test_databricks_router_deployment_sends_its_own_key_to_its_own_base(
 _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset(
     {
         LlmProviders.OPENAI.value,
+        LlmProviders.AZURE.value,
         LlmProviders.OPENROUTER.value,
         LlmProviders.HOSTED_VLLM.value,
         LlmProviders.LM_STUDIO.value,

tests/unit/llms/azure/decisions/test_transformation.py

@@ -0,0 +1,226 @@
+import json
+import threading
+import time
+from collections.abc import Iterator, Mapping
+from typing import Final
+from unittest.mock import Mock
+
+import pytest
+import respx
+from azure.core.credentials import AccessToken
+
+import litellm
+from litellm.llms.azure import common_utils
+from litellm.llms.azure.decisions.transformation import AzureDecisionsConfig
+from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
+
+_BASE: Final = "https://resource.openai.azure.com"
+_URL: Final = f"{_BASE}/openai/v1/decisions"
+_RESPONSE: Final = {
+    "model": "luna-deployment",
+    "answers": [{"type": "predicate", "name": "defect", "probability": 0.9}],
+    "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
+}
+
+
+@pytest.fixture(autouse=True)
+def clear_credentials(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    for name in (
+        "AZURE_API_KEY",
+        "AZURE_OPENAI_API_KEY",
+        "AZURE_API_BASE",
+        "AZURE_AD_TOKEN",
+        "AZURE_CLIENT_ID",
+        "AZURE_TENANT_ID",
+        "AZURE_CLIENT_SECRET",
+        "AZURE_SCOPE",
+        "AZURE_USERNAME",
+        "AZURE_PASSWORD",
+    ):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setattr(litellm, "api_key", None)
+    monkeypatch.setattr(litellm, "azure_key", None)
+    monkeypatch.setattr(litellm, "api_base", None)
+    monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False)
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    common_utils._cached_azure_ad_token_refresh_provider.cache_clear()
+    yield
+    common_utils._cached_azure_ad_token_refresh_provider.cache_clear()
+
+
+async def call_decisions(
+    is_async: bool,
+    systemone: bool,
+    credentials: Mapping[str, object],
+    api_base: str = _BASE,
+) -> DecisionsResponse | OpenAIDecisionResponse:
+    params = {
+        "model": "azure/luna-deployment",
+        "api_base": api_base,
+        **credentials,
+    }
+    if systemone:
+        params["state"] = "The export job hangs."
+        params["questions"] = {"defect": {"type": "noul", "instructions": "Is this a defect?"}}
+        if is_async:
+            return await litellm.asystemone(**params)
+        return litellm.systemone(**params)
+    params["input"] = "The export job hangs."
+    params["questions"] = [{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}]
+    if is_async:
+        return await litellm.adecisions(**params)
+    return litellm.decisions(**params)
+
+
+@pytest.mark.parametrize("is_async", (False, True))
+@pytest.mark.parametrize("systemone", (False, True))
+@pytest.mark.parametrize("suffix", ("", "/", "/openai", "/openai/v1/"))
+async def test_azure_decisions_uses_entra_token_and_openai_wire_format(
+    respx_mock: respx.MockRouter, is_async: bool, systemone: bool, suffix: str
+) -> None:
+    route = respx_mock.post(_URL).respond(json=_RESPONSE)
+    response = await call_decisions(is_async, systemone, {"azure_ad_token": "synthetic-entra"}, _BASE + suffix)
+    request = route.calls.last.request
+    assert request.headers["authorization"] == "Bearer synthetic-entra"
+    assert "api-key" not in request.headers
+    assert json.loads(request.content) == {
+        "model": "luna-deployment",
+        "input": "The export job hangs.",
+        "questions": [{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}],
+    }
+    if systemone:
+        assert response.model_dump()["answers"]["defect"]["noul"] == 0.9
+    else:
+        assert response.model_dump()["answers"][0]["probability"] == 0.9
+    assert response._hidden_params["custom_llm_provider"] == "azure"
+
+
+@pytest.mark.parametrize("is_async", (False, True))
+@pytest.mark.parametrize("systemone", (False, True))
+async def test_default_credential_refreshes_without_opt_in(
+    monkeypatch: pytest.MonkeyPatch,
+    respx_mock: respx.MockRouter,
+    is_async: bool,
+    systemone: bool,
+) -> None:
+    credential = Mock(spec=["get_token"])
+    credential.get_token.side_effect = [
+        AccessToken("synthetic-first", int(time.time()) + 120),
+        AccessToken("synthetic-refreshed", int(time.time()) + 3600),
+    ]
+    factory = Mock(return_value=credential)
+    monkeypatch.setattr("azure.identity.DefaultAzureCredential", factory)
+    monkeypatch.setenv("OPENAI_API_KEY", "synthetic-unrelated-openai-key")
+    route = respx_mock.post(_URL).respond(json=_RESPONSE)
+    await call_decisions(is_async, systemone, {})
+    await call_decisions(is_async, systemone, {})
+    assert [call.request.headers["authorization"] for call in route.calls] == [
+        "Bearer synthetic-first",
+        "Bearer synthetic-refreshed",
+    ]
+    factory.assert_called_once_with()
+    assert credential.get_token.call_count == 2
+    assert all(
+        call.args == ("https://cognitiveservices.azure.com/.default",) for call in credential.get_token.call_args_list
+    )
+
+
+@pytest.mark.parametrize("kind", ("service-principal", "oidc", "provider"))
+async def test_azure_decisions_reuses_existing_azure_credential_resolution(
+    monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, kind: str
+) -> None:
+    provider = Mock(return_value="synthetic-entra")
+    if kind == "service-principal":
+        factory = Mock(return_value=provider)
+        monkeypatch.setattr(common_utils, "get_azure_ad_token_from_entra_id", factory)
+        credentials = {
+            "tenant_id": "tenant",
+            "client_id": "client",
+            "client_secret": "synthetic-secret",
+        }
+    elif kind == "oidc":
+        monkeypatch.setattr(common_utils, "get_azure_ad_token_from_oidc", provider)
+        credentials = {
+            "tenant_id": "tenant",
+            "client_id": "client",
+            "azure_ad_token": "oidc/azure/https://cognitiveservices.azure.com/.default",
+        }
+    else:
+        credentials = {"azure_ad_token_provider": provider}
+    route = respx_mock.post(_URL).respond(json=_RESPONSE)
+    await call_decisions(True, False, credentials)
+    assert route.calls.last.request.headers["authorization"] == "Bearer synthetic-entra"
+    provider.assert_called_once()
+    if kind == "service-principal":
+        factory.assert_called_once_with(
+            tenant_id="tenant",
+            client_id="client",
+            client_secret="synthetic-secret",
+            scope="https://cognitiveservices.azure.com/.default",
+        )
+    if kind == "oidc":
+        provider.assert_called_once_with(
+            azure_ad_token=credentials["azure_ad_token"],
+            azure_client_id="client",
+            azure_tenant_id="tenant",
+            scope="https://cognitiveservices.azure.com/.default",
+        )
+
+
+@pytest.mark.parametrize("key_source", ("api_key", "AZURE_API_KEY", "AZURE_OPENAI_API_KEY"))
+async def test_azure_api_key_takes_precedence_over_entra(
+    monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, key_source: str
+) -> None:
+    provider = Mock(side_effect=AssertionError("API key must bypass Entra"))
+    credentials = {"azure_ad_token_provider": provider}
+    if key_source == "api_key":
+        credentials["api_key"] = "synthetic-key"
+    else:
+        monkeypatch.setenv(key_source, "synthetic-key")
+    route = respx_mock.post(_URL).respond(json=_RESPONSE)
+    await call_decisions(True, False, credentials)
+    assert route.calls.last.request.headers["api-key"] == "synthetic-key"
+    assert "authorization" not in route.calls.last.request.headers
+    provider.assert_not_called()
+
+
+async def test_azure_credential_lookup_does_not_block_the_event_loop(
+    respx_mock: respx.MockRouter,
+) -> None:
+    event_loop_thread = threading.get_ident()
+
+    def provider() -> str:
+        assert threading.get_ident() != event_loop_thread
+        return "synthetic-entra"
+
+    respx_mock.post(_URL).respond(json=_RESPONSE)
+    await call_decisions(True, False, {"azure_ad_token_provider": provider})
+
+
+def test_azure_decisions_requires_azure_endpoint(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("OPENAI_API_BASE", "https://unrelated.example")
+    with pytest.raises(litellm.BadRequestError, match="api_base is required"):
+        litellm.decisions(
+            model="azure/luna-deployment",
+            input="Text",
+            questions=[
+                {
+                    "type": "predicate",
+                    "name": "defect",
+                    "instructions": "Is this a defect?",
+                }
+            ],
+        )
+
+
+def test_azure_decisions_does_not_echo_invalid_credentials() -> None:
+    with pytest.raises(litellm.BadRequestError) as error:
+        AzureDecisionsConfig().validate_environment(
+            {},
+            "luna-deployment",
+            None,
+            {"azure_ad_token": {"secret": "sensitive-value"}},
+        )
+    assert "sensitive-value" not in str(error.value)