fix: enforce scheme for Azure AI rerank api_base

This commit is contained in:
Emerson Gomes
2026-02-02 18:13:04 +05:30
committed by Sameer Kankute
parent 92763a14a9
commit bb5397d9b2
2 changed files with 106 additions and 0 deletions
@@ -30,6 +30,12 @@ class AzureAIRerankConfig(CohereRerankConfig):
"Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var."
)
original_url = httpx.URL(api_base)
if not original_url.is_absolute_url:
raise ValueError(
"Azure AI API Base must be an absolute URL including scheme (e.g. "
"'https://<resource>.services.ai.azure.com'). "
f"Got api_base={api_base!r}."
)
normalized_path = original_url.path.rstrip("/")
# Allow callers to pass either full v1/v2 rerank endpoints:
@@ -0,0 +1,100 @@
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.azure_ai.rerank.transformation import AzureAIRerankConfig
class TestAzureAIRerankConfigGetCompleteUrl:
def setup_method(self):
self.config = AzureAIRerankConfig()
self.model = "azure_ai/cohere-rerank-v3-english"
def test_api_base_required(self):
with pytest.raises(ValueError) as exc_info:
self.config.get_complete_url(api_base=None, model=self.model)
assert "api_base=None" in str(exc_info.value)
@pytest.mark.parametrize(
"api_base",
[
"example.com",
"example.com/v1",
"//example.com/v1",
"/v1/rerank",
],
)
def test_api_base_requires_scheme(self, api_base):
with pytest.raises(ValueError) as exc_info:
self.config.get_complete_url(api_base=api_base, model=self.model)
error_message = str(exc_info.value).lower()
assert "absolute url" in error_message
assert "scheme" in error_message
@pytest.mark.parametrize(
"api_base, expected_url",
[
(
"https://my-resource.services.ai.azure.com/v1/rerank/",
"https://my-resource.services.ai.azure.com/v1/rerank",
),
(
"https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank/",
"https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank",
),
],
)
def test_preserves_full_rerank_endpoint(self, api_base, expected_url):
url = self.config.get_complete_url(api_base=api_base, model=self.model)
assert url == expected_url
@pytest.mark.parametrize(
"api_base, expected_url",
[
(
"https://my-resource.services.ai.azure.com/v1",
"https://my-resource.services.ai.azure.com/v1/rerank",
),
(
"https://my-resource.services.ai.azure.com/v2/",
"https://my-resource.services.ai.azure.com/v2/rerank",
),
(
"https://my-resource.services.ai.azure.com/providers/cohere/v2",
"https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank",
),
(
"https://my-resource.services.ai.azure.com/providers/cohere/v2/",
"https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank",
),
],
)
def test_appends_rerank_for_version_paths(self, api_base, expected_url):
url = self.config.get_complete_url(api_base=api_base, model=self.model)
assert url == expected_url
@pytest.mark.parametrize(
"api_base",
[
"https://my-resource.services.ai.azure.com",
"https://my-resource.services.ai.azure.com/",
],
)
def test_defaults_to_v1_rerank_when_base_has_no_path(self, api_base):
url = self.config.get_complete_url(api_base=api_base, model=self.model)
assert url == "https://my-resource.services.ai.azure.com/v1/rerank"
def test_preserves_query_params(self):
url = self.config.get_complete_url(
api_base="https://my-resource.services.ai.azure.com/v1?r=1",
model=self.model,
)
assert url == "https://my-resource.services.ai.azure.com/v1/rerank?r=1"