chore: initial public snapshot for github upload

llm-gateway-competitors/litellm-wheel-src/litellm/llms/azure_ai/rerank/handler.py

@@ -0,0 +1,5 @@
+"""
+Azure AI Rerank - uses `llm_http_handler.py` to make httpx requests
+
+Request/Response transformation is handled in `transformation.py`
+"""

llm-gateway-competitors/litellm-wheel-src/litellm/llms/azure_ai/rerank/transformation.py

@@ -0,0 +1,125 @@
+"""
+Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format. 
+"""
+
+from typing import Optional
+
+import httpx
+
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+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: Optional[str],
+        model: str,
+        optional_params: Optional[dict] = 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 = 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:
+        # - 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: Optional[str] = None,
+        optional_params: Optional[dict] = None,
+    ) -> dict:
+        if api_key is None:
+            api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key
+
+        if api_key is None:
+            raise ValueError(
+                "Azure AI API key is required. Please set 'AZURE_AI_API_KEY' or 'litellm.azure_key'"
+            )
+
+        default_headers = {
+            "Authorization": f"Bearer {api_key}",
+            "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: Optional[str] = None,
+        request_data: dict = {},
+        optional_params: dict = {},
+        litellm_params: dict = {},
+    ) -> RerankResponse:
+        rerank_response = 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 = 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: Optional[str]) -> Optional[str]:
+        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