From 464a7b5600ffdfc4e851b54f2e1f2e62758753ee Mon Sep 17 00:00:00 2001 From: Roo Code Date: Sat, 13 Dec 2025 08:58:19 +0000 Subject: [PATCH] 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 --- src/api/providers/openai.ts | 61 +++++- .../kimi-tool-call-extractor.spec.ts | 182 ++++++++++++++++++ .../utils/kimi-tool-call-extractor.ts | 132 +++++++++++++ 3 files changed, 373 insertions(+), 2 deletions(-) create mode 100644 src/api/providers/utils/__tests__/kimi-tool-call-extractor.spec.ts create mode 100644 src/api/providers/utils/kimi-tool-call-extractor.ts diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 2a2065edd6..1803b70d20 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -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 { diff --git a/src/api/providers/utils/__tests__/kimi-tool-call-extractor.spec.ts b/src/api/providers/utils/__tests__/kimi-tool-call-extractor.spec.ts new file mode 100644 index 0000000000..d85b6cb519 --- /dev/null +++ b/src/api/providers/utils/__tests__/kimi-tool-call-extractor.spec.ts @@ -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) + }) + }) +}) diff --git a/src/api/providers/utils/kimi-tool-call-extractor.ts b/src/api/providers/utils/kimi-tool-call-extractor.ts new file mode 100644 index 0000000000..d89546e082 --- /dev/null +++ b/src/api/providers/utils/kimi-tool-call-extractor.ts @@ -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") + ) +}