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") + ) +}