tool_parser/parsers/
mistral.rs

1use async_trait::async_trait;
2use openai_protocol::common::Tool;
3use serde_json::Value;
4
5use crate::{
6    errors::{ParserError, ParserResult},
7    parsers::helpers,
8    partial_json::PartialJson,
9    traits::ToolParser,
10    types::{FunctionCall, StreamingParseResult, ToolCall},
11};
12
13/// Mistral format parser for tool calls
14///
15/// Handles the Mistral-specific format:
16/// `[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`
17///
18/// Reference: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3?chat_template=default
19pub struct MistralParser {
20    /// Parser for handling incomplete JSON during streaming
21    partial_json: PartialJson,
22
23    /// Buffer for accumulating incomplete patterns across chunks
24    buffer: String,
25
26    /// Stores complete tool call info (name and arguments) for each tool being parsed
27    prev_tool_call_arr: Vec<Value>,
28
29    /// Index of currently streaming tool call (-1 means no active tool)
30    current_tool_id: i32,
31
32    /// Flag for whether current tool's name has been sent to client
33    current_tool_name_sent: bool,
34
35    /// Tracks raw JSON string content streamed to client for each tool's arguments
36    streamed_args_for_tool: Vec<String>,
37
38    /// Token configuration
39    bot_token: &'static str,
40    eot_token: &'static str,
41    tool_call_separator: &'static str,
42
43    /// Track whether we've already stripped the closing ] bracket
44    array_closed: bool,
45}
46
47impl MistralParser {
48    /// Build structural tag for Mistral tool call format.
49    ///
50    /// Mistral outputs tool calls as a JSON array after `[TOOL_CALLS]`:
51    /// `[TOOL_CALLS] [{"name":"func1", "arguments": {...}}, {"name":"func2", "arguments": {...}}]`
52    ///
53    /// Uses `triggered_tags` with a nested `tags_with_separator` to support
54    /// parallel tool calls (multiple objects separated by `, ` in one array).
55    pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
56        let mut tags = Vec::new();
57        for tool in tools {
58            let name = &tool.function.name;
59            if name.is_empty() {
60                continue;
61            }
62            let schema = &tool.function.parameters;
63            let name_json = serde_json::to_string(name).unwrap_or_default();
64            tags.push(serde_json::json!({
65                "begin": format!("{{\"name\": {name_json}, \"arguments\": "),
66                "content": {
67                    "type": "json_schema",
68                    "json_schema": schema,
69                },
70                "end": "}",
71            }));
72        }
73        serde_json::json!({
74            "format": {
75                "type": "triggered_tags",
76                "triggers": ["[TOOL_CALLS]"],
77                "tags": [{
78                    "begin": "[TOOL_CALLS] [",
79                    "content": {
80                        "type": "tags_with_separator",
81                        "tags": tags,
82                        "separator": ", ",
83                    },
84                    "end": "]",
85                }],
86                "at_least_one": at_least_one,
87            }
88        })
89    }
90
91    /// Create a new Mistral parser
92    pub fn new() -> Self {
93        Self {
94            partial_json: PartialJson::default(),
95            buffer: String::new(),
96            prev_tool_call_arr: Vec::new(),
97            current_tool_id: -1,
98            current_tool_name_sent: false,
99            streamed_args_for_tool: Vec::new(),
100            bot_token: "[TOOL_CALLS] [",
101            eot_token: "]",
102            tool_call_separator: ", ",
103            array_closed: false,
104        }
105    }
106
107    fn extract_json_array_with_pos(text: &str) -> Option<(usize, &str)> {
108        const BOT_TOKEN: &str = "[TOOL_CALLS] [";
109
110        // Find the start of the token
111        let start_idx = text.find(BOT_TOKEN)?;
112
113        // Start from the opening bracket after [TOOL_CALLS]
114        // The -1 is to include the opening bracket that's part of the token
115        let json_start = start_idx + BOT_TOKEN.len() - 1;
116
117        let mut bracket_count = 0;
118        let mut in_string = false;
119        let mut escape_next = false;
120
121        let bytes = text.as_bytes();
122
123        for i in json_start..text.len() {
124            let char = bytes[i];
125
126            if escape_next {
127                escape_next = false;
128                continue;
129            }
130
131            if char == b'\\' {
132                escape_next = true;
133                continue;
134            }
135
136            if char == b'"' && !escape_next {
137                in_string = !in_string;
138                continue;
139            }
140
141            if !in_string {
142                if char == b'[' {
143                    bracket_count += 1;
144                } else if char == b']' {
145                    bracket_count -= 1;
146                    if bracket_count == 0 {
147                        // Found the matching closing bracket
148                        return Some((start_idx, &text[json_start..=i]));
149                    }
150                }
151            }
152        }
153
154        // Incomplete array (no matching closing bracket found)
155        None
156    }
157
158    /// Parse tool calls from a JSON array
159    fn parse_json_array(json_str: &str) -> ParserResult<Vec<ToolCall>> {
160        let value: Value = serde_json::from_str(json_str)
161            .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
162
163        let mut tools = Vec::new();
164
165        if let Value::Array(arr) = value {
166            for item in &arr {
167                if let Some(tool) = Self::parse_single_object(item)? {
168                    tools.push(tool);
169                }
170            }
171        } else {
172            // Single object case (shouldn't happen with Mistral format, but handle it)
173            if let Some(tool) = Self::parse_single_object(&value)? {
174                tools.push(tool);
175            }
176        }
177
178        Ok(tools)
179    }
180
181    /// Parse a single JSON object into a ToolCall
182    fn parse_single_object(obj: &Value) -> ParserResult<Option<ToolCall>> {
183        let name = obj.get("name").and_then(|v| v.as_str());
184
185        if let Some(name) = name {
186            // Get arguments - Mistral uses "arguments" key
187            let empty_obj = Value::Object(serde_json::Map::new());
188            let args = obj.get("arguments").unwrap_or(&empty_obj);
189
190            // Convert arguments to JSON string
191            let arguments = serde_json::to_string(args)
192                .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
193
194            Ok(Some(ToolCall {
195                function: FunctionCall {
196                    name: name.to_string(),
197                    arguments,
198                },
199            }))
200        } else {
201            Ok(None)
202        }
203    }
204}
205
206impl Default for MistralParser {
207    fn default() -> Self {
208        Self::new()
209    }
210}
211
212#[async_trait]
213impl ToolParser for MistralParser {
214    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
215        // Check if text contains Mistral format
216        if !self.has_tool_markers(text) {
217            return Ok((text.to_string(), vec![]));
218        }
219
220        // Extract JSON array from Mistral format with position
221        if let Some((start_idx, json_array)) = Self::extract_json_array_with_pos(text) {
222            // Extract normal text before BOT_TOKEN
223            let normal_text_before = if start_idx > 0 {
224                text[..start_idx].to_string()
225            } else {
226                String::new()
227            };
228
229            match Self::parse_json_array(json_array) {
230                Ok(tools) => Ok((normal_text_before, tools)),
231                Err(e) => {
232                    // If JSON parsing fails, return the original text as normal text
233                    tracing::debug!("Failed to parse tool call: {}", e);
234                    Ok((text.to_string(), vec![]))
235                }
236            }
237        } else {
238            // Markers present but no complete array found
239            Ok((text.to_string(), vec![]))
240        }
241    }
242
243    async fn parse_incremental(
244        &mut self,
245        chunk: &str,
246        tools: &[Tool],
247    ) -> ParserResult<StreamingParseResult> {
248        // Append new text to buffer
249        self.buffer.push_str(chunk);
250        let current_text = &self.buffer.clone();
251
252        // Check if current_text has tool_call
253        let has_tool_start = self.has_tool_markers(current_text)
254            || (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
255
256        if !has_tool_start {
257            // Only clear buffer if we're sure no tool call is starting
258            if helpers::ends_with_partial_token(&self.buffer, self.bot_token).is_none() {
259                let mut normal_text = self.buffer.clone();
260                self.buffer.clear();
261
262                // Strip ] only once (the closing bracket of [TOOL_CALLS] array)
263                // current_tool_id > 0 means we've parsed at least one tool
264                if !self.array_closed
265                    && self.current_tool_id > 0
266                    && normal_text.starts_with(self.eot_token)
267                {
268                    // Safe: starts_with(self.eot_token) check guarantees strip_prefix succeeds
269                    if let Some(stripped) = normal_text.strip_prefix(self.eot_token) {
270                        normal_text = stripped.to_string();
271                    }
272                    self.array_closed = true;
273                }
274
275                return Ok(StreamingParseResult {
276                    normal_text,
277                    calls: vec![],
278                });
279            } else {
280                // Might be partial bot_token, keep buffering
281                return Ok(StreamingParseResult::default());
282            }
283        }
284
285        // Build tool indices
286        let tool_indices = helpers::get_tool_indices(tools);
287
288        // Determine start index for JSON parsing
289        let start_idx = if let Some(pos) = current_text.find(self.bot_token) {
290            pos + self.bot_token.len()
291        } else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
292            self.tool_call_separator.len()
293        } else {
294            0
295        };
296
297        helpers::handle_json_tool_streaming(
298            current_text,
299            start_idx,
300            &mut self.partial_json,
301            &tool_indices,
302            &mut self.buffer,
303            &mut self.current_tool_id,
304            &mut self.current_tool_name_sent,
305            &mut self.streamed_args_for_tool,
306            &mut self.prev_tool_call_arr,
307        )
308    }
309
310    fn has_tool_markers(&self, text: &str) -> bool {
311        text.contains("[TOOL_CALLS]")
312    }
313
314    fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::types::ToolCallItem>> {
315        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
316    }
317
318    fn reset(&mut self) {
319        helpers::reset_parser_state(
320            &mut self.buffer,
321            &mut self.prev_tool_call_arr,
322            &mut self.current_tool_id,
323            &mut self.current_tool_name_sent,
324            &mut self.streamed_args_for_tool,
325        );
326        self.array_closed = false;
327    }
328}
tool_parser/parsers/mistral.rs

tool_parser/parsers/
mistral.rs
