Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

feat: add NVIDIA NIM support for thinking models (Kimi-K2) #11112

Draft
ghost wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: add NVIDIA NIM support for thinking models (Kimi-K2) #11112

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 47 additions & 4 deletions src/api/providers/openai.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,22 @@
const modelId = this.options.openAiModelId ?? ""
const enabledR1Format = this.options.openAiR1FormatEnabled ?? false
const isAzureAiInference = this._isAzureAiInference(modelUrl)
const deepseekReasoner = modelId.includes("deepseek-reasoner") || enabledR1Format
const isNvidiaNim = this._isNvidiaNim(modelUrl)

// Auto-detect reasoning/thinking models that require R1 format:
// - DeepSeek Reasoner models
// - Models with "thinking" suffix (e.g., kimi-k2-thinking)
// - DeepSeek R1 models (deepseek-r1, deepseek/deepseek-r1)
// - QWQ models
// - Or when user explicitly enables R1 format
const modelIdLower = modelId.toLowerCase()
const useR1Format =
enabledR1Format ||
modelIdLower.includes("deepseek-reasoner") ||
modelIdLower.includes("-thinking") ||
modelIdLower.includes("deepseek-r1") ||
modelIdLower.includes("/deepseek-r1") ||
modelIdLower.includes("qwq")

if (modelId.includes("o1") || modelId.includes("o3") || modelId.includes("o4")) {
yield* this.handleO3FamilyMessage(modelId, systemPrompt, messages, metadata)
Expand All @@ -104,7 +119,7 @@
if (this.options.openAiStreamingEnabled ?? true) {
let convertedMessages

if (deepseekReasoner) {
if (useR1Format) {
convertedMessages = convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
} else {
if (modelInfo.supportsPromptCache) {
Expand Down Expand Up @@ -152,16 +167,25 @@

const isGrokXAI = this._isGrokXAI(this.options.openAiBaseUrl)

// Determine if this is a thinking model that needs NVIDIA NIM specific parameters
const isThinkingModel = modelIdLower.includes("-thinking") || modelIdLower.includes("kimi-k2")

const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
model: modelId,
temperature: this.options.modelTemperature ?? (deepseekReasoner ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
messages: convertedMessages,
stream: true as const,
...(isGrokXAI ? {} : { stream_options: { include_usage: true } }),
...(reasoning && reasoning),
tools: this.convertToolsForOpenAI(metadata?.tools),
tool_choice: metadata?.tool_choice,
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
// Add NVIDIA NIM specific parameters for thinking models
// See: https://build.nvidia.com/moonshotai/kimi-k2-thinking
...(isNvidiaNim && isThinkingModel && {
chat_template_kwargs: { thinking: true },
reasoning_effort: "high",
}),
}
Comment on lines +183 to 189

@ghost ghost Jan 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This block is spread after ...(reasoning && reasoning) on line 179, so if a user has configured a different reasoning effort through their model settings, this hardcoded "high" value will silently override their preference. If NVIDIA NIM only supports "high" for thinking models, consider adding a comment explaining this constraint. Otherwise, consider using the user's configured value: reasoning_effort: reasoningEffort ?? "high".

Fix it with Roo Code or mention @roomote and request a fix.


// Add max_tokens if needed
Expand Down Expand Up @@ -221,15 +245,24 @@
yield this.processUsageMetrics(lastUsage, modelInfo)
}
} else {
// Determine if this is a thinking model that needs NVIDIA NIM specific parameters
const isThinkingModel = modelIdLower.includes("-thinking") || modelIdLower.includes("kimi-k2")

const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
model: modelId,
messages: deepseekReasoner
messages: useR1Format
? convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
: [systemMessage, ...convertToOpenAiMessages(messages)],
// Tools are always present (minimum ALWAYS_AVAILABLE_TOOLS)
tools: this.convertToolsForOpenAI(metadata?.tools),
tool_choice: metadata?.tool_choice,
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
// Add NVIDIA NIM specific parameters for thinking models
// See: https://build.nvidia.com/moonshotai/kimi-k2-thinking
...(isNvidiaNim && isThinkingModel && {
chat_template_kwargs: { thinking: true },
reasoning_effort: "high",
}),
}

// Add max_tokens if needed
Expand Down Expand Up @@ -508,6 +541,16 @@
return urlHost.endsWith(".services.ai.azure.com")
}

/**
* Check if the base URL is NVIDIA NIM API.
* NVIDIA NIM uses integrate.api.nvidia.com for their API endpoint.
* See: https://build.nvidia.com/docs/overview
*/
private _isNvidiaNim(baseUrl?: string): boolean {
const urlHost = this._getUrlHost(baseUrl)
return urlHost.includes("api.nvidia.com") || urlHost.includes("nvidia.com")

Check failure on line 551 in src/api/providers/openai.ts

View check run for this annotation

GitHub Advanced Security / CodeQL

Incomplete URL substring sanitization 

'[api.nvidia.com](1)' can be anywhere in the URL, and arbitrary hosts may come before or after it.

Check failure on line 551 in src/api/providers/openai.ts

View check run for this annotation

GitHub Advanced Security / CodeQL

Incomplete URL substring sanitization 

'[nvidia.com](1)' can be anywhere in the URL, and arbitrary hosts may come before or after it.

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization High

'
api.nvidia.com
Loading
' can be anywhere in the URL, and arbitrary hosts may come before or after it.

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization High

'
nvidia.com
Loading
' can be anywhere in the URL, and arbitrary hosts may come before or after it.
}
Comment on lines +544 to +552

@ghost ghost Jan 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR adds new detection methods and parameters but lacks unit tests. The existing openai.spec.ts file has extensive test coverage for similar functionality (Azure, Grok, O3 models). Consider adding tests for:

  • _isNvidiaNim() detection method
  • NVIDIA NIM specific parameters (chat_template_kwargs, reasoning_effort)
  • isThinkingModel detection logic
  • Extended useR1Format auto-detection

Fix it with Roo Code or mention @roomote and request a fix.


/**
* Adds max_completion_tokens to the request body if needed based on provider configuration
* Note: max_tokens is deprecated in favor of max_completion_tokens as per OpenAI documentation
Expand Down
Loading