diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 19f51c92aad..a77bb6f2eed 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,5 +1,6 @@ import asyncio import hashlib +import itertools import json import os import re @@ -9,7 +10,7 @@ from types import MappingProxyType from typing import Final, Literal, NamedTuple, NoReturn, cast import httpx -from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from openai import APIError, AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict @@ -57,6 +58,10 @@ _AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitab ) +def _callable_azure_ad_token_provider(provider: object) -> Callable[[], str | Awaitable[str]] | None: + return _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(provider) if callable(provider) else None + + class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): api_version: ReadOnly[str] base_url: ReadOnly[str] @@ -192,8 +197,22 @@ def _exception_chain(error: BaseException) -> Iterator[BaseException]: link = link.__cause__ or link.__context__ +_REQUEST_FAILURE_TYPES: Final = (BaseLLMException, APIError, httpx.HTTPError) + + +def _belongs_to_the_same_attempt(link: BaseException) -> bool: + return isinstance(link, EntraRejectionError) or not isinstance(link, _REQUEST_FAILURE_TYPES) + + +def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]: + yield error + yield from itertools.takewhile(_belongs_to_the_same_attempt, itertools.islice(_exception_chain(error), 1, None)) + + def entra_rejection_in(error: BaseException) -> EntraRejectionError | None: - return next((link for link in _exception_chain(error) if isinstance(link, EntraRejectionError)), None) + return next( + (link for link in _links_within_the_failed_attempt(error) if isinstance(link, EntraRejectionError)), None + ) def caused_by_entra_rejection(error: BaseException) -> bool: @@ -230,7 +249,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"] if "x-ratelimit-remaining-tokens" in headers: openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"] - llm_response_headers: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()} + llm_response_headers: Final = {f"llm_provider-{k}": v for k, v in headers.items()} return {**llm_response_headers, **openai_headers} @@ -850,9 +869,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_version: str | None, is_async: bool, ) -> dict: - azure_ad_token_provider = _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python( - litellm_params.get("azure_ad_token_provider") - ) + azure_ad_token_provider = _callable_azure_ad_token_provider(litellm_params.get("azure_ad_token_provider")) # If we have api_key, then we have higher priority azure_ad_token = litellm_params.get("azure_ad_token") diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index e4698942c7d..d66ba3f324e 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -2971,3 +2971,23 @@ def test_token_provider_error_without_azure_core_installed_keeps_the_generic_tok get_azure_ad_token( GenericLiteLLMParams(azure_ad_token_provider=_raising_token_provider(ValueError("token service down"))) ) + + +def test_image_generation_with_an_api_key_ignores_a_non_callable_azure_ad_token_provider(): + import respx + + with respx.mock: + image_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]}) + ) + response: Final = litellm.image_generation( + model="azure/gpt-image-1", + prompt="a red dot", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + api_key="configured-key", + azure_ad_token_provider="not-a-callable", + ) + + assert image_route.calls.last.request.headers["api-key"] == "configured-key" + assert response.data[0].b64_json == "aGk=" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 8be841e6224..11b07472eaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7360,6 +7360,60 @@ class TestRouterModelRelayUpstreamContract: assert not deployment_route.called + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_later_deployment_failure_after_an_entra_rejection( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from azure.core.exceptions import ClientAuthenticationError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + "order": 1, + }, + }, + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa-unreachable.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "configured-key", + "order": 2, + }, + }, + ], + num_retries=1, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + unreachable_route: Final = respx.post(url__startswith="https://litellm-qa-unreachable.openai.azure.com").mock( + side_effect=httpx.ConnectError("Connection refused") + ) + with pytest.raises(BaseLLMException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert unreachable_route.called + assert excinfo.value.status_code == 500 + + NIM_INFER_BODY = { "input": [ {"type": "image_url", "url": "data:image/png;base64,AAAA"}, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 2655da87485..aa2efbfd04f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -25791,6 +25791,63 @@ async def test_router_auth_failure_entra_did_not_cause_cools_down_despite_a_call assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_fallback_auth_failure_after_an_entra_rejection_still_cools_down_the_fallback(monkeypatch): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + }, + "model_info": {"id": "entra-deployment"}, + }, + { + "model_name": "key-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa-fallback.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "model_info": {"id": "key-deployment"}, + }, + ], + fallbacks=[{"entra-model": ["key-model"]}], + ) + + with respx.mock: + fallback_route: Final = respx.post(url__startswith="https://litellm-qa-fallback.openai.azure.com").mock( + return_value=httpx.Response( + 401, json={"error": {"code": "401", "message": "Access denied due to invalid subscription key."}} + ) + ) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + client_secret="not-the-credential", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="key-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert fallback_route.call_count == 1 + assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router(