Merge pull request #38975 from BerriAI/litellm_fix_azure_ai_reclassify

fix(azure_ai): don't reclassify Foundry deployments as azure provider

litellm/constants.py

@@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran
 
 DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
 AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
+AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
 ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
 ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
 RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(

litellm/llms/azure_ai/chat/transformation.py

@@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     filter_value_from_dict,
 )
 from litellm.llms.azure.common_utils import BaseAzureLLM
+from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
 from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
 from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
 from litellm.llms.openai.openai import OpenAIConfig
@@ -207,20 +208,18 @@ class AzureAIStudioConfig(OpenAIConfig):
                 message["content"] = texts
         return stripped_messages
 
-    def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
-        try:
-            if "/" in model:
-                model = model.split("/", 1)[1]
-            if (
-                model in litellm.open_ai_chat_completion_models
-                or model in litellm.open_ai_text_completion_models
-                or model in litellm.open_ai_embedding_models
-            ):
-                return True
+    def _is_foundry_model_inference_base(self, api_base: str) -> bool:
+        return is_foundry_model_inference_base(api_base)
 
-        except Exception:
+    def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
+        if api_base is None or self._is_foundry_model_inference_base(api_base):
             return False
-        return False
+        stripped_model: Final = model.split("/", 1)[1] if "/" in model else model
+        return (
+            stripped_model in litellm.open_ai_chat_completion_models
+            or stripped_model in litellm.open_ai_text_completion_models
+            or stripped_model in litellm.open_ai_embedding_models
+        )
 
     def _get_openai_compatible_provider_info(
         self,

litellm/llms/azure_ai/common_utils.py

@@ -1,5 +1,6 @@
 from collections.abc import Mapping
 from typing import Final, Literal
+from urllib.parse import urlparse
 
 import litellm
 from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
@@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams
 AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
 
 
+def is_foundry_model_inference_base(api_base: str) -> bool:
+    parsed: Final = urlparse(api_base)
+    host: Final = parsed.hostname
+    if host is None or not host.endswith(".services.ai.azure.com"):
+        return False
+    return "/openai/deployments" not in parsed.path
+
+
 def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
     """
     Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment.

litellm/llms/azure_ai/embed/handler.py

@@ -1,8 +1,10 @@
 from typing import Final
+from urllib.parse import urlsplit, urlunsplit
 
 from openai import OpenAI
 
 import litellm
+from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object
 from .cohere_transformation import AzureAICohereConfig
 
 
+def _foundry_models_route_base(api_base: str | None) -> str | None:
+    if api_base is None or not is_foundry_model_inference_base(api_base):
+        return api_base
+    parts: Final = urlsplit(api_base)
+    path: Final = parts.path.rstrip("/")
+    if path.endswith("/models"):
+        return api_base
+    return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment))
+
+
 class AzureAIEmbedding(OpenAIChatCompletion):
     def _process_response(
         self,
@@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
 
         assemble result in-order, and return
         """
+        resolved_api_base: Final = _foundry_models_route_base(api_base)
         if aembedding is True:
             return self.async_embedding(
                 model,
@@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response,
                 optional_params,
                 api_key,
-                api_base,
+                resolved_api_base,
                 client,
             )
 
@@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response=model_response,
                 optional_params=optional_params,
                 api_key=api_key,
-                api_base=api_base,
+                api_base=resolved_api_base,
                 client=client,
             )
 
@@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
                 model_response,
                 optional_params,
                 api_key,
-                api_base,
+                resolved_api_base,
                 client=(client if client is not None and isinstance(client, OpenAI) else None),
                 aembedding=aembedding,
                 shared_session=shared_session,

litellm/main.py

@@ -61,6 +61,7 @@ if TYPE_CHECKING:
     from litellm.types.utils import TokenCountResponse
 
 from litellm.constants import (
+    AZURE_OPENAI_AUDIO_PROVIDERS,
     DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
     DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
 )
@@ -7777,7 +7778,7 @@ def transcription(
         provider=LlmProviders(custom_llm_provider),
     )
 
-    if custom_llm_provider == "azure" and provider_config is None:
+    if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
         # azure configs
         api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
 
@@ -8064,7 +8065,10 @@ def speech(
         custom_llm_provider=custom_llm_provider,
     )
     response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None
