fix(azure): map context_length_exceeded error code to ContextWindowExceededError (#48616)

litellm/exceptions.py

@@ -232,7 +232,7 @@ class BadRequestError(openai.BadRequestError):
         litellm_debug_info: str | None = None,
         max_retries: int | None = None,
         num_retries: int | None = None,
-        body: dict | None = None,
+        body: Mapping[str, object] | None = None,
     ):
         self.status_code = 400
         self.message = f"litellm.BadRequestError: {message}"
@@ -276,7 +276,7 @@ class ImageFetchError(BadRequestError):
         litellm_debug_info: str | None = None,
         max_retries: int | None = None,
         num_retries: int | None = None,
-        body: dict | None = None,
+        body: Mapping[str, object] | None = None,
     ):
         super().__init__(
             message=message,
@@ -556,6 +556,7 @@ class ContextWindowExceededError(BadRequestError):
         llm_provider,
         response: httpx.Response | None = None,
         litellm_debug_info: str | None = None,
+        body: Mapping[str, object] | None = None,
     ):
         self.status_code = 400
         self.model = model
@@ -567,6 +568,7 @@ class ContextWindowExceededError(BadRequestError):
             llm_provider=self.llm_provider,
             response=response,
             litellm_debug_info=self.litellm_debug_info,
+            body=body,
         )  # Call the base class constructor with the parameters it needs
 
         # set after, to make it clear the raised error is a context window exceeded error
@@ -685,7 +687,7 @@ class ContentPolicyViolationError(BadRequestError):
         response: httpx.Response | None = None,
         litellm_debug_info: str | None = None,
         provider_specific_fields: dict | None = None,
-        body: dict | None = None,
+        body: Mapping[str, object] | None = None,
     ):
         self.status_code = 400
         self.message = f"litellm.ContentPolicyViolationError: {message}"

litellm/litellm_core_utils/exception_mapping_utils.py

@@ -2036,14 +2036,6 @@ def _map_azure_exception(
             litellm_debug_info=extra_information,
             response=getattr(original_exception, "response", None),
         )
