diff --git a/litellm/llms/nvidia_nim/chat/transformation.py b/litellm/llms/nvidia_nim/chat/transformation.py index 8993c0e6b15..a2e807a0d18 100644 --- a/litellm/llms/nvidia_nim/chat/transformation.py +++ b/litellm/llms/nvidia_nim/chat/transformation.py @@ -10,6 +10,7 @@ API calling is done using the OpenAI SDK with an api_base from typing import Final +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig @@ -107,9 +108,13 @@ class NvidiaNimConfig(OpenAIGPTConfig): drop_params: bool, ) -> dict: supported_openai_params: Final = self.get_supported_openai_params(model=model) + asks_logprobs: Final = "logprobs" in non_default_params and non_default_params["logprobs"] is True + drops_top_logprobs: Final = not asks_logprobs and drops_unsupported_params(drop_params) for param, value in non_default_params.items(): if param == "max_completion_tokens": optional_params["max_tokens"] = value + elif param == "top_logprobs" and drops_top_logprobs: + continue elif param in supported_openai_params: optional_params[param] = value return optional_params diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index c1bbbda8794..f5937947f2a 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -93,9 +93,7 @@ def test_chat_completion_nvidia_nim_with_tools(): }, ] - with patch.object( - client, "post" - ) as mock_client: + with patch.object(client, "post") as mock_client: try: completion( model=model_name, @@ -284,3 +282,56 @@ def test_completion_nvidia_nim_forwards_logprobs_and_returns_them(): assert isinstance(response, ModelResponse) assert response.choices[0].logprobs.content[0].token == token_logprob.token assert response.choices[0].logprobs.content[0].logprob == token_logprob.logprob + + +@pytest.mark.parametrize( + ("logprobs", "drop_params", "global_drop_params", "expected_body"), + [ + (None, True, False, {}), + (None, None, True, {}), + (False, True, False, {"logprobs": False}), + (None, None, False, {"top_logprobs": 2}), + (True, True, False, {"logprobs": True, "top_logprobs": 2}), + ], +) +def test_completion_nvidia_nim_drops_top_logprobs_without_logprobs_only_under_drop_params( + monkeypatch: pytest.MonkeyPatch, + logprobs: bool | None, + drop_params: bool | None, + global_drop_params: bool, + expected_body: dict[str, bool | int], +): + from openai import OpenAI + from openai.types.chat import ChatCompletion + from openai.types.chat.chat_completion import Choice + from openai.types.chat.chat_completion_message import ChatCompletionMessage + + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + transport: Final = CapturingTransport( + ChatCompletion( + id="chatcmpl-nim", + object="chat.completion", + created=1, + model="openai/gpt-oss-20b", + choices=[ + Choice( + index=0, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="hi"), + ) + ], + ) + ) + client: Final = OpenAI(api_key="fake-api-key", http_client=httpx.Client(transport=transport)) + + completion( + model="nvidia_nim/openai/gpt-oss-20b", + messages=[{"role": "user", "content": "Say hi"}], + logprobs=logprobs, + top_logprobs=2, + drop_params=drop_params, + client=client, + ) + + request_body: Final = transport.request_bodies[0] + assert {key: request_body[key] for key in ("logprobs", "top_logprobs") if key in request_body} == expected_body