diff --git a/litellm/__init__.py b/litellm/__init__.py index c71a984be3..c5ecf81877 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1098,6 +1098,7 @@ from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig from .llms.infinity.rerank.transformation import InfinityRerankConfig from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig +from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from .llms.clarifai.chat.transformation import ClarifaiConfig diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 9d6a797afd..8316e923df 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -156,7 +156,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}" ) - return RerankResponse(**raw_response_json) + return self._transform_response(raw_response_json) def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] diff --git a/litellm/utils.py b/litellm/utils.py index 7fa9de70a6..e12387899b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7294,6 +7294,8 @@ class ProviderConfigManager: return litellm.InfinityRerankConfig() elif litellm.LlmProviders.JINA_AI == provider: return litellm.JinaAIRerankConfig() + elif litellm.LlmProviders.HOSTED_VLLM == provider: + return litellm.HostedVLLMRerankConfig() elif litellm.LlmProviders.HUGGINGFACE == provider: return litellm.HuggingFaceRerankConfig() elif litellm.LlmProviders.DEEPINFRA == provider: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index dde2b45aee..b30b048995 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -845,6 +845,16 @@ def test_check_provider_match(): model_info = {"litellm_provider": "bedrock"} assert litellm.utils._check_provider_match(model_info, "openai") is False +def test_get_provider_rerank_config(): + """ + Test the get_provider_rerank_config function for various providers + """ + from litellm import HostedVLLMRerankConfig + from litellm.utils import LlmProviders, ProviderConfigManager + + # Test for hosted_vllm provider + config = ProviderConfigManager.get_provider_rerank_config("my_model", LlmProviders.HOSTED_VLLM, 'http://localhost', []) + assert isinstance(config, HostedVLLMRerankConfig) # Models that should be skipped during testing OLD_PROVIDERS = ["aleph_alpha", "palm"]