fix(azure_ai): route Foundry embeddings to the /models inference route

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
@@ -208,11 +209,7 @@ class AzureAIStudioConfig(OpenAIConfig):
         return stripped_messages
 
     def _is_foundry_model_inference_base(self, 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
+        return is_foundry_model_inference_base(api_base)
 
     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):

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,

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