feat(azure-ai): support Entra ID auth for Microsoft-Decision-1 on /v1/decisions and /v1/systemone (#46539)

Backport of #46539 to lts/1.108.x.
Cherry-picked from caad4d6204 (main).
This commit is contained in:
jesus-berri 2026-11-12 21:14:05 +00:00 • committed by Devin AI
parent 1b769ce7d8
commit e18b03d84c
6 changed files with 490 additions and 10 deletions

View file

@ -1,11 +1,21 @@
import asyncio
import re
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import ValidationError
import litellm
from litellm.llms.azure_ai.common_utils import (
AZURE_ENTRA_LITELLM_PARAM_KEYS,
get_azure_ai_auth_headers,
)
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
_FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$")
_AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"})
class AzureAIDecisionsConfig(BaseDecisionsConfig):
@ -14,10 +24,69 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig):
path = "/providers/microsoft/v1/systemone"
api_key_env = ("AZURE_AI_API_KEY",)
api_base_env = ("AZURE_AI_API_BASE",)
api_key_required = False
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://<resource>.services.ai.azure.com"
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
credential_params: Final = MappingProxyType(
{key: litellm_params[key] for key in _AZURE_AI_CREDENTIAL_PARAM_KEYS if key in litellm_params}
)
try:
auth_headers: Final = get_azure_ai_auth_headers(api_key, credential_params)
except ValidationError as error:
field_names: Final = tuple(
dict.fromkeys(
".".join(str(field) for field in detail["loc"])
for detail in error.errors(include_input=False, include_context=False, include_url=False)
)
)
raise litellm.BadRequestError(
message=f"Invalid Azure AI credential parameter(s): {', '.join(f'`{name}`' for name in field_names)}",
model=model,
llm_provider="azure_ai",
) from error
except ValueError as error:
raise litellm.AuthenticationError(message=str(error), model=model, llm_provider="azure_ai") from error
return {
**super().validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
),
**auth_headers,
}
async def avalidate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
if api_key:
return self.validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
return await asyncio.to_thread(
self.validate_environment,
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
def get_complete_url(self, api_base: str, model: str) -> str:
url: Final = httpx.URL(api_base)
resource_path: Final = _FOUNDRY_ROUTE_SUFFIX.sub("", url.path.rstrip("/"))

View file

@ -510,13 +510,33 @@ class BaseDecisionsConfig(ABC):
def request_model(self, model: str) -> str:
return model
def validate_environment(self, headers: Mapping[str, str], model: str, api_key: str | None) -> dict[str, str]:
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
return {
**{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS},
**({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}),
"Content-Type": "application/json",
}
async def avalidate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
return self.validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
def get_complete_url(self, api_base: str, model: str) -> str:
return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}"

View file

