From 604f1fde5014d3a966e1a2ee814e28f83bf97549 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:46:33 -0700 Subject: [PATCH] fix(azure_ai): route Foundry embeddings to the /models inference route --- litellm/llms/azure_ai/chat/transformation.py | 7 +- litellm/llms/azure_ai/common_utils.py | 9 +++ litellm/llms/azure_ai/embed/handler.py | 19 ++++- .../embed/test_azure_ai_embed_handler.py | 69 +++++++++++++++++++ 4 files changed, 96 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 8797f4dc400..f2d405e9a17 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/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): diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 26a90157455..aa34bab5b2e 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/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. diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 65c3997c099..c65edbf56e6 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/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, diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py new file mode 100644 index 00000000000..0401629171d --- /dev/null +++ b/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