From e5757934fd7e1e37ecf566cd712ea66313315ccf Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Mon, 26 Oct 2026 20:09:58 -0700 Subject: [PATCH] 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 | 3 + litellm/_lazy_imports_registry.py | 2 + .../llms/azure/decisions/transformation.py | 77 ++++++ litellm/utils.py | 8 +- tests/unit/decisions/test_main.py | 1 + tests/unit/llms/azure/decisions/__init__.py | 0 .../azure/decisions/test_transformation.py | 226 ++++++++++++++++++ 7 files changed, 315 insertions(+), 2 deletions(-) create mode 100644 litellm/llms/azure/decisions/transformation.py create mode 100644 tests/unit/llms/azure/decisions/__init__.py create mode 100644 tests/unit/llms/azure/decisions/test_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index c2c13a12eea..9ec11dd1291 100644 --- a/litellm/__init__.py +++ b/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, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 6aecfd2ff8c..0f8288941e2 100644 --- a/litellm/_lazy_imports_registry.py +++ b/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", diff --git a/litellm/llms/azure/decisions/transformation.py b/litellm/llms/azure/decisions/transformation.py new file mode 100644 index 00000000000..3fb40e39702 --- /dev/null +++ b/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") diff --git a/litellm/utils.py b/litellm/utils.py index fdfad7fc665..781aa5f4fce 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index d13d212b4de..0bc18080ef4 100644 --- a/tests/unit/decisions/test_main.py +++ b/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, diff --git a/tests/unit/llms/azure/decisions/__init__.py b/tests/unit/llms/azure/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/decisions/test_transformation.py b/tests/unit/llms/azure/decisions/test_transformation.py new file mode 100644 index 00000000000..b6de6084b9c --- /dev/null +++ b/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)