-    elif "This model's maximum context length is" in error_str:
-        raise ContextWindowExceededError(
-            message=f"AzureException ContextWindowExceededError - {message}",
-            llm_provider="azure",
-            model=model,
-            litellm_debug_info=extra_information,
-            response=getattr(original_exception, "response", None),
-        )
     elif "DeploymentNotFound" in error_str:
         raise NotFoundError(
             message=f"AzureException NotFoundError - {message}",
@@ -2085,6 +2077,18 @@ def _map_azure_exception(
             response=getattr(original_exception, "response", None),
             body=getattr(original_exception, "body", None),
         )
+    elif getattr(original_exception, "status_code", None) in (None, 400) and (
+        azure_error_code == "context_length_exceeded"
+        or ExceptionCheckers.is_error_str_context_window_exceeded(error_str)
+    ):
+        raise ContextWindowExceededError(
+            message=f"AzureException ContextWindowExceededError - {message}",
+            llm_provider="azure",
+            model=model,
+            litellm_debug_info=extra_information,
+            response=getattr(original_exception, "response", None),
+            body=getattr(original_exception, "body", None),
+        )
     elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400):
         raise BadRequestError(
             message=f"AzureException BadRequestError - {message}",

tests/unit/llms/azure/test_azure_exception_mapping.py

@@ -1,13 +1,16 @@
 import json
+from collections.abc import Awaitable, Callable, Mapping
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import pytest
 import httpx
+import openai
 import respx
 
 
 import litellm
-from litellm.exceptions import ContentPolicyViolationError
+from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError
 from litellm.litellm_core_utils.exception_mapping_utils import exception_type
 
 
@@ -490,3 +493,142 @@ class TestAzureExceptionMapping:
         error = exc_info.value
         assert "encrypted_content_affinity" in error.message
         assert "enable_pre_call_checks" in error.message
+
+
+_AZURE_API_BASE: Final = "https://azure-context.openai.azure.com"
+_AZURE_CALL_KWARGS: Final = {
+    "api_base": _AZURE_API_BASE,
+    "api_version": "2024-10-21",
+    "api_key": "test-azure-key",
+    "max_retries": 0,
+}
+
+
+_CONTEXT_LENGTH_ERROR: Final = {
+    "message": "Your input exceeds the context window of this model. Please adjust your input and try again.",
+    "type": "invalid_request_error",
+    "param": "input",
+    "code": "context_length_exceeded",
+}
+
+
+def _azure_error_response(status_code: int, error: Mapping[str, object]) -> httpx.Response:
+    return httpx.Response(
+        status_code, text=json.dumps({"error": error}, indent=2), headers={"content-type": "application/json"}
+    )
+
+
+async def _azure_chat(*, stream: bool) -> None:
+    messages: Final = [{"role": "user", "content": "hi"}]
+    if stream:
+        stream_response: Final = await litellm.acompletion(
+            model="azure/gpt-4o", messages=messages, stream=True, **_AZURE_CALL_KWARGS
+        )
+        assert isinstance(stream_response, litellm.CustomStreamWrapper)
+        async for _chunk in stream_response:
+            pass
+    else:
+        litellm.completion(model="azure/gpt-4o", messages=messages, **_AZURE_CALL_KWARGS)
+
+
+async def _azure_responses(*, is_async: bool) -> None:
+    if is_async:
+        await litellm.aresponses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS)
+    else:
+        litellm.responses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS)
+
+
+async def _azure_embedding() -> None:
+    litellm.embedding(model="azure/text-embedding-3-small", input="hi", **_AZURE_CALL_KWARGS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "call",
+    [
+        pytest.param(lambda: _azure_chat(stream=False), id="chat"),
+        pytest.param(lambda: _azure_chat(stream=True), id="async-chat-stream"),
+        pytest.param(lambda: _azure_responses(is_async=False), id="responses"),
+        pytest.param(lambda: _azure_responses(is_async=True), id="async-responses"),
+        pytest.param(_azure_embedding, id="embedding"),
+    ],
+)
+@pytest.mark.parametrize(
+    ("code", "expected"),
+    [("context_length_exceeded", ContextWindowExceededError), ("invalid_value", BadRequestError)],
+)
+@respx.mock
+async def test_azure_400_is_classified_by_error_code(
+    monkeypatch: pytest.MonkeyPatch, call: Callable[[], Awaitable[None]], code: str, expected: type[Exception]
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    respx.post(url__startswith=_AZURE_API_BASE).mock(
+        return_value=_azure_error_response(400, {**_CONTEXT_LENGTH_ERROR, "code": code})
+    )
+
+    with pytest.raises(BadRequestError) as exc_info:
+        await call()
+
+    assert type(exc_info.value) is expected
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_azure_context_length_code_triggers_context_window_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    route: Final = respx.post(url__startswith=_AZURE_API_BASE).mock(
+        return_value=_azure_error_response(400, _CONTEXT_LENGTH_ERROR)
+    )
+    router: Final = litellm.Router(
+        model_list=[
+            {"model_name": "small", "litellm_params": {"model": "azure/gpt-4o", **_AZURE_CALL_KWARGS}},
+            {"model_name": "large", "litellm_params": {"model": "gpt-4.1", "mock_response": "served by large"}},
+        ],
+        context_window_fallbacks=[{"small": ["large"]}],
+        num_retries=0,
+    )
+
+    response: Final = await router.acompletion(  # pyright: ignore[reportUnknownMemberType]  # untyped **kwargs
+        model="small", messages=[{"role": "user", "content": "hi"}]
+    )
+
+    assert isinstance(response, litellm.ModelResponse)
+    assert response.choices[0].message.content == "served by large"
+    assert len(route.calls) == 1
+
+
+def _azure_chat_error(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> openai.APIError:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    respx.post(url__startswith=_AZURE_API_BASE).mock(return_value=response)
+    with pytest.raises(openai.APIError) as exc_info:
+        litellm.completion(model="azure/gpt-4o", messages=[{"role": "user", "content": "hi"}], **_AZURE_CALL_KWARGS)
+    return exc_info.value
+
+
+@respx.mock
+def test_azure_context_window_error_keeps_upstream_code_and_param(monkeypatch: pytest.MonkeyPatch) -> None:
+    error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, _CONTEXT_LENGTH_ERROR))
+
+    assert isinstance(error, ContextWindowExceededError)
+    assert error.code == "context_length_exceeded"
+    assert error.param == "input"
+
+
+@pytest.mark.parametrize("status_code", [429, 500])
+@respx.mock
+def test_azure_non_400_with_context_length_code_keeps_its_status(
+    monkeypatch: pytest.MonkeyPatch, status_code: int
+) -> None:
+    error: Final = _azure_chat_error(monkeypatch, _azure_error_response(status_code, _CONTEXT_LENGTH_ERROR))
+
+    assert not isinstance(error, ContextWindowExceededError)
+    assert getattr(error, "status_code", None) == status_code
+
+
+@respx.mock
+def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest.MonkeyPatch) -> None:
+    policy_error: Final = {**_CONTEXT_LENGTH_ERROR, "inner_error": {"code": "ResponsibleAIPolicyViolation"}}
+
+    error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error))
+
+    assert type(error) is ContentPolicyViolationError