From f3640447906f07c5bb0dc59ab454b0ac6d98f665 Mon Sep 17 00:00:00 2001 From: atomic <5234009+atomic@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:14:30 -0700 Subject: [PATCH] fix(nvidia_nim): scope image passages to ranking route --- .../nvidia_nim/rerank/ranking_transformation.py | 2 ++ litellm/llms/nvidia_nim/rerank/transformation.py | 15 +++++++-------- tests/llm_translation/test_nvidia_nim.py | 11 +++++++---- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 1b211faf8b3..f3b7ef7d8b0 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -37,6 +37,8 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): }' """ + SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text", "image") + def __init__(self) -> None: super().__init__() # top_n captured in transform_rerank_request and applied in diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 65cb818295a..bb07f9ec74f 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -54,10 +54,9 @@ 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") + # The legacy retrieval rerank route accepts text passages only. The native + # ranking subclass expands this tuple for VL models that accept images. + SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text",) def __init__(self) -> None: pass @@ -212,12 +211,12 @@ class NvidiaNimRerankConfig(BaseRerankConfig): if isinstance(doc, str): passages.append({"text": doc}) elif isinstance(doc, dict): - # Preserve structured passages (text, image, or mixed) so - # VL rerank models receive image passages intact + # Preserve only the structured passage fields supported by the + # selected rerank route. supported_fields: NvidiaNimPassageObject = {} # mutable-ok: assembling a request TypedDict - if "text" in doc: + if "text" in self.SUPPORTED_PASSAGE_FIELDS and "text" in doc: supported_fields["text"] = doc["text"] - if "image" in doc: + if "image" in self.SUPPORTED_PASSAGE_FIELDS and "image" in doc: supported_fields["image"] = doc["image"] if supported_fields: passages.append(supported_fields) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 5448fba4d0b..4bad94e0834 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -525,13 +525,16 @@ class TestNvidiaNimRetrievalRerankRequestTransform: request_data = self._build_request([TEXT_DOC]) assert request_data["passages"] == [TEXT_DOC] - def test_image_object_documents_are_preserved(self): + def test_image_object_documents_keep_retrieval_behavior(self): request_data = self._build_request([IMAGE_DOC, TEXT_DOC]) - assert request_data["passages"] == [IMAGE_DOC, TEXT_DOC] + assert request_data["passages"] == [ + {"text": json.dumps(IMAGE_DOC)}, + TEXT_DOC, + ] - def test_mixed_text_image_documents_are_preserved(self): + def test_mixed_text_image_documents_keep_text_only(self): request_data = self._build_request([MIXED_DOC]) - assert request_data["passages"] == [MIXED_DOC] + assert request_data["passages"] == [{"text": MIXED_DOC["text"]}] def test_unsupported_dict_documents_are_stringified(self): doc = {"title": "no supported fields here"}