diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 6671ba09a8a..58f26b2eb75 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -6,15 +6,25 @@ Use this by passing "nvidia_nim/ranking/" to force the /v1/ranking endpoi Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy """ -from typing import Final +from typing import Any, Final +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig +from litellm.types.rerank import RerankResponse class NvidiaNimRankingConfig(NvidiaNimRerankConfig): """ Configuration for NVIDIA NIM models that use the /v1/ranking endpoint. - + + The native /v1/ranking request schema accepts only 'model', 'query', + 'passages', and 'truncate' -- requests containing 'top_k' are rejected + with a 400 validation error. Cohere-compatible 'top_n' is therefore + applied client-side by truncating the converted response instead of + being forwarded to the endpoint. + Example: curl -X "POST" 'https://ai.api.nvidia.com/v1/ranking' \ -H 'Accept: application/json' \ @@ -27,6 +37,14 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): }' """ + def __init__(self) -> None: + super().__init__() + # top_n captured in transform_rerank_request and applied in + # transform_rerank_response. The provider config is instantiated + # per-request (see ProviderConfigManager.get_provider_rerank_config), + # so this does not leak across requests. + self._client_side_top_n: int | None = None + def _get_clean_model_name(self, model: str) -> str: """Strip 'nvidia_nim/' and 'ranking/' prefixes from model name.""" # First strip nvidia_nim/ prefix if present @@ -58,6 +76,47 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): return f"{api_base}/v1/ranking" + def map_cohere_rerank_params( + self, + non_default_params: dict | None, + model: str, + drop_params: bool, + query: str, + documents: list[str | dict[str, Any]], + custom_llm_provider: str | None = None, + top_n: int | None = None, + rank_fields: list[str] | None = None, + return_documents: bool | None = True, + max_chunks_per_doc: int | None = None, + max_tokens_per_doc: int | None = None, + instruction: str | None = None, + ) -> dict: + """ + Keep Cohere's top_n as-is instead of mapping it to top_k. + + The native /v1/ranking endpoint rejects top_k, so top_n is applied + client-side after the response is converted. + """ + optional_params = super().map_cohere_rerank_params( + non_default_params=non_default_params, + model=model, + drop_params=drop_params, + query=query, + documents=documents, + custom_llm_provider=custom_llm_provider, + top_n=None, # do not map top_n -> top_k for /v1/ranking + rank_fields=rank_fields, + return_documents=return_documents, + max_chunks_per_doc=max_chunks_per_doc, + max_tokens_per_doc=max_tokens_per_doc, + instruction=instruction, + ) + # /v1/ranking rejects top_k even when passed as a provider-specific param + optional_params.pop("top_k", None) + if top_n is not None: + optional_params["top_n"] = top_n + return optional_params + def transform_rerank_request( self, model: str, @@ -67,11 +126,62 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): ) -> dict: """ Transform request, using clean model name without 'ranking/' prefix. + + top_n / top_k are stripped from the outgoing request: the native + /v1/ranking endpoint accepts only model, query, passages, and + truncate. top_n is stashed and applied client-side in + transform_rerank_response. """ + top_n = optional_rerank_params.get("top_n") + if top_n is not None: + if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1: + raise ValueError(f"top_n must be a positive integer, got: {top_n!r}") + self._client_side_top_n = top_n + clean_model: Final = self._get_clean_model_name(model) + filtered_params: Final = { + k: v for k, v in optional_rerank_params.items() if k not in ("top_n", "top_k") + } return super().transform_rerank_request( model=clean_model, - optional_rerank_params=optional_rerank_params, + optional_rerank_params=filtered_params, headers=headers, litellm_params=litellm_params, ) + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: str | None = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> RerankResponse: + """ + Convert the native ranking response, then apply top_n client-side. + + /v1/ranking returns rankings sorted by relevance, but sort before + truncating in case a server returns them unsorted. + """ + response = super().transform_rerank_response( + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + request_data=request_data, + optional_params=optional_params, + litellm_params=litellm_params, + ) + + top_n = optional_params.get("top_n") or self._client_side_top_n + if top_n is not None and response.results is not None and len(response.results) > top_n: + response.results = sorted( + response.results, + key=lambda result: result["relevance_score"], + reverse=True, + )[:top_n] + return response diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index aeb1190d0a5..7d4ecdb4cdc 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -21,8 +21,9 @@ class NvidiaNimQueryObject(TypedDict): text: Required[str] -class NvidiaNimPassageObject(TypedDict): - text: Required[str] +class NvidiaNimPassageObject(TypedDict, total=False): + text: str + image: str class NvidiaNimRerankRequest(TypedDict, total=False): @@ -53,6 +54,11 @@ class NvidiaNimRerankConfig(BaseRerankConfig): DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com" + # Structured document fields forwarded to the ranking API as-is. + # VL rerank models (e.g. nvidia/llama-nemotron-rerank-vl-1b-v2) accept + # image passages alongside text passages. + SUPPORTED_PASSAGE_FIELDS = ("text", "image") + def __init__(self) -> None: pass @@ -206,11 +212,17 @@ class NvidiaNimRerankConfig(BaseRerankConfig): if isinstance(doc, str): passages.append({"text": doc}) elif isinstance(doc, dict): - # If document is already a dict, check if it has 'text' field - if "text" in doc: - passages.append({"text": doc["text"]}) + # Preserve structured passages (text, image, or mixed) so + # VL rerank models receive image passages intact + supported_fields: NvidiaNimPassageObject = { + field: doc[field] # type: ignore[misc] + for field in self.SUPPORTED_PASSAGE_FIELDS + if field in doc + } + if supported_fields: + passages.append(supported_fields) else: - # Otherwise, stringify the dict + # No supported fields - stringify the dict import json passages.append({"text": json.dumps(doc)}) @@ -304,9 +316,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig): "relevance_score": ranking["logit"], } - # Include document if it was in the original request + # Include document if it was in the original request. + # Image-only passages carry no 'text' field, so guard the lookup. index: int = ranking["index"] - if index < len(original_passages): + if index < len(original_passages) and "text" in original_passages[index]: result_item["document"] = {"text": original_passages[index]["text"]} results.append(result_item) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 80e764147bb..60d2a960d5b 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -303,3 +303,183 @@ class TestNvidiaNim(BaseLLMRerankTest): ), ): await super().test_basic_rerank(sync_mode=sync_mode) + + +# --------------------------------------------------------------------------- +# Regression tests for https://github.com/BerriAI/litellm/issues/34165 +# +# The native /v1/ranking endpoint accepts only model, query, passages, and +# truncate. Two defects are covered here: +# 1. structured image documents were json.dumps-stringified into text passages +# 2. Cohere top_n was mapped to top_k, which /v1/ranking rejects with a 400 +# --------------------------------------------------------------------------- + +from litellm.llms.nvidia_nim.rerank.ranking_transformation import ( + NvidiaNimRankingConfig, +) +from litellm.types.rerank import RerankResponse + +RANKING_MODEL = "ranking/nvidia/llama-nemotron-rerank-vl-1b-v2" +IMAGE_DOC = {"image": "data:image/jpeg;base64,/9j/4AAQSkZJRg=="} +TEXT_DOC = {"text": "a plain text passage"} +MIXED_DOC = {"text": "caption for the image", "image": "data:image/png;base64,iVBORw0KGgo="} + + +def _build_ranking_request(documents, top_n=None, non_default_params=None): + """Run map_cohere_rerank_params + transform_rerank_request for /v1/ranking.""" + config = NvidiaNimRankingConfig() + optional_params = config.map_cohere_rerank_params( + non_default_params=non_default_params, + model=RANKING_MODEL, + drop_params=False, + query="which passage shows a cat?", + documents=documents, + top_n=top_n, + ) + request_data = config.transform_rerank_request( + model=RANKING_MODEL, + optional_rerank_params=optional_params, + headers={}, + ) + return config, request_data + + +def _build_ranking_response(config, request_data, rankings): + """Run transform_rerank_response against a mocked raw ranking response.""" + raw_response = MagicMock() + raw_response.json.return_value = {"rankings": rankings} + return config.transform_rerank_response( + model=RANKING_MODEL, + raw_response=raw_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + + +class TestNvidiaNimRankingRequestTransform: + def test_string_documents(self): + _, request_data = _build_ranking_request(["passage one", "passage two"]) + assert request_data["passages"] == [ + {"text": "passage one"}, + {"text": "passage two"}, + ] + + def test_text_object_documents(self): + _, request_data = _build_ranking_request([TEXT_DOC]) + assert request_data["passages"] == [TEXT_DOC] + + def test_image_object_documents_are_preserved(self): + _, request_data = _build_ranking_request([IMAGE_DOC, TEXT_DOC]) + assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + + def test_mixed_text_image_documents_are_preserved(self): + _, request_data = _build_ranking_request([MIXED_DOC]) + assert request_data["passages"] == [MIXED_DOC] + + def test_unsupported_dict_documents_are_stringified(self): + doc = {"title": "no supported fields here"} + _, request_data = _build_ranking_request([doc]) + assert request_data["passages"] == [{"text": json.dumps(doc)}] + + def test_top_n_is_not_sent_to_the_ranking_endpoint(self): + _, request_data = _build_ranking_request(["a", "b"], top_n=1) + assert "top_k" not in request_data + assert "top_n" not in request_data + + def test_provider_specific_top_k_is_stripped(self): + _, request_data = _build_ranking_request(["a", "b"], non_default_params={"top_k": 2}) + assert "top_k" not in request_data + + @pytest.mark.parametrize("invalid_top_n", [0, -1, 1.5, "2", True]) + def test_invalid_top_n_raises_value_error(self, invalid_top_n): + with pytest.raises(ValueError, match="top_n"): + _build_ranking_request(["a", "b"], top_n=invalid_top_n) + + +class TestNvidiaNimRankingResponseTransform: + RANKINGS = [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.75}, + {"index": 2, "logit": 0.55}, + ] + + def test_top_n_one_truncates_to_best_result(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=1) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 1 + assert response.results[0]["index"] == 0 + + def test_top_n_equal_to_document_count_keeps_all_results(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=3) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 3 + + def test_top_n_greater_than_document_count_keeps_all_results(self): + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=10) + response = _build_ranking_response(config, request_data, self.RANKINGS) + assert len(response.results) == 3 + + def test_top_n_truncation_keeps_most_relevant_results(self): + unsorted_rankings = [ + {"index": 0, "logit": 0.10}, + {"index": 1, "logit": 0.90}, + {"index": 2, "logit": 0.50}, + ] + config, request_data = _build_ranking_request(["a", "b", "c"], top_n=2) + response = _build_ranking_response(config, request_data, unsorted_rankings) + assert [result["index"] for result in response.results] == [1, 2] + + def test_image_only_passages_do_not_break_document_echo(self): + config, request_data = _build_ranking_request([IMAGE_DOC, TEXT_DOC]) + response = _build_ranking_response(config, request_data, self.RANKINGS[:2]) + assert len(response.results) == 2 + # Image-only passage has no text to echo back + assert "document" not in response.results[0] + assert response.results[1]["document"] == {"text": TEXT_DOC["text"]} + + +@pytest.mark.asyncio() +async def test_nvidia_nim_ranking_endpoint_image_documents_and_top_n(): + """ + End-to-end (mocked transport): image documents reach /v1/ranking intact + and top_n is applied client-side instead of being sent as top_k. + """ + mock_response = AsyncMock() + + def return_val(): + return { + "rankings": [ + {"index": 0, "logit": 0.95}, + {"index": 1, "logit": 0.75}, + ], + } + + mock_response.json = return_val + mock_response.headers = {"key": "value"} + mock_response.status_code = 200 + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=mock_response, + ) as mock_post: + response = await litellm.arerank( + model="nvidia_nim/ranking/nvidia/llama-nemotron-rerank-vl-1b-v2", + query="which passage shows a cat?", + documents=[IMAGE_DOC, TEXT_DOC], + top_n=1, + api_key="fake-api-key", + ) + + mock_post.assert_called_once() + request_data = json.loads(mock_post.call_args.kwargs["data"]) + + assert mock_post.call_args.kwargs["url"] == "https://ai.api.nvidia.com/v1/ranking" + # Image passage preserved as-is, not stringified into text + assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + # Neither top_k nor top_n is sent to the native endpoint + assert "top_k" not in request_data + assert "top_n" not in request_data + # top_n applied client-side on the converted response + assert len(response.results) == 1 + assert response.results[0]["index"] == 0