scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source 
- Auto-detect reasoning/thinking models that require R1 format
- Add NVIDIA NIM specific parameters (chat_template_kwargs, reasoning_effort)
- Add _isNvidiaNim method to detect NVIDIA NIM API endpoint
- Support models with -thinking suffix (e.g., kimi-k2-thinking)

This validates the user feedback that NVIDIA NIM works correctly with
Kimi-K2 Thinking models when proper parameters are passed.

Related to #10969
This commit is contained in:
Roo Code 2026-01-30 15:25:44 +00:00
parent d8f65b655d
commit 2075e1b9dc
1 changed files with 47 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
src/api/providers/openai.ts
View file

@ -89,7 +89,22 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
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)
@ -104,7 +119,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
if (this.options.openAiStreamingEnabled ?? true) {
let convertedMessages
if (deepseekReasoner) {
if (useR1Format) {
convertedMessages = convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
} else {
if (modelInfo.supportsPromptCache) {
@ -152,9 +167,12 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
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 } }),
@ -162,6 +180,12 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
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
@ -221,15 +245,24 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
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
@ -508,6 +541,16 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
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")
}
/**
* 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