From 2ab34f9a5272ad2ef6083aa2a40ca430df125d4e Mon Sep 17 00:00:00 2001 From: Jack Cherng Date: Sat, 8 Nov 2025 11:11:59 +0800 Subject: [PATCH] Fix HostedVLLMRerankConfig will not be used (#16352) * Fix HostedVLLMRerankConfig will not be used Signed-off-by: Jun-Fei Cherng * Fix no usage statistics in rerank with hosted_vllm Signed-off-by: Jun-Fei Cherng * Revise typo in comment Signed-off-by: Jun-Fei Cherng --------- Signed-off-by: Jun-Fei Cherng --- litellm/__init__.py | 1 + litellm/llms/hosted_vllm/rerank/transformation.py | 2 +- litellm/utils.py | 2 ++ tests/test_litellm/test_utils.py | 10 ++++++++++ 4 files changed, 14 insertions(+), 1 deletion(-) 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"]