From ee2015d5ff8721a86bd866b60c43741640fad84d Mon Sep 17 00:00:00 2001 From: Roo Code Date: Tue, 18 Nov 2025 19:55:35 +0000 Subject: [PATCH] fix: add support for Kimi K2 tool call format in Chutes.AI provider - Created KimiToolCallParser to handle the unique <|tool_calls_section_begin|> format - Updated ChutesHandler to use the parser for Kimi K2 models - Added comprehensive tests for the new parser and integration - Fixes issue #9366 where Kimi K2 Thinking model was not working properly --- src/api/providers/__tests__/chutes.spec.ts | 109 +++++++++++++ .../__tests__/kimi-tool-parser.spec.ts | 148 ++++++++++++++++++ src/api/providers/chutes.ts | 55 ++++++- src/api/providers/kimi-tool-parser.ts | 145 +++++++++++++++++ 4 files changed, 456 insertions(+), 1 deletion(-) create mode 100644 src/api/providers/__tests__/kimi-tool-parser.spec.ts create mode 100644 src/api/providers/kimi-tool-parser.ts diff --git a/src/api/providers/__tests__/chutes.spec.ts b/src/api/providers/__tests__/chutes.spec.ts index b4c933d4cc..a8bc30c2cc 100644 --- a/src/api/providers/__tests__/chutes.spec.ts +++ b/src/api/providers/__tests__/chutes.spec.ts @@ -256,4 +256,113 @@ describe("ChutesHandler", () => { // The default model is DeepSeek-R1, so it returns DEEP_SEEK_DEFAULT_TEMPERATURE expect(model.info.temperature).toBe(DEEP_SEEK_DEFAULT_TEMPERATURE) }) + + it("should handle Kimi K2 tool call format", async () => { + // Mock Kimi K2 model response with tool call format + mockCreate.mockImplementationOnce(async () => ({ + [Symbol.asyncIterator]: async function* () { + yield { + choices: [ + { + delta: { + content: + 'I\'ll help you with that. <|tool_calls_section_begin|> <|tool_call_begin|> functions.codebase_search:12 <|tool_call_argument_begin|> {"query": "TeamSelect scene", "path": "ouroboros"} <|tool_call_end|> <|tool_calls_section_end|>', + }, + index: 0, + }, + ], + usage: null, + } + yield { + choices: [ + { + delta: { content: " Let me search for that." }, + index: 0, + }, + ], + usage: null, + } + yield { + choices: [ + { + delta: {}, + index: 0, + }, + ], + usage: { prompt_tokens: 15, completion_tokens: 25 }, + } + }, + })) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Find TeamSelect scene" }] + mockFetchModel.mockResolvedValueOnce({ + id: "moonshotai/Kimi-K2-Instruct-75k", + info: { maxTokens: 32768, temperature: 0.5 }, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + // Should parse the Kimi K2 tool call format correctly + expect(chunks).toEqual([ + { type: "text", text: "I'll help you with that. " }, + { + type: "tool_call", + id: "tool_call_12", + name: "codebase_search", + arguments: '{"query":"TeamSelect scene","path":"ouroboros"}', + }, + { type: "text", text: " Let me search for that." }, + { type: "usage", inputTokens: 15, outputTokens: 25 }, + ]) + }) + + it("should handle Kimi K2 model without tool calls", async () => { + // Mock Kimi K2 model response without tool calls + mockCreate.mockImplementationOnce(async () => ({ + [Symbol.asyncIterator]: async function* () { + yield { + choices: [ + { + delta: { content: "This is a regular response without tool calls." }, + index: 0, + }, + ], + usage: null, + } + yield { + choices: [ + { + delta: {}, + index: 0, + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 8 }, + } + }, + })) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + mockFetchModel.mockResolvedValueOnce({ + id: "moonshotai/Kimi-K2-Instruct-0905", + info: { maxTokens: 32768, temperature: 0.5 }, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + // Should handle regular text without tool calls + expect(chunks).toEqual([ + { type: "text", text: "This is a regular response without tool calls." }, + { type: "usage", inputTokens: 10, outputTokens: 8 }, + ]) + }) }) diff --git a/src/api/providers/__tests__/kimi-tool-parser.spec.ts b/src/api/providers/__tests__/kimi-tool-parser.spec.ts new file mode 100644 index 0000000000..0142683cb8 --- /dev/null +++ b/src/api/providers/__tests__/kimi-tool-parser.spec.ts @@ -0,0 +1,148 @@ +// npx vitest run api/providers/__tests__/kimi-tool-parser.spec.ts + +import { describe, it, expect, beforeEach } from "vitest" +import { KimiToolCallParser } from "../kimi-tool-parser" + +describe("KimiToolCallParser", () => { + let parser: KimiToolCallParser + + beforeEach(() => { + parser = new KimiToolCallParser() + }) + + describe("processChunk", () => { + it("should parse a complete tool call section", () => { + const chunk = `I'll help you with that. <|tool_calls_section_begin|> <|tool_call_begin|> functions.codebase_search:12 <|tool_call_argument_begin|> {"query": "TeamSelect scene", "path": "ouroboros"} <|tool_call_end|> <|tool_calls_section_end|>` + + const results = parser.processChunk(chunk) + + expect(results).toHaveLength(2) + expect(results[0]).toEqual({ + type: "text", + content: "I'll help you with that. ", + }) + expect(results[1]).toEqual({ + type: "tool_call", + toolCall: { + id: "tool_call_12", + name: "codebase_search", + arguments: '{"query":"TeamSelect scene","path":"ouroboros"}', + }, + }) + }) + + it("should handle tool calls without functions prefix", () => { + const chunk = `<|tool_calls_section_begin|> <|tool_call_begin|> read_file:5 <|tool_call_argument_begin|> {"path": "test.ts"} <|tool_call_end|> <|tool_calls_section_end|>` + + const results = parser.processChunk(chunk) + + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "tool_call", + toolCall: { + id: "tool_call_5", + name: "read_file", + arguments: '{"path":"test.ts"}', + }, + }) + }) + + it("should handle partial tool call sections across chunks", () => { + const chunk1 = `Some text before <|tool_calls_section_begin|> <|tool_call_begin|> functions.` + const chunk2 = `search_files:15 <|tool_call_argument_begin|> {"query": "test"} <|tool_call_end|> <|tool_calls_section_end|> and text after` + + const results1 = parser.processChunk(chunk1) + expect(results1).toHaveLength(1) + expect(results1[0]).toEqual({ + type: "text", + content: "Some text before ", + }) + + const results2 = parser.processChunk(chunk2) + expect(results2).toHaveLength(2) + expect(results2[0]).toEqual({ + type: "tool_call", + toolCall: { + id: "tool_call_15", + name: "search_files", + arguments: '{"query":"test"}', + }, + }) + expect(results2[1]).toEqual({ + type: "text", + content: " and text after", + }) + }) + + it("should handle multiple tool calls in sequence", () => { + const chunk = `<|tool_calls_section_begin|> <|tool_call_begin|> functions.read_file:1 <|tool_call_argument_begin|> {"path": "file1.ts"} <|tool_call_end|> <|tool_calls_section_end|> Then <|tool_calls_section_begin|> <|tool_call_begin|> functions.write_file:2 <|tool_call_argument_begin|> {"path": "file2.ts", "content": "test"} <|tool_call_end|> <|tool_calls_section_end|>` + + const results = parser.processChunk(chunk) + + expect(results).toHaveLength(3) + expect(results[0].type).toBe("tool_call") + expect(results[0].toolCall?.name).toBe("read_file") + expect(results[1].type).toBe("text") + expect(results[1].content).toBe(" Then ") + expect(results[2].type).toBe("tool_call") + expect(results[2].toolCall?.name).toBe("write_file") + }) + + it("should handle text without tool calls", () => { + const chunk = "This is just regular text without any tool calls." + + const results = parser.processChunk(chunk) + + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "text", + content: "This is just regular text without any tool calls.", + }) + }) + + it("should handle malformed tool call sections gracefully", () => { + const chunk = `<|tool_calls_section_begin|> <|tool_call_begin|> invalid_format <|tool_call_end|> <|tool_calls_section_end|>` + + const results = parser.processChunk(chunk) + + // Should not parse as tool call due to invalid format + expect(results).toHaveLength(0) + }) + + it("should handle invalid JSON in tool arguments", () => { + const chunk = `<|tool_calls_section_begin|> <|tool_call_begin|> functions.test:1 <|tool_call_argument_begin|> {invalid json} <|tool_call_end|> <|tool_calls_section_end|>` + + const results = parser.processChunk(chunk) + + // Should not parse as tool call due to invalid JSON + expect(results).toHaveLength(0) + }) + }) + + describe("flush", () => { + it("should return remaining buffer content", () => { + const chunk = "Some incomplete text <|tool_calls_section_begin" + + parser.processChunk(chunk) + const results = parser.flush() + + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "text", + content: "Some incomplete text <|tool_calls_section_begin", + }) + }) + + it("should return empty array when buffer is empty", () => { + const results = parser.flush() + expect(results).toHaveLength(0) + }) + + it("should clear buffer after flush", () => { + parser.processChunk("Some text") + parser.flush() + const secondFlush = parser.flush() + expect(secondFlush).toHaveLength(0) + }) + }) +}) diff --git a/src/api/providers/chutes.ts b/src/api/providers/chutes.ts index d19c55abce..9f9b048652 100644 --- a/src/api/providers/chutes.ts +++ b/src/api/providers/chutes.ts @@ -11,6 +11,7 @@ import { ApiStream } from "../transform/stream" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { RouterProvider } from "./router-provider" +import { KimiToolCallParser } from "./kimi-tool-parser" export class ChutesHandler extends RouterProvider implements SingleCompletionHandler { constructor(options: ApiHandlerOptions) { @@ -100,8 +101,60 @@ export class ChutesHandler extends RouterProvider implements SingleCompletionHan for (const processedChunk of matcher.final()) { yield processedChunk } + } else if (model.id.includes("Kimi-K2")) { + // Special handling for Kimi K2 models with their unique tool call format + const stream = await this.client.chat.completions.create(this.getCompletionParams(systemPrompt, messages)) + const kimiParser = new KimiToolCallParser() + + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + + if (delta?.content) { + // Process through Kimi parser to extract tool calls + const parsed = kimiParser.processChunk(delta.content) + for (const item of parsed) { + if (item.type === "text" && item.content) { + yield { type: "text", text: item.content } + } else if (item.type === "tool_call" && item.toolCall) { + yield { + type: "tool_call", + id: item.toolCall.id, + name: item.toolCall.name, + arguments: item.toolCall.arguments, + } + } + } + } + + if (delta && "reasoning_content" in delta && delta.reasoning_content) { + yield { type: "reasoning", text: (delta.reasoning_content as string | undefined) || "" } + } + + if (chunk.usage) { + yield { + type: "usage", + inputTokens: chunk.usage.prompt_tokens || 0, + outputTokens: chunk.usage.completion_tokens || 0, + } + } + } + + // Flush any remaining content from the parser + const remaining = kimiParser.flush() + for (const item of remaining) { + if (item.type === "text" && item.content) { + yield { type: "text", text: item.content } + } else if (item.type === "tool_call" && item.toolCall) { + yield { + type: "tool_call", + id: item.toolCall.id, + name: item.toolCall.name, + arguments: item.toolCall.arguments, + } + } + } } else { - // For non-DeepSeek-R1 models, use standard OpenAI streaming + // For non-DeepSeek-R1 and non-Kimi models, use standard OpenAI streaming const stream = await this.client.chat.completions.create(this.getCompletionParams(systemPrompt, messages)) for await (const chunk of stream) { diff --git a/src/api/providers/kimi-tool-parser.ts b/src/api/providers/kimi-tool-parser.ts new file mode 100644 index 0000000000..113c156172 --- /dev/null +++ b/src/api/providers/kimi-tool-parser.ts @@ -0,0 +1,145 @@ +/** + * Parser for Kimi K2 model's tool call format + * Handles the <|tool_calls_section_begin|> style format + */ +export class KimiToolCallParser { + private buffer: string = "" + private inToolCallSection: boolean = false + private currentToolCall: { + id?: string + name?: string + arguments?: string + } = {} + + /** + * Process a chunk of text and extract tool calls if present + * @param chunk The text chunk to process + * @returns Array of parsed content (text or tool calls) + */ + public processChunk(chunk: string): Array<{ type: "text" | "tool_call"; content?: string; toolCall?: any }> { + const results: Array<{ type: "text" | "tool_call"; content?: string; toolCall?: any }> = [] + this.buffer += chunk + + // Check for tool call section markers + const toolCallStartPattern = /<\|tool_calls_section_begin\|>/g + const toolCallSectionEndPattern = /<\|tool_calls_section_end\|>/g + const toolCallBeginPattern = /<\|tool_call_begin\|>/g + const toolCallEndPattern = /<\|tool_call_end\|>/g + const toolCallArgPattern = /<\|tool_call_argument_begin\|>/g + + let processedBuffer = this.buffer + + // Process tool call sections + while (true) { + const startMatch = processedBuffer.match(toolCallStartPattern) + if (!startMatch) break + + const startIndex = startMatch.index! + const beforeToolCall = processedBuffer.substring(0, startIndex) + + // Add any text before the tool call section + if (beforeToolCall.trim()) { + results.push({ type: "text", content: beforeToolCall }) + } + + // Find the end of the tool call section + const endMatch = processedBuffer.substring(startIndex).match(toolCallSectionEndPattern) + if (!endMatch) { + // Tool call section not complete yet, keep it in buffer + this.buffer = processedBuffer.substring(startIndex) + return results + } + + const endIndex = startIndex + endMatch.index! + endMatch[0].length + const toolCallSection = processedBuffer.substring(startIndex, endIndex) + + // Parse the tool call section + const toolCall = this.parseToolCallSection(toolCallSection) + if (toolCall) { + results.push({ type: "tool_call", toolCall }) + } + + // Continue processing the rest of the buffer + processedBuffer = processedBuffer.substring(endIndex) + } + + // Handle remaining text + if (processedBuffer.trim()) { + // Check if we might be in the middle of a tool call marker + const partialMarkers = ["<|tool_calls_section_begin", "<|tool_call_begin", "<|tool_call_argument_begin"] + const hasPartialMarker = partialMarkers.some((marker) => processedBuffer.includes(marker)) + + if (hasPartialMarker) { + // Keep partial markers in buffer for next chunk + this.buffer = processedBuffer + } else { + // Output remaining text + results.push({ type: "text", content: processedBuffer }) + this.buffer = "" + } + } else { + this.buffer = "" + } + + return results + } + + /** + * Parse a complete tool call section + */ + private parseToolCallSection(section: string): any | null { + // Extract tool call details + const toolCallMatch = section.match( + /<\|tool_call_begin\|>\s*functions\.(\w+):(\d+)\s*<\|tool_call_argument_begin\|>\s*({[^}]*})\s*<\|tool_call_end\|>/, + ) + + if (toolCallMatch) { + const [, functionName, callId, argumentsJson] = toolCallMatch + try { + const args = JSON.parse(argumentsJson) + return { + id: `tool_call_${callId}`, + name: functionName, + arguments: JSON.stringify(args), + } + } catch (e) { + console.error("Failed to parse Kimi tool call arguments:", e) + return null + } + } + + // Alternative format without explicit function prefix + const altMatch = section.match( + /<\|tool_call_begin\|>\s*(\w+):(\d+)\s*<\|tool_call_argument_begin\|>\s*({[^}]*})\s*<\|tool_call_end\|>/, + ) + + if (altMatch) { + const [, functionName, callId, argumentsJson] = altMatch + try { + const args = JSON.parse(argumentsJson) + return { + id: `tool_call_${callId}`, + name: functionName, + arguments: JSON.stringify(args), + } + } catch (e) { + console.error("Failed to parse Kimi tool call arguments:", e) + return null + } + } + + return null + } + + /** + * Get any remaining buffered content + */ + public flush(): Array<{ type: "text" | "tool_call"; content?: string; toolCall?: any }> { + const results: Array<{ type: "text" | "tool_call"; content?: string; toolCall?: any }> = [] + if (this.buffer.trim()) { + results.push({ type: "text", content: this.buffer }) + this.buffer = "" + } + return results + } +}