diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 1c8d3c7d9d..ca1d7a4197 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 { KimiK2ToolCallParser } from "./utils/kimi-k2-tool-parser" // TODO: Rename this to OpenAICompatibleHandler. Also, I think the // `OpenAINativeHandler` can subclass from this, since it's obviously @@ -193,13 +194,40 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl let lastUsage const toolCallAccumulator = new Map() + // Initialize Kimi K2 parser if needed + const isKimiK2 = KimiK2ToolCallParser.isKimiK2Model(modelId) + const kimiK2Parser = isKimiK2 ? new KimiK2ToolCallParser() : null + for await (const chunk of stream) { const delta = chunk.choices?.[0]?.delta ?? {} const finishReason = chunk.choices?.[0]?.finish_reason if (delta.content) { - for (const chunk of matcher.update(delta.content)) { - yield chunk + // For Kimi K2 models, parse special tool call tokens from content + if (kimiK2Parser) { + const parsed = kimiK2Parser.processChunk(delta.content) + + // Yield any extracted tool calls + for (const toolCall of parsed.toolCalls) { + yield { + type: "tool_call", + id: toolCall.id, + name: toolCall.name, + arguments: toolCall.arguments, + } + } + + // Process remaining content through think matcher + if (parsed.content) { + for (const chunk of matcher.update(parsed.content)) { + yield chunk + } + } + } else { + // Regular processing for non-Kimi K2 models + for (const chunk of matcher.update(delta.content)) { + yield chunk + } } } @@ -260,6 +288,19 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl toolCallAccumulator.clear() } + // Flush any pending Kimi K2 tool calls + if (kimiK2Parser) { + const flushed = kimiK2Parser.flush() + for (const toolCall of flushed.toolCalls) { + yield { + type: "tool_call", + id: toolCall.id, + name: toolCall.name, + arguments: toolCall.arguments, + } + } + } + for (const chunk of matcher.final()) { yield chunk } diff --git a/src/api/providers/utils/__tests__/kimi-k2-tool-parser.spec.ts b/src/api/providers/utils/__tests__/kimi-k2-tool-parser.spec.ts new file mode 100644 index 0000000000..3ad19f78a6 --- /dev/null +++ b/src/api/providers/utils/__tests__/kimi-k2-tool-parser.spec.ts @@ -0,0 +1,231 @@ +import { describe, it, expect } from "vitest" +import { KimiK2ToolCallParser } from "../kimi-k2-tool-parser" + +describe("KimiK2ToolCallParser", () => { + describe("isKimiK2Model", () => { + it("should identify Kimi K2 models correctly", () => { + expect(KimiK2ToolCallParser.isKimiK2Model("Kimi-K2-Thinking")).toBe(true) + expect(KimiK2ToolCallParser.isKimiK2Model("kimi-k2-thinking")).toBe(true) + expect(KimiK2ToolCallParser.isKimiK2Model("Kimi-K2")).toBe(true) + expect(KimiK2ToolCallParser.isKimiK2Model("kimi-thinking")).toBe(true) + expect(KimiK2ToolCallParser.isKimiK2Model("KIMI-K2-THINKING")).toBe(true) + }) + + it("should not identify non-Kimi K2 models", () => { + expect(KimiK2ToolCallParser.isKimiK2Model("gpt-4")).toBe(false) + expect(KimiK2ToolCallParser.isKimiK2Model("claude-3")).toBe(false) + expect(KimiK2ToolCallParser.isKimiK2Model("kimi-v1")).toBe(false) + expect(KimiK2ToolCallParser.isKimiK2Model("thinking-model")).toBe(false) + }) + }) + + describe("processChunk", () => { + it("should parse single tool call correctly", () => { + const parser = new KimiK2ToolCallParser() + + const chunk1 = "Let me read the file.\n<|tool_calls_section_begin|>" + const chunk2 = "<|tool_call_begin|>functions.read_file:0" + const chunk3 = '<|tool_call_argument_begin|>{"files":[{"path":"test.txt"}]}' + const chunk4 = "<|tool_call_end|><|tool_calls_section_end|>" + + let result = parser.processChunk(chunk1) + expect(result.content).toBe("Let me read the file.\n") + expect(result.toolCalls).toHaveLength(0) + expect(result.isBuffering).toBe(true) + + result = parser.processChunk(chunk2) + expect(result.content).toBe("") + expect(result.toolCalls).toHaveLength(0) + expect(result.isBuffering).toBe(true) + + result = parser.processChunk(chunk3) + expect(result.content).toBe("") + expect(result.toolCalls).toHaveLength(0) + expect(result.isBuffering).toBe(true) + + result = parser.processChunk(chunk4) + expect(result.content).toBe("") + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0]).toEqual({ + id: "tool_0", + name: "read_file", + arguments: '{"files":[{"path":"test.txt"}]}', + }) + expect(result.isBuffering).toBe(false) + }) + + it("should parse multiple tool calls", () => { + const parser = new KimiK2ToolCallParser() + + const input = `<|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.write_to_file:1<|tool_call_argument_begin|>{"path":"file2.txt","content":"Test","line_count":1}<|tool_call_end|> +<|tool_calls_section_end|>` + + const result = parser.processChunk(input) + expect(result.content).toBe("") + expect(result.toolCalls).toHaveLength(2) + expect(result.toolCalls[0]).toEqual({ + id: "tool_0", + name: "read_file", + arguments: '{"files":[{"path":"file1.txt"}]}', + }) + expect(result.toolCalls[1]).toEqual({ + id: "tool_1", + name: "write_to_file", + arguments: '{"path":"file2.txt","content":"Test","line_count":1}', + }) + expect(result.isBuffering).toBe(false) + }) + + it("should handle mixed content and tool calls", () => { + const parser = new KimiK2ToolCallParser() + + const input = `I'll help you with that. Let me first read the file. +<|tool_calls_section_begin|> +<|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>{"files":[{"path":"config.json"}]}<|tool_call_end|> +<|tool_calls_section_end|> +Now let me process the data.` + + const result = parser.processChunk(input) + expect(result.content).toContain("I'll help you with that") + expect(result.content).toContain("Now let me process the data") + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0]).toEqual({ + id: "tool_0", + name: "read_file", + arguments: '{"files":[{"path":"config.json"}]}', + }) + }) + + it("should handle chunked input correctly", () => { + const parser = new KimiK2ToolCallParser() + + // Simulate streaming chunks that split across boundaries + const chunks = [ + "Starting task\n<|tool_", + "calls_section_begin|><|tool_call", + "_begin|>functions.wr", + "ite_to_file:0<|tool_call_argu", + 'ment_begin|>{"path":"test.md",', + '"content":"# Test","line_count":1}<|tool_', + "call_end|><|tool_calls_section_end|>\nDone!", + ] + + let allContent = "" + let allToolCalls: any[] = [] + + for (const chunk of chunks) { + const result = parser.processChunk(chunk) + allContent += result.content + allToolCalls.push(...result.toolCalls) + } + + expect(allContent).toBe("Starting task\n\nDone!") + expect(allToolCalls).toHaveLength(1) + expect(allToolCalls[0]).toEqual({ + id: "tool_0", + name: "write_to_file", + arguments: '{"path":"test.md","content":"# Test","line_count":1}', + }) + }) + + it("should handle tool calls with complex JSON arguments", () => { + const parser = new KimiK2ToolCallParser() + + const input = `<|tool_calls_section_begin|> +<|tool_call_begin|>functions.apply_diff:0<|tool_call_argument_begin|>{ + "path": "src/main.ts", + "diff": "<<<<<<< SEARCH\\nold code\\n=======\\nnew code\\n>>>>>>> REPLACE" +}<|tool_call_end|> +<|tool_calls_section_end|>` + + const result = parser.processChunk(input) + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0].name).toBe("apply_diff") + expect(result.toolCalls[0].arguments).toContain("<<<<<<< SEARCH") + expect(result.toolCalls[0].arguments).toContain(">>>>>>> REPLACE") + }) + }) + + describe("flush", () => { + it("should flush pending tool calls", () => { + const parser = new KimiK2ToolCallParser() + + // Start a tool call but don't close it + const chunk = + '<|tool_calls_section_begin|><|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>{"files":[{"path":"test.txt"}]}' + + let result = parser.processChunk(chunk) + expect(result.toolCalls).toHaveLength(0) + expect(result.isBuffering).toBe(true) + + // Force flush + const flushed = parser.flush() + expect(flushed.toolCalls).toHaveLength(1) + expect(flushed.toolCalls[0]).toEqual({ + id: "tool_0", + name: "read_file", + arguments: '{"files":[{"path":"test.txt"}]}', + }) + }) + + it("should reset state after flush", () => { + const parser = new KimiK2ToolCallParser() + + // Process partial tool call + parser.processChunk("<|tool_calls_section_begin|><|tool_call_begin|>functions.test:0") + + // Flush + parser.flush() + + // Process new content - should not be buffering + const result = parser.processChunk("Regular content") + expect(result.content).toBe("Regular content") + expect(result.isBuffering).toBe(false) + }) + }) + + describe("edge cases", () => { + it("should handle empty tool calls section", () => { + const parser = new KimiK2ToolCallParser() + + const input = "<|tool_calls_section_begin|><|tool_calls_section_end|>" + const result = parser.processChunk(input) + + expect(result.content).toBe("") + expect(result.toolCalls).toHaveLength(0) + expect(result.isBuffering).toBe(false) + }) + + it("should handle newlines in tool name/id section", () => { + const parser = new KimiK2ToolCallParser() + + const input = `<|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|>` + + const result = parser.processChunk(input) + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0].name).toBe("read_file") + expect(result.toolCalls[0].id).toBe("tool_0") + }) + + it("should handle tool calls without arguments", () => { + const parser = new KimiK2ToolCallParser() + + const input = + "<|tool_calls_section_begin|><|tool_call_begin|>functions.get_status:0<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + + const result = parser.processChunk(input) + expect(result.toolCalls).toHaveLength(1) + expect(result.toolCalls[0]).toEqual({ + id: "tool_0", + name: "get_status", + arguments: "{}", + }) + }) + }) +}) diff --git a/src/api/providers/utils/kimi-k2-tool-parser.ts b/src/api/providers/utils/kimi-k2-tool-parser.ts new file mode 100644 index 0000000000..5e5e1bfc0e --- /dev/null +++ b/src/api/providers/utils/kimi-k2-tool-parser.ts @@ -0,0 +1,262 @@ +/** + * Parser for Kimi K2 Thinking model's native tool call format. + * Kimi K2 uses special tokens in the content stream rather than the standard OpenAI tool_calls field. + * + * Token format: + * <|tool_calls_section_begin|> + * <|tool_call_begin|> + * functions.tool_name:call_id + * <|tool_call_argument_begin|> + * {"arg": "value"} + * <|tool_call_end|> + * <|tool_calls_section_end|> + */ +export class KimiK2ToolCallParser { + private contentBuffer = "" + private isInToolCallSection = false + private isInToolCall = false + private isInArguments = false + private currentToolCall: { + id: string + name: string + arguments: string + } | null = null + private pendingToolCalls: Array<{ + id: string + name: string + arguments: string + }> = [] + + // Special tokens used by Kimi K2 + private readonly TOOL_CALLS_BEGIN = "<|tool_calls_section_begin|>" + private readonly TOOL_CALLS_END = "<|tool_calls_section_end|>" + private readonly TOOL_CALL_BEGIN = "<|tool_call_begin|>" + private readonly TOOL_CALL_END = "<|tool_call_end|>" + private readonly TOOL_ARG_BEGIN = "<|tool_call_argument_begin|>" + + /** + * Process incoming content chunk and extract tool calls + * @param chunk - The content chunk to process + * @returns Object containing remaining content and extracted tool calls + */ + processChunk(chunk: string): { + content: string + toolCalls: Array<{ id: string; name: string; arguments: string }> + isBuffering: boolean + } { + this.contentBuffer += chunk + const extractedToolCalls: Array<{ id: string; name: string; arguments: string }> = [] + let processedContent = "" + + // Process the buffer + let i = 0 + while (i < this.contentBuffer.length) { + let tokenFound = false + + // Check if we might be at the start of a special token + // We need to check if we have enough characters for the smallest token + const remainingChars = this.contentBuffer.length - i + + // Check for special tokens only if we have enough characters + if ( + !this.isInToolCallSection && + remainingChars >= this.TOOL_CALLS_BEGIN.length && + this.contentBuffer.substring(i).startsWith(this.TOOL_CALLS_BEGIN) + ) { + // Start of tool calls section + this.isInToolCallSection = true + i += this.TOOL_CALLS_BEGIN.length + tokenFound = true + } else if ( + !this.isInToolCallSection && + remainingChars < this.TOOL_CALLS_BEGIN.length && + this.TOOL_CALLS_BEGIN.startsWith(this.contentBuffer.substring(i)) + ) { + // Might be the start of TOOL_CALLS_BEGIN but we don't have enough characters yet + break // Buffer the rest for next chunk + } else if ( + this.isInToolCallSection && + remainingChars >= this.TOOL_CALLS_END.length && + this.contentBuffer.substring(i).startsWith(this.TOOL_CALLS_END) + ) { + // End of tool calls section + this.isInToolCallSection = false + // Flush any pending tool calls + if (this.pendingToolCalls.length > 0) { + extractedToolCalls.push(...this.pendingToolCalls) + this.pendingToolCalls = [] + } + i += this.TOOL_CALLS_END.length + tokenFound = true + } else if ( + this.isInToolCallSection && + remainingChars < this.TOOL_CALLS_END.length && + this.TOOL_CALLS_END.startsWith(this.contentBuffer.substring(i)) + ) { + // Might be the start of TOOL_CALLS_END but we don't have enough characters yet + break // Buffer the rest for next chunk + } else if ( + this.isInToolCallSection && + !this.isInToolCall && + remainingChars >= this.TOOL_CALL_BEGIN.length && + this.contentBuffer.substring(i).startsWith(this.TOOL_CALL_BEGIN) + ) { + // Start of individual tool call + this.isInToolCall = true + this.currentToolCall = { id: "", name: "", arguments: "" } + i += this.TOOL_CALL_BEGIN.length + tokenFound = true + } else if ( + this.isInToolCallSection && + !this.isInToolCall && + remainingChars < this.TOOL_CALL_BEGIN.length && + this.TOOL_CALL_BEGIN.startsWith(this.contentBuffer.substring(i)) + ) { + // Might be the start of TOOL_CALL_BEGIN but we don't have enough characters yet + break // Buffer the rest for next chunk + } else if ( + this.isInToolCall && + remainingChars >= this.TOOL_CALL_END.length && + this.contentBuffer.substring(i).startsWith(this.TOOL_CALL_END) + ) { + // End of individual tool call + this.isInToolCall = false + this.isInArguments = false + if (this.currentToolCall) { + this.pendingToolCalls.push(this.currentToolCall) + this.currentToolCall = null + } + i += this.TOOL_CALL_END.length + tokenFound = true + } else if ( + this.isInToolCall && + remainingChars < this.TOOL_CALL_END.length && + this.TOOL_CALL_END.startsWith(this.contentBuffer.substring(i)) + ) { + // Might be the start of TOOL_CALL_END but we don't have enough characters yet + break // Buffer the rest for next chunk + } else if ( + this.isInToolCall && + !this.isInArguments && + remainingChars >= this.TOOL_ARG_BEGIN.length && + this.contentBuffer.substring(i).startsWith(this.TOOL_ARG_BEGIN) + ) { + // Start of arguments section + this.isInArguments = true + i += this.TOOL_ARG_BEGIN.length + tokenFound = true + } else if ( + this.isInToolCall && + !this.isInArguments && + remainingChars < this.TOOL_ARG_BEGIN.length && + this.TOOL_ARG_BEGIN.startsWith(this.contentBuffer.substring(i)) + ) { + // Might be the start of TOOL_ARG_BEGIN but we don't have enough characters yet + break // Buffer the rest for next chunk + } + + if (!tokenFound) { + // Process content based on current state + if (this.isInToolCall && this.currentToolCall) { + if (!this.isInArguments) { + // Parsing tool name and ID (format: functions.tool_name:call_id) + const char = this.contentBuffer[i] + if (char === "\n" || char === "\r") { + // Skip newlines + i++ + continue + } + + // Buffer the tool name/id string + const toolInfo = this.currentToolCall.name + char + this.currentToolCall.name = toolInfo + + // Check if we've reached the end of tool name/id + if (toolInfo.includes(":")) { + // Parse the format: functions.tool_name:call_id + const parts = toolInfo.match(/^functions\.([^:]+):(.+)$/) + if (parts) { + this.currentToolCall.name = parts[1] + this.currentToolCall.id = `tool_${parts[2]}` + } + } + } else { + // Parsing arguments JSON + this.currentToolCall.arguments += this.contentBuffer[i] + } + } else if (!this.isInToolCallSection) { + // Check if this might be the start of a tool token when not in a section + const possibleTokenStarts = [this.TOOL_CALLS_BEGIN] + let mightBeToken = false + + for (const token of possibleTokenStarts) { + if (token.startsWith(this.contentBuffer.substring(i))) { + // This might be the start of a token + mightBeToken = true + break + } + } + + if (mightBeToken && remainingChars < this.TOOL_CALLS_BEGIN.length) { + // Buffer this for next chunk + break + } else { + // Regular content outside tool calls + processedContent += this.contentBuffer[i] + } + } + + i++ + } + } + + // Update the buffer to only contain unprocessed content + this.contentBuffer = this.contentBuffer.substring(i) + + // If we're still in a tool call section, we need to buffer more content + const isBuffering = this.isInToolCallSection || this.isInToolCall + + return { + content: processedContent, + toolCalls: extractedToolCalls, + isBuffering, + } + } + + /** + * Force flush any buffered content and pending tool calls + * Used when the stream ends + */ + flush(): { + content: string + toolCalls: Array<{ id: string; name: string; arguments: string }> + } { + const toolCalls = [...this.pendingToolCalls] + + // If we have a current tool call in progress, add it + if (this.currentToolCall) { + toolCalls.push(this.currentToolCall) + } + + // Reset state + this.contentBuffer = "" + this.isInToolCallSection = false + this.isInToolCall = false + this.isInArguments = false + this.currentToolCall = null + this.pendingToolCalls = [] + + return { + content: "", + toolCalls, + } + } + + /** + * Check if the model is likely Kimi K2 based on model ID + */ + static isKimiK2Model(modelId: string): boolean { + const lowerModelId = modelId.toLowerCase() + return lowerModelId.includes("kimi") && (lowerModelId.includes("k2") || lowerModelId.includes("thinking")) + } +}