scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(azure): keep the streaming Entra rejection inside the failed attempt

Browse source
This commit is contained in:
mateo-berri 2026-11-04 08:38:40 -08:00
parent d005db4662
commit d1d44fa0e8
2 changed files with 62 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/azure/common_utils.py
View file

@ -197,16 +197,20 @@ def _exception_chain(error: BaseException) -> Iterator[BaseException]:
link = link.__cause__ or link.__context__
_REQUEST_FAILURE_TYPES: Final = (BaseLLMException, APIError, httpx.HTTPError)
_REQUEST_FAILURE_TYPES: Final = (APIError, httpx.HTTPError)
def _belongs_to_the_same_attempt(link: BaseException) -> bool:
return isinstance(link, EntraRejectionError) or not isinstance(link, _REQUEST_FAILURE_TYPES)
def _ends_the_failed_attempt(link: BaseException, failed_attempt: BaseException) -> bool:
return not isinstance(link, EntraRejectionError) and (
isinstance(link, _REQUEST_FAILURE_TYPES) or type(link) is type(failed_attempt)
)
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))
yield from itertools.takewhile(
lambda link: not _ends_the_failed_attempt(link, error), itertools.islice(_exception_chain(error), 1, None)
)
def entra_rejection_in(error: BaseException) -> EntraRejectionError | None:

54
tests/unit/test_router/test_router.py
View file

@ -25848,6 +25848,60 @@ async def test_router_fallback_auth_failure_after_an_entra_rejection_still_cools
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")
@pytest.mark.parametrize(
"stream_request",
[
pytest.param(
lambda router, **caller_fields: router.acompletion(
model="entra-model", messages=[{"role": "user", "content": "hi"}], stream=True, **caller_fields
),
id="chat_completions",
),
pytest.param(
lambda router, **caller_fields: router.aanthropic_messages(
model="entra-model",
messages=[{"role": "user", "content": "hi"}],
max_tokens=5,
stream=True,
**caller_fields,
),
id="messages",
),
],
)
async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deployment_routable(
monkeypatch: pytest.MonkeyPatch,
stream_request: Callable[..., Awaitable[object]],
) -> None:
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: Final = 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"},
}
]
)
with pytest.raises(litellm.AuthenticationError):
await stream_request(router, client_secret="not-the-credential")
with pytest.raises(litellm.AuthenticationError):
await stream_request(router)
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_router_get_async_openai_model_client():
router = Router(