From c4dd22c07966bc7fa2f0de8471f26092aef90a06 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 16 Dec 2025 19:31:10 -0600 Subject: [PATCH] fix: broaden Azure AI rerank URL handling --- .../llms/azure_ai/rerank/transformation.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index a47b6082c3..085f36f74b 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import RerankResponse +from litellm.utils import _add_path_to_api_base class AzureAIRerankConfig(CohereRerankConfig): @@ -28,9 +29,24 @@ class AzureAIRerankConfig(CohereRerankConfig): raise ValueError( "Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var." ) - if not api_base.endswith("/v1/rerank"): - api_base = f"{api_base}/v1/rerank" - return api_base + original_url = httpx.URL(api_base) + normalized_path = original_url.path.rstrip("/") + + # Allow callers to pass either full v1/v2 rerank endpoints: + # - https://.services.ai.azure.com/v1/rerank + # - https://.services.ai.azure.com/providers/cohere/v2/rerank + if normalized_path.endswith("/v1/rerank") or normalized_path.endswith("/v2/rerank"): + return str(original_url.copy_with(path=normalized_path or "/")) + + # If callers pass just the version path (e.g. ".../v2"), append "/rerank" + if normalized_path.endswith("/v1") or normalized_path.endswith("/v2"): + return _add_path_to_api_base( + api_base=str(original_url.copy_with(path=normalized_path or "/")), + ending_path="/rerank", + ) + + # Backwards compatible default: Azure AI rerank was originally exposed under /v1/rerank + return _add_path_to_api_base(api_base=api_base, ending_path="/v1/rerank") def validate_environment( self,