mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(nvidia_nim): preserve image passages and stop sending top_k to /v1/ranking
The NVIDIA NIM native /v1/ranking endpoint accepts only model, query, passages, and truncate. The rerank transform stringified structured image documents into text passages, so VL rerank models scored serialized JSON instead of the image, and it mapped Cohere top_n to top_k, which /v1/ranking rejects with a 400 validation error. - preserve structured documents (text, image, mixed) as passages - for nvidia_nim/ranking/ models, keep top_n out of the provider request and truncate the converted response client-side - guard the response document echo for image-only passages Fixes #34165
This commit is contained in:
atomic
parent
8db2fbaad0
commit
108b0f935a
3 changed files with 314 additions and 11 deletions
|
|
@ -6,15 +6,25 @@ Use this by passing "nvidia_nim/ranking/<model>" 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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue