Skip to main content

tool_parser/parsers/
kimik2.rs

1use async_trait::async_trait;
2use openai_protocol::common::Tool;
3use regex::Regex;
4use serde_json::Value;
5
6use crate::{
7    errors::{ParserError, ParserResult},
8    parsers::helpers,
9    traits::ToolParser,
10    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
11};
12
13/// Kimi K2 format parser for tool calls
14///
15/// Handles the Kimi K2 specific format:
16/// `<|tool_calls_section_begin|><|tool_call_begin|>functions.{name}:{index}<|tool_call_argument_begin|>{json_args}<|tool_call_end|><|tool_calls_section_end|>`
17///
18/// Features:
19/// - Token-based delimiters
20/// - Function calls with explicit indexing
21/// - JSON arguments
22///
23/// Reference: <https://huggingface.co/moonshotai/Kimi-K2-Instruct/blob/main/docs/tool_call_guidance.md>
24pub struct KimiK2Parser {
25    /// Regex for extracting complete tool calls
26    tool_call_extractor: Regex,
27    /// Regex for extracting partial tool calls (streaming)
28    stream_tool_call_extractor: Regex,
29    /// Regex pattern for removing completed tool calls from buffer
30    tool_call_end_pattern: Regex,
31    /// Robust parser for ids like "functions.search:0" or fallback "search:0"
32    tool_call_id_regex: Regex,
33
34    /// Buffer for accumulating incomplete patterns across chunks
35    buffer: String,
36
37    /// Stores complete tool call info (name and arguments) for each tool being parsed
38    prev_tool_call_arr: Vec<Value>,
39
40    /// Index of currently streaming tool call (-1 means no active tool)
41    current_tool_id: i32,
42
43    /// Flag for whether current tool's name has been sent to client
44    current_tool_name_sent: bool,
45
46    /// Tracks raw JSON string content streamed to client for each tool's arguments
47    streamed_args_for_tool: Vec<String>,
48
49    /// Tracks the last arguments sent for incremental diffing
50    last_arguments: String,
51}
52
53impl KimiK2Parser {
54    /// Build structural tag for Kimi K2 tool call format.
55    ///
56    /// Uses dual triggers following the serving engine's approach:
57    /// - `<|tool_calls_section_begin|>` for the first tool call
58    /// - `<|tool_call_begin|>` for subsequent parallel tool calls
59    pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
60        let mut tags = Vec::new();
61        for (index, tool) in tools.iter().enumerate() {
62            let name = &tool.function.name;
63            if name.is_empty() {
64                continue;
65            }
66            let schema = &tool.function.parameters;
67
68            // Tag for FIRST tool call (includes outer section wrapper)
69            tags.push(serde_json::json!({
70                "begin": format!(
71                    "<|tool_calls_section_begin|><|tool_call_begin|>functions.{name}:{index}<|tool_call_argument_begin|>"
72                ),
73                "content": {
74                    "type": "json_schema",
75                    "json_schema": schema,
76                },
77                "end": "<|tool_call_end|>",
78            }));
79
80            // Tag for SUBSEQUENT tool calls (no outer section wrapper)
81            tags.push(serde_json::json!({
82                "begin": format!(
83                    "<|tool_call_begin|>functions.{name}:{index}<|tool_call_argument_begin|>"
84                ),
85                "content": {
86                    "type": "json_schema",
87                    "json_schema": schema,
88                },
89                "end": "<|tool_call_end|>",
90            }));
91        }
92        serde_json::json!({
93            "format": {
94                "type": "triggered_tags",
95                "triggers": ["<|tool_calls_section_begin|>", "<|tool_call_begin|>"],
96                "tags": tags,
97                "at_least_one": at_least_one,
98            }
99        })
100    }
101
102    /// Create a new Kimi K2 parser
103    #[expect(
104        clippy::expect_used,
105        reason = "regex patterns are compile-time string literals"
106    )]
107    pub fn new() -> Self {
108        // Pattern for complete tool calls
109        let tool_call_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*?\})\s*<\|tool_call_end\|>";
110        let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
111
112        // Pattern for streaming (partial) tool calls
113        let stream_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*)";
114        let stream_tool_call_extractor = Regex::new(stream_pattern).expect("Valid regex pattern");
115
116        // Pattern for removing completed tool calls
117        let end_pattern = r"(?s)<\|tool_call_begin\|>.*?<\|tool_call_end\|>";
118        let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern");
119
120        // Robust parser for ids like "functions.search:0" or fallback "search:0"
121        let id_pattern = r"^(?:functions\.)?(?P<name>[\w\.]+):(?P<index>\d+)$";
122        let tool_call_id_regex = Regex::new(id_pattern).expect("Valid regex pattern");
123
124        Self {
125            tool_call_extractor,
126            stream_tool_call_extractor,
127            tool_call_end_pattern,
128            tool_call_id_regex,
129            buffer: String::new(),
130            prev_tool_call_arr: Vec::new(),
131            current_tool_id: -1,
132            current_tool_name_sent: false,
133            streamed_args_for_tool: Vec::new(),
134            last_arguments: String::new(),
135        }
136    }
137
138    /// Parse function ID to extract name and index
139    fn parse_function_id(&self, id: &str) -> Option<(String, usize)> {
140        if let Some(captures) = self.tool_call_id_regex.captures(id) {
141            let name = captures.name("name")?.as_str().to_string();
142            let index = captures.name("index")?.as_str().parse::<usize>().ok()?;
143            Some((name, index))
144        } else {
145            None
146        }
147    }
148}
149
150impl Default for KimiK2Parser {
151    fn default() -> Self {
152        Self::new()
153    }
154}
155
156#[async_trait]
157impl ToolParser for KimiK2Parser {
158    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
159        if !self.has_tool_markers(text) {
160            return Ok((text.to_string(), vec![]));
161        }
162
163        // Find where tool calls begin
164        // Safe: has_tool_markers() already confirmed the marker exists
165        let idx = text
166            .find("<|tool_calls_section_begin|>")
167            .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
168        let normal_text = text[..idx].to_string();
169
170        // Try to extract tool calls
171        let mut tools = Vec::new();
172        for captures in self.tool_call_extractor.captures_iter(text) {
173            if let (Some(id_match), Some(args_match)) = (
174                captures.name("tool_call_id"),
175                captures.name("function_arguments"),
176            ) {
177                let function_id = id_match.as_str();
178                let function_args = args_match.as_str();
179
180                // Parse function ID
181                if let Some((func_name, _index)) = self.parse_function_id(function_id) {
182                    // Try to parse JSON arguments
183                    match serde_json::from_str::<Value>(function_args) {
184                        Ok(_) => {
185                            tools.push(ToolCall {
186                                function: FunctionCall {
187                                    name: func_name,
188                                    arguments: function_args.to_string(),
189                                },
190                            });
191                        }
192                        Err(e) => {
193                            tracing::debug!(
194                                "Failed to parse JSON arguments for {}: {}",
195                                func_name,
196                                e
197                            );
198                            continue;
199                        }
200                    }
201                } else {
202                    tracing::debug!("Failed to parse function ID: {}", function_id);
203                    continue;
204                }
205            }
206        }
207
208        // If no tools were successfully parsed despite having markers, return entire text as fallback
209        if tools.is_empty() {
210            return Ok((text.to_string(), vec![]));
211        }
212
213        Ok((normal_text, tools))
214    }
215
216    async fn parse_incremental(
217        &mut self,
218        chunk: &str,
219        tools: &[Tool],
220    ) -> ParserResult<StreamingParseResult> {
221        self.buffer.push_str(chunk);
222        let current_text = &self.buffer.clone();
223
224        // Check if we have a tool call (either the start token or individual tool call)
225        let has_tool_call =
226            self.has_tool_markers(current_text) || current_text.contains("<|tool_call_begin|>");
227
228        if !has_tool_call {
229            // No tool markers detected - return all buffered content as normal text
230            let mut normal_text = std::mem::take(&mut self.buffer);
231            // Remove end tokens if present
232            for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>"] {
233                normal_text = normal_text.replace(e_token, "");
234            }
235            return Ok(StreamingParseResult {
236                normal_text,
237                calls: vec![],
238            });
239        }
240
241        // Build tool indices for validation
242        let tool_indices = helpers::get_tool_indices(tools);
243
244        let mut calls: Vec<ToolCallItem> = Vec::new();
245
246        // Try to match streaming pattern
247        if let Some(captures) = self.stream_tool_call_extractor.captures(current_text) {
248            if let (Some(id_match), Some(args_match)) = (
249                captures.name("tool_call_id"),
250                captures.name("function_arguments"),
251            ) {
252                let function_id = id_match.as_str();
253                let function_args = args_match.as_str();
254
255                // Parse function ID
256                if let Some((func_name, _index)) = self.parse_function_id(function_id) {
257                    // Validate tool name
258                    if !tool_indices.contains_key(&func_name) {
259                        // Invalid tool name - skip this tool, preserve indexing for next tool
260                        tracing::debug!("Invalid tool name '{}' - skipping", func_name);
261                        helpers::reset_current_tool_state(
262                            &mut self.buffer,
263                            &mut self.current_tool_name_sent,
264                            &mut self.streamed_args_for_tool,
265                            &self.prev_tool_call_arr,
266                        );
267                        return Ok(StreamingParseResult::default());
268                    }
269
270                    // Initialize state if this is the first tool call
271                    if self.current_tool_id == -1 {
272                        self.current_tool_id = 0;
273                        self.prev_tool_call_arr = Vec::new();
274                        self.streamed_args_for_tool = vec![String::new()];
275                    }
276
277                    // Ensure we have enough entries in our tracking arrays
278                    helpers::ensure_capacity(
279                        self.current_tool_id,
280                        &mut self.prev_tool_call_arr,
281                        &mut self.streamed_args_for_tool,
282                    );
283
284                    // Send tool name if not sent yet
285                    if self.current_tool_name_sent {
286                        // Compute incremental diff
287                        let argument_diff = if function_args.starts_with(&self.last_arguments) {
288                            &function_args[self.last_arguments.len()..]
289                        } else {
290                            function_args
291                        };
292
293                        // Split by end token before sending (like Python does)
294                        let parsed_args_diff =
295                            if let Some(pos) = argument_diff.find("<|tool_call_end|>") {
296                                &argument_diff[..pos]
297                            } else {
298                                argument_diff
299                            };
300
301                        if !parsed_args_diff.is_empty() {
302                            calls.push(ToolCallItem {
303                                tool_index: self.current_tool_id as usize,
304                                name: None,
305                                parameters: parsed_args_diff.to_string(),
306                            });
307                            // Note: Python adds full diff to _last_arguments, not just parsed part
308                            self.last_arguments.push_str(argument_diff);
309                            let tool_id = self.current_tool_id as usize;
310                            if tool_id < self.streamed_args_for_tool.len() {
311                                self.streamed_args_for_tool[tool_id].push_str(parsed_args_diff);
312                            }
313                        }
314
315                        // Check completeness - split by end token first
316                        let parsed_args = if let Some(pos) = function_args.find("<|tool_call_end|>")
317                        {
318                            &function_args[..pos]
319                        } else {
320                            function_args
321                        };
322
323                        if helpers::is_complete_json(parsed_args) {
324                            // Update the stored arguments
325                            if let Ok(parsed_args_value) =
326                                serde_json::from_str::<Value>(parsed_args)
327                            {
328                                let tool_id = self.current_tool_id as usize;
329                                if tool_id < self.prev_tool_call_arr.len() {
330                                    if let Some(obj) =
331                                        self.prev_tool_call_arr[tool_id].as_object_mut()
332                                    {
333                                        obj.insert("arguments".to_string(), parsed_args_value);
334                                    }
335                                }
336                            }
337
338                            // Find the end of the current tool call and remove only that part from buffer
339                            if let Some(mat) = self.tool_call_end_pattern.find(current_text) {
340                                // Remove the completed tool call from buffer, keep any remaining content
341                                self.buffer = current_text[mat.end()..].to_string();
342                            } else {
343                                self.buffer.clear();
344                            }
345
346                            let result = StreamingParseResult {
347                                normal_text: String::new(),
348                                calls,
349                            };
350
351                            self.current_tool_id += 1;
352                            self.last_arguments.clear();
353                            self.current_tool_name_sent = false;
354                            return Ok(result);
355                        }
356                    } else {
357                        calls.push(ToolCallItem {
358                            tool_index: self.current_tool_id as usize,
359                            name: Some(func_name.clone()),
360                            parameters: String::new(),
361                        });
362                        self.current_tool_name_sent = true;
363
364                        // Store the tool call info for serving layer completions endpoint
365                        let tool_id = self.current_tool_id as usize;
366                        if self.prev_tool_call_arr.len() <= tool_id {
367                            self.prev_tool_call_arr
368                                .resize_with(tool_id + 1, || Value::Null);
369                        }
370                        self.prev_tool_call_arr[tool_id] = serde_json::json!({
371                            "name": func_name,
372                            "arguments": {},
373                        });
374                    }
375                }
376            }
377        }
378
379        Ok(StreamingParseResult {
380            normal_text: String::new(),
381            calls,
382        })
383    }
384
385    fn has_tool_markers(&self, text: &str) -> bool {
386        text.contains("<|tool_calls_section_begin|>")
387    }
388
389    fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
390        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
391    }
392
393    fn reset(&mut self) {
394        self.buffer.clear();
395        self.prev_tool_call_arr.clear();
396        self.current_tool_id = -1;
397        self.current_tool_name_sent = false;
398        self.streamed_args_for_tool.clear();
399        self.last_arguments.clear();
400    }
401}