Merge b04dc3b997 into 72228db5c6

litellm/llms/azure_ai/README.md

@@ -1 +1,39 @@
-`/chat/completion` calls routed via `openai.py`. 
+# Azure AI
+
+## Native Cohere embeddings on Microsoft Foundry
+
+Use `azure_ai/Cohere-Embed-V5-Fast` or `azure_ai/Cohere-Embed-V5-Pro` with the
+Foundry resource URL. LiteLLM sends the native Cohere request to
+`/providers/cohere/v2/embed` and calculates cost using the existing Azure AI
+pricing entries and the provider's billed text and image tokens.
+
+```python
+import os
+
+from litellm import embedding
+
+response = embedding(
+    model="azure_ai/Cohere-Embed-V5-Fast",
+    api_base="https://my-resource.services.ai.azure.com",
+    api_key=os.environ["AZURE_AI_API_KEY"],
+    input=["Searchable document text"],
+    input_type="search_document",
+    dimensions=256,
+)
+```
+
+The same configuration works with `aembedding`. Azure authentication also
+supports Entra access tokens, service principals, OIDC federation and managed
+identity through the existing Azure credential options. Omit the API key when
+using those credentials.
+
+The API base can be the resource root, `/models`, `/providers/cohere`,
+`/providers/cohere/v2`, or the complete `/providers/cohere/v2/embed` endpoint.
+For a custom deployment name, set the API base to one of the explicit
+`/providers/cohere` paths. Query parameters are preserved.
+
+Existing generic Azure embeddings and legacy Cohere v3/v4 deployments retain
+their current endpoints unless an explicit native Cohere API base is supplied.
+
+The native endpoint format is documented in
+[Cohere's Azure examples](https://github.com/cohere-ai/cohere-azure-workshops).

litellm/llms/azure_ai/embed/cohere_transformation.py

@@ -9,14 +9,71 @@ Convers
 Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html
 """
 
+from collections.abc import Mapping, Sequence
 from typing import Final
+from urllib.parse import urlsplit, urlunsplit
 
+from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo, get_azure_ai_auth_headers
+from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig
 from litellm.types.llms.azure_ai import ImageEmbeddingInput, ImageEmbeddingRequest
-from litellm.types.llms.openai import EmbeddingCreateParams
+from litellm.types.llms.openai import AllMessageValues, EmbeddingCreateParams
 from litellm.types.utils import EmbeddingResponse, Usage
 from litellm.utils import is_base64_encoded
 
 
+class AzureAICohereNativeEmbeddingConfig(CohereEmbeddingConfig):
+    @staticmethod
+    def supports_model(model: str, api_base: str | None = None) -> bool:
+        path: Final = urlsplit(api_base or "").path.rstrip("/")
+        name: Final = model.lower().removeprefix("azure_ai/")
+        return path.endswith(("/providers/cohere", "/providers/cohere/v2", "/providers/cohere/v2/embed")) or (
+            name.startswith("cohere-embed-") and not name.startswith(("cohere-embed-v3", "cohere-embed-v4"))
+        )
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        api_key: str | None,
+        model: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        stream: bool | None = None,
+    ) -> str:
+        resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
+        if resolved_base is None:
+            raise ValueError("Azure AI API base is required for native Cohere embeddings")
+        parts: Final = urlsplit(resolved_base)
+        path: Final = parts.path.rstrip("/")
+        endpoint: Final = (
+            path
+            if path.endswith("/providers/cohere/v2/embed")
+            else f"{path}/embed"
+            if path.endswith("/providers/cohere/v2")
+            else f"{path}/v2/embed"
+            if path.endswith("/providers/cohere")
+            else f"{path.removesuffix('/models')}/providers/cohere/v2/embed"
+        )
+        return urlunsplit((parts.scheme, parts.netloc, endpoint, parts.query, parts.fragment))
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        messages: Sequence[AllMessageValues],
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> dict[str, str]:  # mutable-ok: BaseEmbeddingConfig requires a mutable header dictionary
+        return {
+            "Content-Type": "application/json",
+            **get_azure_ai_auth_headers(
+                api_key=AzureFoundryModelInfo.get_api_key(api_key), litellm_params=litellm_params
+            ),
+            **headers,
+        }
+
+
 class AzureAICohereConfig:
     def __init__(self) -> None:
         pass

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1401,7 +1401,7 @@ class BaseLLMHTTPHandler:
         headers: dict[str, Any] | None = None,
     ) -> EmbeddingResponse:
         provider_config: Final = ProviderConfigManager.get_provider_embedding_config(
-            model=model, provider=litellm.LlmProviders(custom_llm_provider)
+            model=model, provider=litellm.LlmProviders(custom_llm_provider), api_base=api_base
         )
         if provider_config is None:
             raise ValueError(f"Provider {custom_llm_provider} does not support embedding")

litellm/main.py

@@ -6987,6 +6987,7 @@ def embedding(
         encoding_format=encoding_format,
         custom_llm_provider=custom_llm_provider,
         allowed_openai_params=allowed_openai_params,
+        api_base=api_base,
         **non_default_params,
     )
 
@@ -7198,7 +7199,13 @@ def embedding(
                 aembedding=aembedding,
                 drop_params=drop_unsupported_params,
             )
-        elif custom_llm_provider == "oci":
+        elif custom_llm_provider == "oci" or (
+            custom_llm_provider == "azure_ai"
+            and ProviderConfigManager.get_provider_embedding_config(
+                model=model, provider=LlmProviders.AZURE_AI, api_base=api_base
+            )
+            is not None
+        ):
             response = base_llm_http_handler.embedding(
                 model=model,
                 input=input,
@@ -7702,7 +7709,6 @@ def embedding(
                 or get_azure_ai_entra_token(litellm_params=litellm_params_dict)
             )
 
-            ## EMBEDDING CALL
             response = azure_ai_embedding.embedding(
                 model=model,
                 input=input,

litellm/utils.py

@@ -4769,6 +4769,7 @@ def get_optional_params_embeddings(
     drop_params: bool | None = None,
     additional_drop_params: list[str] | None = None,
     allowed_openai_params: list[str] | None = None,
+    api_base: str | None = None,
     **kwargs,
 ):
     # Lazy load get_supported_openai_params
@@ -4777,6 +4778,7 @@ def get_optional_params_embeddings(
     # retrieve all parameters passed to the function
     passed_params: Final = locals()
     passed_params.pop("custom_llm_provider", None)
+    passed_params.pop("api_base", None)
     special_params: Final = passed_params.pop("kwargs")
 
     passed_params.pop("drop_params", None)
@@ -4818,6 +4820,7 @@ def get_optional_params_embeddings(
         provider_config = ProviderConfigManager.get_provider_embedding_config(
             model=model,
             provider=LlmProviders(custom_llm_provider),
+            api_base=api_base,
         )
 
     if provider_config is not None:
@@ -10168,6 +10171,7 @@ class ProviderConfigManager:
     def get_provider_embedding_config(
         model: str,
         provider: LlmProviders,
+        api_base: str | None = None,
     ) -> BaseEmbeddingConfig | None:
         if (
             litellm.LlmProviders.VOYAGE == provider
@@ -10195,6 +10199,12 @@ class ProviderConfigManager:
             from litellm.llms.oci.embed.transformation import OCIEmbedConfig
 
             return OCIEmbedConfig()
+        elif litellm.LlmProviders.AZURE_AI == provider:
+            from litellm.llms.azure_ai.embed.cohere_transformation import AzureAICohereNativeEmbeddingConfig
+
+            if AzureAICohereNativeEmbeddingConfig.supports_model(model, api_base):
+                return AzureAICohereNativeEmbeddingConfig()
+            return None
         elif litellm.LlmProviders.COHERE == provider or litellm.LlmProviders.COHERE_CHAT == provider:
             from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig
 

tests/unit/llms/azure_ai/embed/test_azure_ai_cohere_transformation.py

@@ -1,9 +1,124 @@
+import json
+from collections.abc import Iterator
 from typing import Final
 
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm.caching.llm_caching_handler import LLMClientCache
 from litellm.llms.azure_ai.embed.cohere_transformation import AzureAICohereConfig
 from litellm.types.utils import EmbeddingResponse, Usage
 
 
+@pytest.fixture
+def native_azure_embedding_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    clients: Final = LLMClientCache()
+    monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", clients)
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "force_ipv4", False)
+    monkeypatch.setattr(litellm, "api_key", None)
+    for key in ("AZURE_AI_API_KEY", "AZURE_API_KEY", "AZURE_AD_TOKEN", "AZURE_CLIENT_SECRET"):
+        monkeypatch.delenv(key, raising=False)
+    monkeypatch.setenv("COHERE_API_KEY", "unrelated-cohere-key")
+    yield
+    clients.flush_cache()
+
+
+@pytest.mark.usefixtures("native_azure_embedding_transport", "local_model_cost_map")
+@pytest.mark.respx(assert_all_called=True)
+@pytest.mark.parametrize("model", ["Cohere-Embed-V5-Fast", "Cohere-Embed-V5-Pro"])
+@pytest.mark.parametrize("sync_mode", [False, True])
+@pytest.mark.parametrize("auth", [{"api_key": "azure-key"}, {"azure_ad_token": "entra-token"}])
+@pytest.mark.parametrize(("text_tokens", "image_tokens"), [(5, 0), (0, 6), (3, 6)])
+async def test_native_foundry_embeddings_keep_azure_auth_and_price_billed_usage(
+    respx_mock: respx.MockRouter,
+    model: str,
+    sync_mode: bool,
+    auth: dict[str, str],
+    text_tokens: int,
+    image_tokens: int,
+) -> None:
+    route: Final = respx_mock.post("https://foundry.services.ai.azure.com/providers/cohere/v2/embed").mock(
+        return_value=httpx.Response(
+            200,
+            json={
+                "embeddings": {"float": [[0.125, -0.375]]},
+                "meta": {"billed_units": {"input_tokens": text_tokens, "image_tokens": image_tokens}},
+            },
+        )
+    )
+    input_value: Final = "data:image/png;base64,aGVsbG8=" if image_tokens else "hello world"
+    params: Final = {
+        "model": f"azure_ai/{model}",
+        "api_base": "https://foundry.services.ai.azure.com",
+        "input": [input_value],
+        "encoding_format": "float",
+        "dimensions": 256,
+        **auth,
+    }
+
+    response: Final = litellm.embedding(**params) if sync_mode else await litellm.aembedding(**params)
+
+    assert json.loads(route.calls.last.request.content) == {
+        "model": model,
+        "images" if image_tokens else "texts": [input_value],
+        "input_type": "image" if image_tokens else litellm.COHERE_DEFAULT_EMBEDDING_INPUT_TYPE,
+        "embedding_types": ["float"],
+        "output_dimension": 256,
+    }
+    assert route.calls.last.request.headers["Authorization"] == f"Bearer {next(iter(auth.values()))}"
+    assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.125, -0.375]}]
+    assert response.usage is not None
+    assert response.usage.prompt_tokens == text_tokens + image_tokens
+    assert response.usage.prompt_tokens_details is not None
+    assert response.usage.prompt_tokens_details.text_tokens == text_tokens
+    assert response.usage.prompt_tokens_details.image_tokens == (image_tokens or None)
+    prices: Final = litellm.get_model_info(f"azure_ai/{model}")
+    assert response._hidden_params["custom_llm_provider"] == "azure_ai"
+    assert response._hidden_params["response_cost"] == pytest.approx(
+        text_tokens * prices["input_cost_per_token"] + image_tokens * prices["input_cost_per_image_token"]
+    )
+
+
+@pytest.mark.usefixtures("native_azure_embedding_transport")
+@pytest.mark.respx(assert_all_called=True)
+@pytest.mark.parametrize(
+    ("model", "base_path"),
+    [
+        ("Cohere-Embed-V5-Fast", "/models"),
+        ("Cohere-Embed-V5-Fast", "/"),
+        ("custom-embed-deployment", "/providers/cohere"),
+        ("custom-embed-deployment", "/providers/cohere/v2/"),
+        ("custom-embed-deployment", "/providers/cohere/v2/embed/"),
+    ],
+)
+async def test_native_foundry_embeddings_accept_deployment_names_and_endpoint_forms(
+    respx_mock: respx.MockRouter, model: str, base_path: str
+) -> None:
+    route: Final = respx_mock.post(
+        "https://foundry.services.ai.azure.com/providers/cohere/v2/embed?api-version=preview"
+    ).mock(return_value=httpx.Response(200, json={"embeddings": {"float": [[0.5]]}}))
+
+    response: Final = await litellm.aembedding(
+        model=f"azure_ai/{model}",
+        input=["query"],
+        api_base=f"https://foundry.services.ai.azure.com{base_path}?api-version=preview",
+        api_key="azure-key",
+        input_type="search_query",
+        dimensions=256,
+    )
+
+    assert json.loads(route.calls.last.request.content) == {
+        "model": model,
+        "texts": ["query"],
+        "input_type": "search_query",
+        "output_dimension": 256,
+    }
+    assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.5]}]
+
+
 def test_azure_ai_cohere_request_splits_images_from_text_and_keeps_order():
     config: Final = AzureAICohereConfig()
     image_request, text_request, image_indexes = config.transform_request(
@@ -15,9 +130,7 @@ def test_azure_ai_cohere_request_splits_images_from_text_and_keeps_order():
         model="Cohere-embed-v3-multilingual",
     )
 
-    assert image_request == {
-        "input": [{"image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"}]
-    }
+    assert image_request == {"input": [{"image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"}]}
     assert text_request == {"input": ["describe this image"], "model": "Cohere-embed-v3-multilingual"}
     assert image_indexes == [1]