@ -1527,8 +1527,65 @@ class BaseLLMHTTPHandler:
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
) -> tuple[str, dict[str, str], dict[str, object]]:
outbound_headers: Final = provider_config.validate_environment(headers=headers, model=model, api_key=api_key)
litellm_params: Mapping[str, object],
) -> tuple[str, Mapping[str, str], Mapping[str, object]]:
outbound_headers: Final = provider_config.validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
return self._prepare_decisions_request_with_headers(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
logged_body=logged_body,
api_base=api_base,
api_key=api_key,
outbound_headers=outbound_headers,
)
async def _aprepare_decisions_request(
self,
model: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseDecisionsConfig,
body: Mapping[str, object],
logged_body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
litellm_params: Mapping[str, object],
) -> tuple[str, Mapping[str, str], Mapping[str, object]]:
outbound_headers: Final = await provider_config.avalidate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
return self._prepare_decisions_request_with_headers(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
logged_body=logged_body,
api_base=api_base,
api_key=api_key,
outbound_headers=outbound_headers,
)
def _prepare_decisions_request_with_headers(
self,
model: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseDecisionsConfig,
body: Mapping[str, object],
logged_body: Mapping[str, object],
api_base: str,
api_key: str | None,
outbound_headers: Mapping[str, str],
) -> tuple[str, Mapping[str, str], Mapping[str, object]]:
url: Final = provider_config.get_complete_url(api_base=api_base, model=model)
data: Final = dict(body)
if logging_obj is not None:
@ -1565,6 +1622,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
signed_headers, signed_body = provider_config.sign_request(
headers=outbound_headers,
@ -1579,7 +1637,7 @@ class BaseLLMHTTPHandler:
response: Final = sync_httpx_client.post(
url,
content=signed_body,
json=data if signed_body is None else None,
json=dict(data) if signed_body is None else None,
headers=dict(signed_headers),
timeout=timeout,
logging_obj=logging_obj,
@ -1606,7 +1664,7 @@ class BaseLLMHTTPHandler:
litellm_params: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> DecisionsIRResponse:
url, outbound_headers, data = self._prepare_decisions_request(
url, outbound_headers, data = await self._aprepare_decisions_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
@ -1615,6 +1673,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
sign: Final = partial(
provider_config.sign_request,
@ -1637,7 +1696,7 @@ class BaseLLMHTTPHandler:
response: Final = await async_httpx_client.post(
url,
content=signed_body,
json=data if signed_body is None else None,
json=dict(data) if signed_body is None else None,
headers=dict(signed_headers),
timeout=timeout,
logging_obj=logging_obj,

View file

@ -11,9 +11,9 @@ import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.ports import allocate_port
from integration._support.process import owned_proxy_process
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration._support.ports import allocate_port
from integration.cost_calculation.cost_tracking_case import JsonResponse
from pydantic import JsonValue, TypeAdapter
@ -21,6 +21,13 @@ import litellm
_API_KEY: Final = "synthetic-decisions-key"
_ENV_KEY: Final = "synthetic-decisions-env-key"
_AZURE_AI_ENTRA_SCENARIO: Final = "providers"
_AZURE_AI_ENTRA_TOKEN: Final = "synthetic-decisions-entra-token"
_AZURE_AI_MISSING_CREDENTIALS_MESSAGE: Final = (
"Missing Azure AI credentials - set an API key (`api_key` or AZURE_AI_API_KEY), or Entra ID / OAuth "
"credentials (`tenant_id` + `client_id` + `client_secret`, `azure_ad_token`, an OIDC token, or a managed "
"identity with `litellm.enable_azure_ad_token_refresh = True`)"
)
_PASS_THROUGH_MODEL: Final = "gpt-6-luna"
_PASS_THROUGH_AUTHORIZATION: Final = "Bearer customer-held-upstream-key"
_PASS_THROUGH_NEIGHBOUR: Final = "decisions-beside-a-pass-through"
@ -100,9 +107,7 @@ _PROVIDERS: Final = (
_Provider(
"bespoke", "bespoke/nimble-latest", "/v1/systemone", "nimble-latest", None, False, "bespoke/nimble-latest"
),
_Provider(
"hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None
),
_Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None),
_Provider(
"fireworks_ai",
"fireworks_ai/accounts/fireworks/routers/rune",
@ -235,6 +240,56 @@ def _assert_connection_error(response: httpx.Response) -> None:
assert _GENERIC_API_ERROR not in response.text, response.text
def test_azure_ai_entra_auth_reaches_the_scripted_upstream_for_both_decisions_routes(
gateway: Gateway, scenario: Scenario
) -> None:
provider: Final = next(provider for provider in _PROVIDERS if provider.name == "azure_ai")
handle: Final = register_scenario(
_AZURE_AI_ENTRA_SCENARIO,
JsonResponse(content_type="application/json", body=_answer_body(provider)),
)
scenario.cleanups.callback(delete_scenario, handle)
model: Final = "azure_ai/decision-1-entra"
body: Final = {"model": model, "state": _STATE, "questions": _QUESTIONS}
decisions_body: Final = {
"model": model,
"input": "The export job hangs at 99%",
"questions": [{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"}],
}
no_credentials_body: Final = {**body, "model": "azure_ai/decision-1-no-credentials"}
systemone_auth_error: Final = gateway.request("POST", "/v1/systemone", no_credentials_body)
assert systemone_auth_error.status_code == 401, systemone_auth_error.text
assert _AZURE_AI_MISSING_CREDENTIALS_MESSAGE in string_value(
object_value(object_value(systemone_auth_error.json())["error"])["message"]
)
no_credentials_decisions_body: Final = {
**decisions_body,
"model": "azure_ai/decision-1-no-credentials",
}
decisions_auth_error: Final = gateway.request("POST", "/v1/decisions", no_credentials_decisions_body)
assert decisions_auth_error.status_code == 401, decisions_auth_error.text
assert _AZURE_AI_MISSING_CREDENTIALS_MESSAGE in string_value(
object_value(object_value(decisions_auth_error.json())["error"])["message"]
)
assert not _upstream_calls(gateway, handle)
systemone_response: Final = gateway.request("POST", "/v1/systemone", body)
assert systemone_response.status_code == 200, systemone_response.text
decisions_response: Final = gateway.request("POST", "/v1/decisions", decisions_body)
assert decisions_response.status_code == 200, decisions_response.text
calls: Final = _upstream_calls(gateway, handle)
assert tuple(string_value(call["authorization"]) for call in calls) == (
f"Bearer {_AZURE_AI_ENTRA_TOKEN}",
f"Bearer {_AZURE_AI_ENTRA_TOKEN}",
)
assert tuple(string_value(call["path"]) for call in calls) == (
"/providers/microsoft/v1/systemone",
"/providers/microsoft/v1/systemone",
)
def _pass_through_config(directory: Path, pass_through_target: str, native_api_base: str) -> Path:
base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {

View file

@ -625,6 +625,19 @@ model_list:
api_key: synthetic-azure-ai-key
model_info:
base_model: azure_ai/Microsoft-Decision-1
- model_name: azure_ai/decision-1-entra
litellm_params:
model: azure_ai/decision-1
api_base: os.environ/INTEGRATION_UPSTREAM_URL
azure_ad_token: synthetic-decisions-entra-token
model_info:
base_model: azure_ai/Microsoft-Decision-1
- model_name: azure_ai/decision-1-no-credentials
litellm_params:
model: azure_ai/decision-1
api_base: os.environ/INTEGRATION_UPSTREAM_URL
model_info:
base_model: azure_ai/Microsoft-Decision-1
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL

View file

@ -24,6 +24,7 @@ from litellm.constants import AWS_SIGNING_MAX_THREADS
from litellm.cost_calculator import get_response_cost_from_hidden_params
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers
from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.decisions import (
@ -2681,6 +2682,269 @@ async def test_azure_ai_decision_requires_api_base_before_http(
assert len(respx_mock.calls) == 0
def _clear_azure_ai_auth_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("AZURE_AI_API_KEY", raising=False)
monkeypatch.delenv("AZURE_AD_TOKEN", raising=False)
monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
monkeypatch.delenv("AZURE_CLIENT_ID", raising=False)
monkeypatch.delenv("AZURE_CLIENT_SECRET", raising=False)
monkeypatch.delenv("AZURE_USERNAME", raising=False)
monkeypatch.delenv("AZURE_PASSWORD", raising=False)
monkeypatch.delenv("AZURE_AI_API_BASE", raising=False)
monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False)
def test_azure_ai_systemone_uses_entra_token_without_api_key(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.systemone(
model="azure_ai/decision-1",
state="review",
questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
api_base="https://res.services.ai.azure.com",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token"
@pytest.mark.asyncio
async def test_azure_ai_decisions_use_entra_token_without_api_key(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
await litellm.adecisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token"
def test_azure_ai_systemone_environment_api_key_wins_over_deployment_entra(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
monkeypatch.setenv("AZURE_AI_API_KEY", "env-key")
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.systemone(
model="azure_ai/decision-1",
state="review",
questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
api_base="https://res.services.ai.azure.com",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer env-key"
@pytest.mark.asyncio
async def test_azure_ai_decisions_environment_api_key_wins_over_deployment_entra(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
monkeypatch.setenv("AZURE_AI_API_KEY", "env-key")
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
await litellm.adecisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer env-key"
def test_azure_ai_decisions_uses_environment_api_key_without_entra_credentials(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
monkeypatch.setenv("AZURE_AI_API_KEY", "env-key")
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer env-key"
def test_azure_ai_decisions_explicit_api_key_precedes_entra_and_environment(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
monkeypatch.setenv("AZURE_AI_API_KEY", "env-key")
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
api_key="call-key",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer call-key"
def test_azure_ai_decisions_does_not_validate_noncredential_params_as_auth(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
azure_ad_token="entra-access-token",
max_retries="three",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer entra-access-token"
def test_azure_ai_decisions_malformed_credential_is_a_sanitized_bad_request(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
with pytest.raises(litellm.BadRequestError) as caught:
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
tenant_id="tenant-id",
client_id="client-id",
client_secret=123,
)
assert caught.value.status_code == 400
assert "`client_secret`" in str(caught.value)
assert "123" not in str(caught.value)
assert not route.called
def test_azure_ai_decisions_prefer_api_key_when_both_are_set(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
api_key="foundry-key",
azure_ad_token="entra-access-token",
)
assert route.called
assert route.calls[0].request.headers["authorization"] == "Bearer foundry-key"
def test_azure_ai_decisions_without_credentials_raise_authentication_error(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
with pytest.raises(ValueError, match="Missing Azure AI credentials") as helper_error:
get_azure_ai_auth_headers(None, {})
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
with pytest.raises(litellm.AuthenticationError) as caught:
litellm.decisions(
model="azure_ai/decision-1",
input="review",
questions=_PREDICATE_QUESTIONS,
api_base="https://res.services.ai.azure.com",
)
assert caught.value.status_code == 401
assert str(helper_error.value) in str(caught.value)
assert not route.called
@pytest.mark.asyncio
async def test_azure_ai_systemone_without_credentials_raises_authentication_error(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
_clear_azure_ai_auth_environment(monkeypatch)
with pytest.raises(ValueError, match="Missing Azure AI credentials") as helper_error:
get_azure_ai_auth_headers(None, {})
route: Final = respx_mock.post("https://res.services.ai.azure.com/providers/microsoft/v1/systemone").respond(
json={**_RESPONSE, "model": "microsoft-decision-1"}
)
with pytest.raises(litellm.AuthenticationError) as caught:
await litellm.asystemone(
model="azure_ai/decision-1",
state="review",
questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
api_base="https://res.services.ai.azure.com",
)
assert caught.value.status_code == 401
assert str(helper_error.value) in str(caught.value)
assert not route.called
_CLEF_OMNI_UPSTREAM: Final = "https://example.com/ai/run/@cf/cloudflare/clef-omni"
_CLEF_OMNI_RESPONSE: Final[Mapping[str, object]] = {
"success": True,