diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index a820ac7f345..21de0406864 100644 --- a/litellm/llms/custom_llm.py +++ b/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, diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index e27585116ce..057236f0d5d 100644 --- a/litellm/rerank_api/main.py +++ b/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, diff --git a/tests/test_litellm/rerank_api/test_rerank_custom_provider.py b/tests/test_litellm/rerank_api/test_rerank_custom_provider.py new file mode 100644 index 00000000000..da2399dc5b1 --- /dev/null +++ b/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 +