diff --git a/litellm/llms/azure_ai/README.md b/litellm/llms/azure_ai/README.md index 8c521519da1..694fe77d4ed 100644 --- a/litellm/llms/azure_ai/README.md +++ b/litellm/llms/azure_ai/README.md @@ -1 +1,39 @@ -`/chat/completion` calls routed via `openai.py`. \ No newline at end of file +# 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). diff --git a/litellm/llms/azure_ai/embed/cohere_transformation.py b/litellm/llms/azure_ai/embed/cohere_transformation.py index 82e855d8706..352d9c73ecb 100644 --- a/litellm/llms/azure_ai/embed/cohere_transformation.py +++ b/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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..619e719ef09 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1385,7 +1385,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") diff --git a/litellm/main.py b/litellm/main.py index 6b9964a450f..d210b46a99f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6961,6 +6961,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, ) @@ -7172,7 +7173,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, @@ -7676,7 +7683,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, diff --git a/litellm/utils.py b/litellm/utils.py index 53b02ebddac..709c50db309 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4759,6 +4759,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 @@ -4767,6 +4768,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) @@ -4808,6 +4810,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: @@ -10137,6 +10140,7 @@ class ProviderConfigManager: def get_provider_embedding_config( model: str, provider: LlmProviders, + api_base: str | None = None, ) -> BaseEmbeddingConfig | None: if ( litellm.LlmProviders.VOYAGE == provider @@ -10164,6 +10168,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 diff --git a/tests/unit/llms/azure_ai/embed/test_azure_ai_cohere_transformation.py b/tests/unit/llms/azure_ai/embed/test_azure_ai_cohere_transformation.py index 19ae0cb1f8e..b23b5fd00af 100644 --- a/tests/unit/llms/azure_ai/embed/test_azure_ai_cohere_transformation.py +++ b/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]