feat: add support for Kimi K2 Thinking model embedded tool calls

This adds support for extracting tool calls from Kimi K2 Thinking model
responses when using the OpenAI Compatible provider. The Kimi K2 Thinking
model embeds tool calls in the reasoning_content field using special tags
rather than the standard tool_calls field.

Changes:
- Added kimi-tool-call-extractor utility to parse embedded tool calls
- Updated OpenAI handler to detect and extract embedded tool calls
- Added comprehensive tests for the new functionality

Fixes #10064
This commit is contained in:
Roo Code 2025-12-13 08:58:19 +00:00
parent a3b258ad62
commit 464a7b5600
3 changed files with 373 additions and 2 deletions

View file

@ -25,6 +25,7 @@ import { BaseProvider } from "./base-provider"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
import { getApiRequestTimeout } from "./utils/timeout-config"
import { handleOpenAIError } from "./utils/openai-error-handler"
import { hasKimiEmbeddedToolCalls, extractKimiToolCalls, isKimiThinkingModel } from "./utils/kimi-tool-call-extractor"
// TODO: Rename this to OpenAICompatibleHandler. Also, I think the
// `OpenAINativeHandler` can subclass from this, since it's obviously
@ -195,6 +196,11 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
let lastUsage
// For Kimi K2 Thinking model, accumulate reasoning content to extract embedded tool calls
const isKimiThinking = isKimiThinkingModel(modelId)
let accumulatedReasoningContent = ""
let hasReceivedToolCalls = false
for await (const chunk of stream) {
const delta = chunk.choices?.[0]?.delta ?? {}
@ -205,13 +211,21 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
}
if ("reasoning_content" in delta && delta.reasoning_content) {
const reasoningText = (delta.reasoning_content as string | undefined) || ""
// Accumulate reasoning content for Kimi thinking model to extract tool calls later
if (isKimiThinking) {
accumulatedReasoningContent += reasoningText
}
yield {
type: "reasoning",
text: (delta.reasoning_content as string | undefined) || "",
text: reasoningText,
}
}
if (delta.tool_calls) {
hasReceivedToolCalls = true
for (const toolCall of delta.tool_calls) {
yield {
type: "tool_call_partial",
@ -228,6 +242,24 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
}
}
// For Kimi K2 Thinking model: extract tool calls from accumulated reasoning content
// if no standard tool_calls were received but reasoning contains embedded tool calls
if (isKimiThinking && !hasReceivedToolCalls && hasKimiEmbeddedToolCalls(accumulatedReasoningContent)) {
const { toolCalls } = extractKimiToolCalls(accumulatedReasoningContent)
for (let index = 0; index < toolCalls.length; index++) {
const toolCall = toolCalls[index]
// Emit as complete tool call since we have all the data
yield {
type: "tool_call_partial",
index,
id: toolCall.id,
name: toolCall.function.name,
arguments: toolCall.function.arguments,
}
}
}
for (const chunk of matcher.final()) {
yield chunk
}
@ -265,7 +297,20 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
const message = response.choices?.[0]?.message
if (message?.tool_calls) {
// Check for Kimi K2 Thinking model embedded tool calls in reasoning_content
const isKimiThinking = isKimiThinkingModel(modelId)
const reasoningContent = (message as any)?.reasoning_content as string | undefined
// Emit reasoning content if present
if (reasoningContent) {
yield {
type: "reasoning",
text: reasoningContent,
}
}
// Handle standard tool calls or extract from Kimi embedded format
if (message?.tool_calls && message.tool_calls.length > 0) {
for (const toolCall of message.tool_calls) {
if (toolCall.type === "function") {
yield {
@ -276,6 +321,18 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
}
}
}
} else if (isKimiThinking && reasoningContent && hasKimiEmbeddedToolCalls(reasoningContent)) {
// Extract embedded tool calls from Kimi K2 Thinking model's reasoning_content
const { toolCalls } = extractKimiToolCalls(reasoningContent)
for (const toolCall of toolCalls) {
yield {
type: "tool_call",
id: toolCall.id,
name: toolCall.function.name,
arguments: toolCall.function.arguments,
}
}
}
yield {

View file

@ -0,0 +1,182 @@
import {
hasKimiEmbeddedToolCalls,
extractKimiToolCalls,
isKimiThinkingModel,
type KimiToolCall,
} from "../kimi-tool-call-extractor"
describe("kimi-tool-call-extractor", () => {
describe("hasKimiEmbeddedToolCalls", () => {
it("should return true when content contains tool call markers", () => {
const content = "Some reasoning <|tool_calls_section_begin|> stuff <|tool_calls_section_end|>"
expect(hasKimiEmbeddedToolCalls(content)).toBe(true)
})
it("should return false when content does not contain tool call markers", () => {
const content = "Just regular reasoning content without any tool calls"
expect(hasKimiEmbeddedToolCalls(content)).toBe(false)
})
it("should return false for empty string", () => {
expect(hasKimiEmbeddedToolCalls("")).toBe(false)
})
})
describe("extractKimiToolCalls", () => {
it("should extract single tool call from reasoning content", () => {
const content = `Some reasoning here
<|tool_calls_section_begin|>
<|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>{"files":[{"path":"test.txt"}]}<|tool_call_end|>
<|tool_calls_section_end|>
More content after`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(1)
expect(result.toolCalls[0]).toEqual({
id: "kimi-functions.read_file:0",
type: "function",
function: {
name: "read_file",
arguments: '{"files":[{"path":"test.txt"}]}',
},
})
expect(result.cleanedReasoningContent).not.toContain("<|tool_calls_section_begin|>")
expect(result.cleanedReasoningContent).toContain("Some reasoning here")
expect(result.cleanedReasoningContent).toContain("More content after")
})
it("should extract multiple tool calls from reasoning content", () => {
const content = `Thinking about what to do
<|tool_calls_section_begin|>
<|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>{"files":[{"path":"file1.txt"}]}<|tool_call_end|>
<|tool_call_begin|>functions.execute_command:1<|tool_call_argument_begin|>{"command":"ls -la"}<|tool_call_end|>
<|tool_calls_section_end|>`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(2)
expect(result.toolCalls[0]).toEqual({
id: "kimi-functions.read_file:0",
type: "function",
function: {
name: "read_file",
arguments: '{"files":[{"path":"file1.txt"}]}',
},
})
expect(result.toolCalls[1]).toEqual({
id: "kimi-functions.execute_command:1",
type: "function",
function: {
name: "execute_command",
arguments: '{"command":"ls -la"}',
},
})
})
it("should handle tool calls without functions. prefix", () => {
const content = `<|tool_calls_section_begin|>
<|tool_call_begin|>read_file:0<|tool_call_argument_begin|>{"files":[{"path":"test.txt"}]}<|tool_call_end|>
<|tool_calls_section_end|>`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(1)
expect(result.toolCalls[0].function.name).toBe("read_file")
})
it("should return empty array when no tool calls are present", () => {
const content = "Just regular content without tool calls"
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(0)
expect(result.cleanedReasoningContent).toBe(content)
})
it("should return empty array when tool call section is empty", () => {
const content = `Some content
<|tool_calls_section_begin|>
<|tool_calls_section_end|>`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(0)
})
it("should clean reasoning content by removing tool call sections", () => {
const content =
"Before tool calls\n<|tool_calls_section_begin|>\n<|tool_call_begin|>functions.test:0<|tool_call_argument_begin|>{}<|tool_call_end|>\n<|tool_calls_section_end|>\nAfter tool calls"
const result = extractKimiToolCalls(content)
// The cleaned content should not contain tool call markers
expect(result.cleanedReasoningContent).not.toContain("<|tool_calls_section_begin|>")
expect(result.cleanedReasoningContent).not.toContain("<|tool_calls_section_end|>")
expect(result.cleanedReasoningContent).not.toContain("<|tool_call_begin|>")
expect(result.cleanedReasoningContent).toContain("Before tool calls")
expect(result.cleanedReasoningContent).toContain("After tool calls")
})
it("should handle complex JSON arguments", () => {
const content = `<|tool_calls_section_begin|>
<|tool_call_begin|>functions.write_to_file:0<|tool_call_argument_begin|>{"path":"src/test.ts","content":"function test() {\\n return 'hello';\\n}"}<|tool_call_end|>
<|tool_calls_section_end|>`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(1)
expect(result.toolCalls[0].function.arguments).toBe(
'{"path":"src/test.ts","content":"function test() {\\n return \'hello\';\\n}"}',
)
})
it("should handle multiple tool call sections", () => {
const content = `First reasoning
<|tool_calls_section_begin|>
<|tool_call_begin|>functions.tool1:0<|tool_call_argument_begin|>{"arg":"value1"}<|tool_call_end|>
<|tool_calls_section_end|>
Middle reasoning
<|tool_calls_section_begin|>
<|tool_call_begin|>functions.tool2:1<|tool_call_argument_begin|>{"arg":"value2"}<|tool_call_end|>
<|tool_calls_section_end|>
End reasoning`
const result = extractKimiToolCalls(content)
expect(result.toolCalls).toHaveLength(2)
expect(result.toolCalls[0].function.name).toBe("tool1")
expect(result.toolCalls[1].function.name).toBe("tool2")
})
})
describe("isKimiThinkingModel", () => {
it("should return true for kimi-k2-thinking model", () => {
expect(isKimiThinkingModel("kimi-k2-thinking")).toBe(true)
})
it("should return true for model with kimi-k2-thinking prefix", () => {
expect(isKimiThinkingModel("moonshotai/kimi-k2-thinking")).toBe(true)
})
it("should be case insensitive", () => {
expect(isKimiThinkingModel("Kimi-K2-Thinking")).toBe(true)
expect(isKimiThinkingModel("KIMI-K2-THINKING")).toBe(true)
})
it("should handle underscore variations", () => {
expect(isKimiThinkingModel("kimi_k2_thinking")).toBe(true)
})
it("should return false for non-thinking kimi models", () => {
expect(isKimiThinkingModel("kimi-k2")).toBe(false)
expect(isKimiThinkingModel("kimi-k2-0905-preview")).toBe(false)
expect(isKimiThinkingModel("kimi-k2-turbo-preview")).toBe(false)
})
it("should return false for other models", () => {
expect(isKimiThinkingModel("gpt-4")).toBe(false)
expect(isKimiThinkingModel("claude-3-opus")).toBe(false)
expect(isKimiThinkingModel("deepseek-reasoner")).toBe(false)
})
})
})

View file

@ -0,0 +1,132 @@
/**
* Extracts tool calls from Kimi K2 Thinking model's reasoning_content field.
*
* Kimi K2 Thinking model embeds tool calls in reasoning_content using special tags:
* - <|tool_calls_section_begin|> ... <|tool_calls_section_end|> wraps all tool calls
* - <|tool_call_begin|> ... <|tool_call_end|> wraps each individual tool call
* - <|tool_call_argument_begin|> marks the start of arguments JSON
*
* Format example:
* <|tool_calls_section_begin|>
* <|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>{"files":[{"path":"test.txt"}]}<|tool_call_end|>
* <|tool_calls_section_end|>
*
* @see https://huggingface.co/moonshotai/Kimi-K2-Thinking/blob/main/docs/tool_call_guidance.md
*/
/**
* Represents an extracted tool call from Kimi K2 Thinking model's reasoning_content.
*/
export interface KimiToolCall {
id: string
type: "function"
function: {
name: string
arguments: string
}
}
/**
* Result of extracting tool calls and cleaning reasoning content.
*/
export interface KimiToolCallExtractionResult {
toolCalls: KimiToolCall[]
cleanedReasoningContent: string
}
/**
* Checks if the content contains Kimi K2 Thinking model's embedded tool call markers.
*/
export function hasKimiEmbeddedToolCalls(content: string): boolean {
return content.includes("<|tool_calls_section_begin|>")
}
/**
* Extracts tool calls from Kimi K2 Thinking model's reasoning_content.
*
* @param content - The reasoning_content or combined content to extract tool calls from
* @returns An object containing the extracted tool calls and the cleaned reasoning content
*/
export function extractKimiToolCalls(content: string): KimiToolCallExtractionResult {
if (!hasKimiEmbeddedToolCalls(content)) {
return {
toolCalls: [],
cleanedReasoningContent: content,
}
}
const toolCalls: KimiToolCall[] = []
// Pattern to match tool call sections
const sectionPattern = /<\|tool_calls_section_begin\|>(.*?)<\|tool_calls_section_end\|>/gs
const toolCallSections = content.match(sectionPattern)
if (!toolCallSections || toolCallSections.length === 0) {
return {
toolCalls: [],
cleanedReasoningContent: content,
}
}
// Pattern to extract individual tool calls
// Format: <|tool_call_begin|>functions.tool_name:index<|tool_call_argument_begin|>JSON_ARGS<|tool_call_end|>
const funcCallPattern =
/<\|tool_call_begin\|>\s*([\w.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(.*?)\s*<\|tool_call_end\|>/gs
for (const section of toolCallSections) {
let match
// Reset lastIndex for each section to ensure all matches are found
funcCallPattern.lastIndex = 0
while ((match = funcCallPattern.exec(section)) !== null) {
const [, functionId, functionArgs] = match
// functionId format: functions.tool_name:index (e.g., "functions.read_file:0")
// We need to extract just the tool name
let functionName = functionId
// Handle "functions.tool_name:index" format
if (functionId.includes(".")) {
const parts = functionId.split(".")
// Get the part after the last dot, then remove the :index suffix
const nameWithIndex = parts[parts.length - 1]
functionName = nameWithIndex.split(":")[0]
} else if (functionId.includes(":")) {
// Handle "tool_name:index" format (without functions. prefix)
functionName = functionId.split(":")[0]
}
toolCalls.push({
id: `kimi-${functionId}`,
type: "function",
function: {
name: functionName,
arguments: functionArgs.trim(),
},
})
}
}
// Clean the reasoning content by removing tool call sections
const cleanedReasoningContent = content.replace(sectionPattern, "").trim()
return {
toolCalls,
cleanedReasoningContent,
}
}
/**
* Checks if a model ID corresponds to a Kimi K2 Thinking model that uses embedded tool calls.
* This includes various forms of the model name that users might use.
*/
export function isKimiThinkingModel(modelId: string): boolean {
const normalizedModelId = modelId.toLowerCase()
// Match various forms of kimi-k2-thinking model name
return (
normalizedModelId.includes("kimi-k2-thinking") ||
normalizedModelId.includes("kimi_k2_thinking") ||
normalizedModelId.includes("kimik2thinking")
)
}