From 8f2108d222a901865b56c245af4e8a293867f438 Mon Sep 17 00:00:00 2001 From: Jim Smith Date: Thu, 18 Jun 2026 11:10:37 -0400 Subject: [PATCH] Address review: thread `instruction` as a typed param + cover rerank_utils Per PR review (greptile P2 + codecov): - Make `instruction` a typed, named argument on the rerank provider interface instead of recovering it from the opaque `non_default_params` blob. Adds `instruction: Optional[str] = None` to `BaseRerankConfig.map_cohere_rerank_params` and every provider override, and forwards it explicitly from `get_optional_rerank_params`. hosted_vllm now reads the named param directly. It is still also surfaced in `non_default_params` so providers that read it there (e.g. DeepInfra) keep working now that `rerank()` consumes `instruction` as a named param rather than leaving it in **kwargs. - Add get_optional_rerank_params unit tests (present + absent) to cover the previously-uncovered threading line flagged by codecov. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../llms/base_llm/rerank/transformation.py | 1 + litellm/llms/cohere/rerank/transformation.py | 1 + .../llms/cohere/rerank_v2/transformation.py | 1 + .../llms/dashscope/rerank/transformation.py | 1 + .../llms/deepinfra/rerank/transformation.py | 1 + .../fireworks_ai/rerank/transformation.py | 1 + .../llms/hosted_vllm/rerank/transformation.py | 6 ++-- .../llms/huggingface/rerank/transformation.py | 1 + litellm/llms/jina_ai/rerank/transformation.py | 1 + .../llms/nvidia_nim/rerank/transformation.py | 1 + .../llms/vertex_ai/rerank/transformation.py | 1 + litellm/llms/voyage/rerank/transformation.py | 1 + litellm/llms/watsonx/rerank/transformation.py | 1 + litellm/rerank_api/rerank_utils.py | 4 +++ .../test_hosted_vllm_rerank_transformation.py | 33 ++++++++++++++++++- 15 files changed, 51 insertions(+), 4 deletions(-) diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 166f876ba04..6603c64142b 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -85,6 +85,7 @@ class BaseRerankConfig(ABC): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: pass diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 64ae8e8ffa7..d875f420310 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -57,6 +57,7 @@ class CohereRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map Cohere rerank params diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 4c800d6455d..0dcb10d5664 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -49,6 +49,7 @@ class CohereRerankV2Config(CohereRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map Cohere rerank params diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 629f3cf4af7..745e85de7e3 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -116,6 +116,7 @@ class DashScopeRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: # qwen3-rerank accepts query/documents/top_n/return_documents. The # rest (rank_fields, max_*_per_doc) are silently dropped. diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index e4bfbcb2513..a5c36ca2e5f 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -104,6 +104,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: # Start with the basic parameters optional_rerank_params = {} diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 4a7b64b9b77..27309780c86 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -67,6 +67,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict[str, Any]: """ Map Cohere rerank params to Fireworks AI rerank params diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index efdff780254..d0c96f8b420 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -77,6 +77,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map parameters for Hosted VLLM rerank @@ -93,9 +94,8 @@ class HostedVLLMRerankConfig(BaseRerankConfig): ) # `instruction` is a vLLM-supported passthrough (folded into the model's - # chat_template_kwargs). It arrives via non_default_params; only forward - # it when explicitly set so omitting it leaves the request unchanged. - instruction = (non_default_params or {}).get("instruction") + # chat_template_kwargs). Only forward it when explicitly set so omitting + # it leaves the request unchanged. if instruction is not None: mapped_params["instruction"] = instruction diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 2c847b617ef..4e409f31ed2 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -100,6 +100,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: optional_rerank_params = {} if non_default_params is not None: diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 56be754fc34..0d48ed5edcd 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -45,6 +45,7 @@ class JinaAIRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: optional_params = {} supported_params = self.get_supported_cohere_rerank_params(model) diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index fc317293acc..8eee188bf46 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -117,6 +117,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map Cohere/OpenAI rerank params to Nvidia NIM format. diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 3b84972e946..d2041009efb 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -242,6 +242,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map Cohere rerank params to Vertex AI format diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index d64450a1211..907e5b7e26b 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -39,6 +39,7 @@ class VoyageRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: # Voyage AI uses 'top_k' instead of 'top_n' optional_params: Dict[str, Any] = {"query": query, "documents": documents} diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 202760f68a6..a34358a6be3 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -104,6 +104,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, + instruction: Optional[str] = None, ) -> Dict: """ Map Cohere rerank params to IBM watsonx.ai rerank params diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py index 32a1a34e343..a8a665496fc 100644 --- a/litellm/rerank_api/rerank_utils.py +++ b/litellm/rerank_api/rerank_utils.py @@ -32,6 +32,9 @@ def get_optional_rerank_params( if max_tokens_per_doc is not None: all_non_default_params["max_tokens_per_doc"] = max_tokens_per_doc if instruction is not None: + # Also surfaced in non_default_params so providers that read it from + # there (e.g. DeepInfra) keep working now that `rerank()` consumes + # `instruction` as a named param instead of leaving it in **kwargs. all_non_default_params["instruction"] = instruction return rerank_provider_config.map_cohere_rerank_params( model=model, @@ -44,5 +47,6 @@ def get_optional_rerank_params( return_documents=return_documents, max_chunks_per_doc=max_chunks_per_doc, max_tokens_per_doc=max_tokens_per_doc, + instruction=instruction, non_default_params=all_non_default_params, ) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index da60ea24d6c..983324adc02 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -4,6 +4,7 @@ import sys import pytest from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig +from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( OptionalRerankParams, RerankBilledUnits, @@ -51,11 +52,12 @@ class TestHostedVLLMRerankTransform: def test_map_cohere_rerank_params_passes_instruction_when_set(self): params = self.config.map_cohere_rerank_params( - non_default_params={"instruction": "Rank by relevance to genomics"}, + non_default_params=None, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], + instruction="Rank by relevance to genomics", ) assert params["instruction"] == "Rank by relevance to genomics" @@ -141,3 +143,32 @@ class TestHostedVLLMRerankTransform: } with pytest.raises(ValueError, match="Missing required fields in the result="): self.config._transform_response(response_dict) + + +class TestGetOptionalRerankParamsInstruction: + """`instruction` is threaded through get_optional_rerank_params only when set.""" + + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_instruction_threaded_when_set(self): + params = get_optional_rerank_params( + rerank_provider_config=self.config, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + instruction="Rank by relevance to genomics", + ) + assert params["instruction"] == "Rank by relevance to genomics" + + def test_instruction_absent_when_not_set(self): + params = get_optional_rerank_params( + rerank_provider_config=self.config, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert "instruction" not in params