mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
37f8b9600a
Bring the Entra ID / OAuth auth work for Azure AI Foundry routes up to date with staging and fix the lint-budget regressions the merge surfaced: - widen get_azure_ai_auth_headers return type to Mapping[str, str] (LIT001) - build the azure_ai image_generation request headers into a new Final local instead of rebinding the Final headers dict (reportGeneralTypeIssues) - order HuggingFace rerank validate_environment params to match BaseRerankConfig so litellm_params lines up positionally (reportIncompatibleMethodOverride) - add a match= to the credential-error test and document the handler-boundary patches the auth wiring tests rely on
119 lines
4.6 KiB
Python
119 lines
4.6 KiB
Python
"""
|
|
Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format.
|
|
"""
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Final
|
|
|
|
import httpx
|
|
|
|
import litellm
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers
|
|
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):
|
|
"""
|
|
Azure AI Rerank - Follows the same Spec as Cohere Rerank
|
|
"""
|
|
|
|
def get_complete_url(
|
|
self,
|
|
api_base: str | None,
|
|
model: str,
|
|
optional_params: dict | None = None,
|
|
) -> str:
|
|
if api_base is None:
|
|
raise ValueError(
|
|
"Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var."
|
|
)
|
|
original_url: Final = 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: Final = original_url.path.rstrip("/")
|
|
|
|
# Allow callers to pass either full v1/v2 rerank endpoints:
|
|
# - https://<resource>.services.ai.azure.com/v1/rerank
|
|
# - https://<resource>.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" or ".../providers/cohere/v2"), append "/rerank"
|
|
if (
|
|
normalized_path.endswith("/v1")
|
|
or normalized_path.endswith("/v2")
|
|
or normalized_path.endswith("/providers/cohere/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,
|
|
headers: dict,
|
|
model: str,
|
|
api_key: str | None = None,
|
|
optional_params: dict | None = None,
|
|
litellm_params: Mapping[str, object] | None = None,
|
|
) -> dict:
|
|
if api_key is None:
|
|
api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key
|
|
|
|
default_headers: Final = {
|
|
**get_azure_ai_auth_headers(api_key=api_key, litellm_params=litellm_params),
|
|
"accept": "application/json",
|
|
"content-type": "application/json",
|
|
}
|
|
|
|
# If 'Authorization' is provided in headers, it overrides the default.
|
|
if "Authorization" in headers:
|
|
default_headers["Authorization"] = headers["Authorization"]
|
|
|
|
# Merge other headers, overriding any default ones except Authorization
|
|
return {**default_headers, **headers}
|
|
|
|
def transform_rerank_response(
|
|
self,
|
|
model: str,
|
|
raw_response: httpx.Response,
|
|
model_response: RerankResponse,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
api_key: str | None = None,
|
|
request_data: dict = {},
|
|
optional_params: dict = {},
|
|
litellm_params: dict = {},
|
|
) -> RerankResponse:
|
|
rerank_response: Final = super().transform_rerank_response(
|
|
model=model,
|
|
raw_response=raw_response,
|
|
model_response=model_response,
|
|
logging_obj=logging_obj,
|
|
api_key=api_key,
|
|
request_data=request_data,
|
|
optional_params=optional_params,
|
|
litellm_params=litellm_params,
|
|
)
|
|
base_model: Final = self._get_base_model(rerank_response._hidden_params.get("llm_provider-azureml-model-group"))
|
|
rerank_response._hidden_params["model"] = base_model
|
|
return rerank_response
|
|
|
|
def _get_base_model(self, azure_model_group: str | None) -> str | None:
|
|
if azure_model_group is None:
|
|
return None
|
|
if azure_model_group == "offer-cohere-rerank-mul-paygo":
|
|
return "azure_ai/cohere-rerank-v3-multilingual"
|
|
if azure_model_group == "offer-cohere-rerank-eng-paygo":
|
|
return "azure_ai/cohere-rerank-v3-english"
|
|
return azure_model_group
|