scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(nvidia_nim): scope image passages to ranking route

Browse source
This commit is contained in:
atomic 2026-08-07 17:14:30 -07:00
parent 1d7db7564a
commit f364044790
No known key found for this signature in database
GPG key ID: 678EEB06E6D94F5B
3 changed files with 16 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/nvidia_nim/rerank/ranking_transformation.py
View file

@ -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

15
litellm/llms/nvidia_nim/rerank/transformation.py
View file

@ -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)

11
tests/llm_translation/test_nvidia_nim.py
View file

@ -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"}