Fix handling of custom_provider map in rerank()

litellm/llms/custom_llm.py

@@ -197,6 +197,42 @@ class CustomLLM(BaseLLM):
     ) -> EmbeddingResponse:
         raise CustomLLMError(status_code=500, message="Not implemented yet!")
 
+    def rerank(
+        self,
+        model: str,
+        query: str,
+        documents: list,
+        top_n: Optional[int],
+        rank_fields: Optional[list],
+        return_documents: Optional[bool],
+        max_chunks_per_doc: Optional[int],
+        logging_obj: Any,
+        optional_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+        timeout: Optional[Union[float, httpx.Timeout]] = None,
+        litellm_params=None,
+    ) -> Any:
+        raise CustomLLMError(status_code=500, message="Not implemented yet!")
+
+    async def arerank(
+        self,
+        model: str,
+        query: str,
+        documents: list,
+        top_n: Optional[int],
+        rank_fields: Optional[list],
+        return_documents: Optional[bool],
+        max_chunks_per_doc: Optional[int],
+        logging_obj: Any,
+        optional_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+        timeout: Optional[Union[float, httpx.Timeout]] = None,
+        litellm_params=None,
+    ) -> Any:
+        raise CustomLLMError(status_code=500, message="Not implemented yet!")
+
     def image_edit(
         self,
         model: str,

litellm/rerank_api/main.py

@@ -5,10 +5,12 @@ from typing import Any, Coroutine, Dict, List, Literal, Optional, Union
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.exceptions import LiteLLMUnknownProvider
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
 from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.llms.custom_llm import CustomLLM
 from litellm.llms.together_ai.rerank.handler import TogetherAIRerank
 from litellm.llms.watsonx.common_utils import IBMWatsonXMixin
 from litellm.rerank_api.rerank_utils import get_optional_rerank_params
@@ -134,6 +136,37 @@ def rerank(  # noqa: PLR0915
             api_key=optional_params.api_key,
         )
 
+        # Handle custom providers registered via custom_provider_map before
+        # attempting to cast the provider string to the LlmProviders enum.
+        if _custom_llm_provider in litellm._custom_providers:
+            custom_handler: Optional[CustomLLM] = None
+            for item in litellm.custom_provider_map:
+                if item["provider"] == _custom_llm_provider:
+                    custom_handler = item["custom_handler"]
+
+            if custom_handler is None:
+                raise LiteLLMUnknownProvider(
+                    model=model, custom_llm_provider=_custom_llm_provider
+                )
+
+            handler_fn = custom_handler.arerank if _is_async else custom_handler.rerank
+
+            return handler_fn(
+                model=model,
+                query=query,
+                documents=documents,
+                top_n=top_n,
+                rank_fields=rank_fields,
+                return_documents=return_documents,
+                max_chunks_per_doc=max_chunks_per_doc,
+                logging_obj=litellm_logging_obj,
+                optional_params=optional_params.model_dump(exclude_unset=True),
+                api_key=dynamic_api_key or optional_params.api_key,
+                api_base=dynamic_api_base or optional_params.api_base,
+                timeout=optional_params.timeout,
+                litellm_params=optional_params.model_dump(exclude_unset=True),
+            )
+
         rerank_provider_config: BaseRerankConfig = (
             ProviderConfigManager.get_provider_rerank_config(
                 model=model,

tests/test_litellm/rerank_api/test_rerank_custom_provider.py

@@ -0,0 +1,97 @@
+"""
+Unit tests for custom_provider_map support in rerank() / arerank().
+
+Regression test for: ValueError: 'X' is not a valid LlmProviders
+when a provider registered via custom_provider_map is used with
+mode: rerank (e.g. during health checks).
+"""
+
+import pytest
+
+import litellm
+from litellm import CustomLLM
+from litellm.types.rerank import RerankResponse, RerankResponseResult
+
+
+def test_rerank_custom_provider_dispatches_to_handler():
+    """rerank() must route to the custom handler without crashing on the LlmProviders enum cast."""
+
+    class MyRerankLLM(CustomLLM):
+        def rerank(
+            self,
+            model,
+            query,
+            documents,
+            top_n,
+            rank_fields,
+            return_documents,
+            max_chunks_per_doc,
+            logging_obj,
+            optional_params,
+            api_key=None,
+            api_base=None,
+            timeout=None,
+            litellm_params=None,
+        ):
+            result: RerankResponseResult = {"index": 0, "relevance_score": 0.99}
+            return RerankResponse(id="test-id", results=[result])
+
+    handler = MyRerankLLM()
+    litellm.custom_provider_map = [
+        {"provider": "custom_rerank_llm", "custom_handler": handler}
+    ]
+
+    resp = litellm.rerank(
+        model="custom_rerank_llm/my-fake-model",
+        query="What is the capital of France?",
+        documents=["Paris is the capital.", "London is a city."],
+        top_n=1,
+    )
+
+    assert isinstance(resp, RerankResponse)
+    assert resp.id == "test-id"
+    assert resp.results is not None
+    assert resp.results[0]["relevance_score"] == 0.99
+
+
+@pytest.mark.asyncio
+async def test_arerank_custom_provider_dispatches_to_handler():
+    """arerank() must route to the custom handler without crashing on the LlmProviders enum cast."""
+
+    class MyRerankLLM(CustomLLM):
+        async def arerank(
+            self,
+            model,
+            query,
+            documents,
+            top_n,
+            rank_fields,
+            return_documents,
+            max_chunks_per_doc,
+            logging_obj,
+            optional_params,
+            api_key=None,
+            api_base=None,
+            timeout=None,
+            litellm_params=None,
+        ):
+            result: RerankResponseResult = {"index": 0, "relevance_score": 0.88}
+            return RerankResponse(id="async-test-id", results=[result])
+
+    handler = MyRerankLLM()
+    litellm.custom_provider_map = [
+        {"provider": "custom_rerank_llm_async", "custom_handler": handler}
+    ]
+
+    resp = await litellm.arerank(
+        model="custom_rerank_llm_async/my-fake-model",
+        query="What is the capital of France?",
+        documents=["Paris is the capital.", "London is a city."],
+        top_n=1,
+    )
+
+    assert isinstance(resp, RerankResponse)
+    assert resp.id == "async-test-id"
+    assert resp.results is not None
+    assert resp.results[0]["relevance_score"] == 0.88
+