-    if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers:
+    if custom_llm_provider == "openai" or (
+        custom_llm_provider in litellm.openai_compatible_providers
+        and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS
+    ):
         if voice is None or not (isinstance(voice, str)):
             raise litellm.BadRequestError(
                 message="'voice' is required to be passed as a string for OpenAI TTS",
@@ -8118,7 +8122,7 @@ def speech(
             aspeech=aspeech,
             shared_session=shared_session,
         )
-    elif custom_llm_provider == "azure":
+    elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
         # Check if this is Azure Speech Service (Cognitive Services TTS)
         if model.startswith("speech/"):
             from litellm.llms.azure.text_to_speech.transformation import (

litellm/realtime_api/main.py

@@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast
 
 import litellm
 from litellm.constants import (
+    AZURE_OPENAI_AUDIO_PROVIDERS,
     REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
     REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
     request_timeout,
@@ -400,7 +401,7 @@ async def _arealtime(
             litellm_metadata=_build_litellm_metadata(kwargs),
             query_params=query_params,
         )
-    elif _custom_llm_provider == "azure":
+    elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
         api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
         # set API KEY
         api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY")

tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py

@@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info():
     assert custom_llm_provider == "azure"
 
 
+@pytest.mark.parametrize(
+    "model, api_base, expected_provider",
+    [
+        ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"),
+        ("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"),
+        (
+            "azure_ai/gpt-4o",
+            "https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions"
+            "?api-version=2024-08-01-preview",
+            "azure",
+        ),
+        ("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"),
+        ("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"),
+    ],
+)
+def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str):
+    """Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure."""
+    config = AzureAIStudioConfig()
+    (
+        _,
+        _,
+        custom_llm_provider,
+    ) = config._get_openai_compatible_provider_info(
+        model=model,
+        api_base=api_base,
+        api_key="my-key",
+        custom_llm_provider="azure_ai",
+    )
+    assert custom_llm_provider == expected_provider
+
+
+def test_is_azure_openai_model_without_api_base_keeps_azure_ai():
+    """Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider."""
+    config = AzureAIStudioConfig()
+    assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False
+    assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True
+
+
 def test_azure_ai_validate_environment():
     config = AzureAIStudioConfig()
     headers = config.validate_environment(

tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py

@@ -0,0 +1,69 @@
+import httpx
+import pytest
+import respx
+
+from litellm import embedding
+from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base
+
+EMBEDDING_PAYLOAD = {
+    "object": "list",
+    "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}],
+    "model": "text-embedding-3-small",
+    "usage": {"prompt_tokens": 2, "total_tokens": 2},
+}
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected"),
+    [
+        (
+            "https://my-foundry.services.ai.azure.com",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview",
+            "https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/models",
+            "https://my-foundry.services.ai.azure.com/models",
+        ),
+        (
+            "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
+            "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
+        ),
+        (
+            "https://my-resource.openai.azure.com",
+            "https://my-resource.openai.azure.com",
+        ),
+        (
+            "https://Mistral-serverless.eastus2.models.ai.azure.com",
+            "https://Mistral-serverless.eastus2.models.ai.azure.com",
+        ),
+        (None, None),
+    ],
+)
+def test_foundry_models_route_base(api_base, expected):
+    assert _foundry_models_route_base(api_base) == expected
+
+
+@respx.mock
+def test_azure_ai_embedding_calls_foundry_models_route():
+    route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock(
+        return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD)
+    )
+
+    response = embedding(
+        model="azure_ai/text-embedding-3-small",
+        input=["hello world"],
+        api_base="https://my-foundry.services.ai.azure.com",
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.data is not None
+    assert len(response.data) == 1

tests/test_litellm/realtime_api/test_main.py

@@ -1,6 +1,7 @@
 import asyncio
 import time
 from types import TracebackType
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 
@@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol():
             model_params={"realtime_protocol": "GA"},
         )
     assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview"
+
+
+class _ConnectThatStopsAfterCapturingTheUrl:
+    url: str | None = None
+
+    def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl":
+        self.url = url
+        return self
+
+    async def __aenter__(self) -> None:
+        raise RuntimeError("backend url captured, nothing to bridge")
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route():
+    connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
+    with patch("websockets.connect", connect):
+        await realtime_main._arealtime.__wrapped__(
+            model="azure_ai/gpt-realtime-mini",
+            websocket=MagicMock(),
+            api_base="https://my-project.services.ai.azure.com",
+            api_key="fake-key",
+            litellm_logging_obj=FakeLogging(),
+        )
+    assert connect.url == (
+        "wss://my-project.services.ai.azure.com/openai/realtime"
+        "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini"
+    )

tests/test_litellm/test_gpt_5_5_model_metadata.py

@@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model):
 
     routed_model, provider, _, _ = get_llm_provider(model=model)
     assert routed_model == model.split("/", 1)[1]
-    # azure_ai/* models resolve under the azure provider in get_llm_provider
-    assert provider == "azure"
+    assert provider == "azure_ai"
 
 
 def test_azure_ai_gpt_5_5_backup_matches_main():

tests/test_litellm/test_main.py

@@ -3335,3 +3335,43 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp
     assert getattr(response.usage, "cost", None) == pytest.approx(0.42)
     assert response._hidden_params.get("response_cost") is None
     assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63)
+
+
+FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com"
+
+
+def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route(
+    respx_mock: respx.MockRouter,
+):
+    route: Final = respx_mock.post(
+        url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+"
+    ).mock(return_value=httpx.Response(200, json={"text": "hello"}))
+
+    response: Final = litellm.transcription(
+        model="azure_ai/whisper-1",
+        file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"),
+        api_base=FOUNDRY_HOST,
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.text == "hello"
+
+
+def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route(
+    respx_mock: respx.MockRouter,
+):
+    route: Final = respx_mock.post(
+        url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+"
+    ).mock(return_value=httpx.Response(200, content=b"mp3-bytes"))
+
+    response: Final = litellm.speech(
+        model="azure_ai/tts-1",
+        input="hello",
+        voice="alloy",
+        api_base=FOUNDRY_HOST,
+        api_key="fake-key",
+    )
+
+    assert route.called
+    assert response.content == b"mp3-bytes"