mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-09 22:31:08 +00:00
feat: add support for Kimi K2 Thinking native tool call format
- Created KimiK2ToolCallParser class to parse special token format - Integrated parser into OpenAI provider for Kimi K2 models - Added comprehensive tests for the parser - Fixes issue where Kimi K2 tool calls were not being executed Fixes #9551
This commit is contained in:
Roo Code
parent
cad6145241
commit
c3afafbd9a
3 changed files with 536 additions and 2 deletions
|
|
@ -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<number, { id: string; name: string; arguments: string }>()
|
||||
|
||||
// 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
|
||||
}
|
||||
|
|
|
|||
231
src/api/providers/utils/__tests__/kimi-k2-tool-parser.spec.ts
Normal file
231
src/api/providers/utils/__tests__/kimi-k2-tool-parser.spec.ts
Normal file
|
|
@ -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: "{}",
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
262
src/api/providers/utils/kimi-k2-tool-parser.ts
Normal file
262
src/api/providers/utils/kimi-k2-tool-parser.ts
Normal file
|
|
@ -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"))
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue