Fix NVIDIA NIM reasoning params for updated API

Alishahryar1 · Alishahryar1 · commit f9e7f65f4c01 · 2026-03-26T12:25:04.000-07:00
Replace dropped params (thinking, reasoning_split, include_reasoning,
return_tokens_as_token_ids, reasoning_effort) with the new API format:
chat_template_kwargs.enable_thinking=True and reasoning_budget=max_tokens.
diff --git a/config/nim.py b/config/nim.py
@@ -1,7 +1,5 @@
 """NVIDIA NIM settings (fixed values, no env config)."""
 
-from typing import Literal
-
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
@@ -22,16 +20,12 @@ class NimSettings(BaseModel):
     stop: str | None = None
 
     parallel_tool_calls: bool = True
-    return_tokens_as_token_ids: bool = False
     ignore_eos: bool = False
 
     min_tokens: int = Field(0, ge=0)
     chat_template: str | None = None
     request_id: str | None = None
 
-    reasoning_effort: Literal["low", "medium", "high"] = "high"
-    include_reasoning: bool = True
-
     model_config = ConfigDict(extra="forbid")
 
     @field_validator("top_k")
diff --git a/providers/nvidia_nim/request.py b/providers/nvidia_nim/request.py
@@ -63,18 +63,8 @@ def build_request_body(request_data: Any, nim: NimSettings) -> dict:
     if request_extra:
         extra_body.update(request_extra)
 
-    # Handle thinking/reasoning mode
-    extra_body.setdefault("thinking", {"type": "enabled"})
-    extra_body.setdefault("reasoning_split", True)
-    extra_body.setdefault(
-        "chat_template_kwargs",
-        {
-            "thinking": True,
-            "enable_thinking": True,
-            "reasoning_split": True,
-            "clear_thinking": False,
-        },
-    )
+    extra_body.setdefault("chat_template_kwargs", {"enable_thinking": True})
+    _set_extra(extra_body, "reasoning_budget", max_tokens)
 
     req_top_k = getattr(request_data, "top_k", None)
     top_k = req_top_k if req_top_k is not None else nim.top_k
@@ -86,10 +76,7 @@ def build_request_body(request_data: Any, nim: NimSettings) -> dict:
     _set_extra(extra_body, "min_tokens", nim.min_tokens, ignore_value=0)
     _set_extra(extra_body, "chat_template", nim.chat_template)
     _set_extra(extra_body, "request_id", nim.request_id)
-    _set_extra(extra_body, "return_tokens_as_token_ids", nim.return_tokens_as_token_ids)
     _set_extra(extra_body, "ignore_eos", nim.ignore_eos)
-    _set_extra(extra_body, "reasoning_effort", nim.reasoning_effort)
-    _set_extra(extra_body, "include_reasoning", nim.include_reasoning)
 
     if extra_body:
         body["extra_body"] = extra_body
diff --git a/tests/config/test_config.py b/tests/config/test_config.py
@@ -125,11 +125,6 @@ def test_top_p_valid(self, top_p):
         s = NimSettings(top_p=top_p)
         assert s.top_p == top_p
 
-    @pytest.mark.parametrize("effort", ["low", "medium", "high"])
-    def test_reasoning_effort_valid(self, effort):
-        s = NimSettings(reasoning_effort=effort)
-        assert s.reasoning_effort == effort
-
     def test_max_tokens_valid(self):
         s = NimSettings(max_tokens=1)
         assert s.max_tokens == 1
@@ -195,12 +190,6 @@ def test_min_tokens_negative(self):
         with pytest.raises(ValidationError):
             NimSettings(min_tokens=-1)
 
-    def test_reasoning_effort_invalid(self):
-        from typing import Any, cast
-
-        with pytest.raises(ValidationError):
-            NimSettings(reasoning_effort=cast(Any, "invalid"))
-
 
 class TestNimSettingsValidators:
     """Test custom field validators in NimSettings."""
diff --git a/tests/providers/test_nvidia_nim.py b/tests/providers/test_nvidia_nim.py
@@ -99,14 +99,9 @@ async def test_build_request_body(nim_provider):
     assert body["messages"][0]["content"] == "System prompt"
 
     assert "extra_body" in body
-    assert "thinking" in body["extra_body"]
-    assert body["extra_body"]["thinking"]["type"] == "enabled"
-
-    # Verify chat_template_kwargs includes both thinking keys for model compatibility
     ctk = body["extra_body"]["chat_template_kwargs"]
-    assert ctk["thinking"] is True
     assert ctk["enable_thinking"] is True
-    assert ctk["clear_thinking"] is False
+    assert body["extra_body"]["reasoning_budget"] == body["max_tokens"]
 
 
 @pytest.mark.asyncio
diff --git a/tests/providers/test_nvidia_nim_request.py b/tests/providers/test_nvidia_nim_request.py
@@ -83,3 +83,49 @@ def test_parallel_tool_calls_included(self, req):
         nim = NimSettings(parallel_tool_calls=False)
         body = build_request_body(req, nim)
         assert body["parallel_tool_calls"] is False
+
+    def test_reasoning_params_in_extra_body(self):
+        req = MagicMock()
+        req.model = "test"
+        req.messages = [MagicMock(role="user", content="hi")]
+        req.max_tokens = 100
+        req.system = None
+        req.temperature = None
+        req.top_p = None
+        req.stop_sequences = None
+        req.tools = None
+        req.tool_choice = None
+        req.extra_body = None
+        req.top_k = None
+
+        nim = NimSettings()
+        body = build_request_body(req, nim)
+        extra = body["extra_body"]
+        assert extra["chat_template_kwargs"] == {"enable_thinking": True}
+        assert extra["reasoning_budget"] == body["max_tokens"]
+
+    def test_no_reasoning_params_in_extra_body(self):
+        req = MagicMock()
+        req.model = "test"
+        req.messages = [MagicMock(role="user", content="hi")]
+        req.max_tokens = 100
+        req.system = None
+        req.temperature = None
+        req.top_p = None
+        req.stop_sequences = None
+        req.tools = None
+        req.tool_choice = None
+        req.extra_body = None
+        req.top_k = None
+
+        nim = NimSettings()
+        body = build_request_body(req, nim)
+        extra = body.get("extra_body", {})
+        for param in (
+            "thinking",
+            "reasoning_split",
+            "return_tokens_as_token_ids",
+            "include_reasoning",
+            "reasoning_effort",
+        ):
+            assert param not in extra
