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                        "at_least_one": at_least_one,
84                    },
85                    "end": "]",
86                }],
87                "at_least_one": at_least_one,
88            }
89        })
90    }
91
92    /// Create a new Mistral parser
93    pub fn new() -> Self {
94        Self {
95            partial_json: PartialJson::default(),
96            buffer: String::new(),
97            prev_tool_call_arr: Vec::new(),
98            current_tool_id: -1,
99            current_tool_name_sent: false,
100            streamed_args_for_tool: Vec::new(),
101            bot_token: "[TOOL_CALLS] [",
102            eot_token: "]",
103            tool_call_separator: ", ",
104            array_closed: false,
105        }
106    }
107
108    fn extract_json_array_with_pos(text: &str) -> Option<(usize, &str)> {
109        const BOT_TOKEN: &str = "[TOOL_CALLS] [";
110
111        // Find the start of the token
112        let start_idx = text.find(BOT_TOKEN)?;
113
114        // Start from the opening bracket after [TOOL_CALLS]
115        // The -1 is to include the opening bracket that's part of the token
116        let json_start = start_idx + BOT_TOKEN.len() - 1;
117
118        let mut bracket_count = 0;
119        let mut in_string = false;
120        let mut escape_next = false;
121
122        let bytes = text.as_bytes();
123
124        for i in json_start..text.len() {
125            let char = bytes[i];
126
127            if escape_next {
128                escape_next = false;
129                continue;
130            }
131
132            if char == b'\\' {
133                escape_next = true;
134                continue;
135            }
136
137            if char == b'"' && !escape_next {
138                in_string = !in_string;
139                continue;
140            }
141
142            if !in_string {
143                if char == b'[' {
144                    bracket_count += 1;
145                } else if char == b']' {
146                    bracket_count -= 1;
147                    if bracket_count == 0 {
148                        // Found the matching closing bracket
149                        return Some((start_idx, &text[json_start..=i]));
150                    }
151                }
152            }
153        }
154
155        // Incomplete array (no matching closing bracket found)
156        None
157    }
158
159    /// Parse tool calls from a JSON array
160    fn parse_json_array(json_str: &str) -> ParserResult<Vec<ToolCall>> {
161        let value: Value = serde_json::from_str(json_str)
162            .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
163
164        let mut tools = Vec::new();
165
166        if let Value::Array(arr) = value {
167            for item in &arr {
168                if let Some(tool) = Self::parse_single_object(item)? {
169                    tools.push(tool);
170                }
171            }
172        } else {
173            // Single object case (shouldn't happen with Mistral format, but handle it)
174            if let Some(tool) = Self::parse_single_object(&value)? {
175                tools.push(tool);
176            }
177        }
178
179        Ok(tools)
180    }
181
182    /// Parse a single JSON object into a ToolCall
183    fn parse_single_object(obj: &Value) -> ParserResult<Option<ToolCall>> {
184        let name = obj.get("name").and_then(|v| v.as_str());
185
186        if let Some(name) = name {
187            // Get arguments - Mistral uses "arguments" key
188            let empty_obj = Value::Object(serde_json::Map::new());
189            let args = obj.get("arguments").unwrap_or(&empty_obj);
190
191            // Convert arguments to JSON string
192            let arguments = serde_json::to_string(args)
193                .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
194
195            Ok(Some(ToolCall {
196                function: FunctionCall {
197                    name: name.to_string(),
198                    arguments,
199                },
200            }))
201        } else {
202            Ok(None)
203        }
204    }
205}
206
207impl Default for MistralParser {
208    fn default() -> Self {
209        Self::new()
210    }
211}
212
213#[async_trait]
214impl ToolParser for MistralParser {
215    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
216        // Check if text contains Mistral format
217        if !self.has_tool_markers(text) {
218            return Ok((text.to_string(), vec![]));
219        }
220
221        // Extract JSON array from Mistral format with position
222        if let Some((start_idx, json_array)) = Self::extract_json_array_with_pos(text) {
223            // Extract normal text before BOT_TOKEN
224            let normal_text_before = if start_idx > 0 {
225                text[..start_idx].to_string()
226            } else {
227                String::new()
228            };
229
230            match Self::parse_json_array(json_array) {
231                Ok(tools) => Ok((normal_text_before, tools)),
232                Err(e) => {
233                    // If JSON parsing fails, return the original text as normal text
234                    tracing::debug!("Failed to parse tool call: {}", e);
235                    Ok((text.to_string(), vec![]))
236                }
237            }
238        } else {
239            // Markers present but no complete array found
240            Ok((text.to_string(), vec![]))
241        }
242    }
243
244    async fn parse_incremental(
245        &mut self,
246        chunk: &str,
247        tools: &[Tool],
248    ) -> ParserResult<StreamingParseResult> {
249        // Append new text to buffer
250        self.buffer.push_str(chunk);
251        let current_text = &self.buffer.clone();
252
253        // Check if current_text has tool_call
254        let has_tool_start = self.has_tool_markers(current_text)
255            || (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
256
257        if !has_tool_start {
258            // Only clear buffer if we're sure no tool call is starting
259            if helpers::ends_with_partial_token(&self.buffer, self.bot_token).is_none() {
260                let mut normal_text = self.buffer.clone();
261                self.buffer.clear();
262
263                // Strip ] only once (the closing bracket of [TOOL_CALLS] array)
264                // current_tool_id > 0 means we've parsed at least one tool
265                if !self.array_closed
266                    && self.current_tool_id > 0
267                    && normal_text.starts_with(self.eot_token)
268                {
269                    // Safe: starts_with(self.eot_token) check guarantees strip_prefix succeeds
270                    if let Some(stripped) = normal_text.strip_prefix(self.eot_token) {
271                        normal_text = stripped.to_string();
272                    }
273                    self.array_closed = true;
274                }
275
276                return Ok(StreamingParseResult {
277                    normal_text,
278                    calls: vec![],
279                });
280            } else {
281                // Might be partial bot_token, keep buffering
282                return Ok(StreamingParseResult::default());
283            }
284        }
285
286        // Build tool indices
287        let tool_indices = helpers::get_tool_indices(tools);
288
289        // Determine start index for JSON parsing
290        let start_idx = if let Some(pos) = current_text.find(self.bot_token) {
291            pos + self.bot_token.len()
292        } else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
293            self.tool_call_separator.len()
294        } else {
295            0
296        };
297
298        helpers::handle_json_tool_streaming(
299            current_text,
300            start_idx,
301            &mut self.partial_json,
302            &tool_indices,
303            &mut self.buffer,
304            &mut self.current_tool_id,
305            &mut self.current_tool_name_sent,
306            &mut self.streamed_args_for_tool,
307            &mut self.prev_tool_call_arr,
308        )
309    }
310
311    fn has_tool_markers(&self, text: &str) -> bool {
312        text.contains("[TOOL_CALLS]")
313    }
314
315    fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::types::ToolCallItem>> {
316        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
317    }
318
319    fn take_unstreamed_normal_text(&mut self) -> String {
320        helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id)
321    }
322
323    fn reset(&mut self) {
324        helpers::reset_parser_state(
325            &mut self.buffer,
326            &mut self.prev_tool_call_arr,
327            &mut self.current_tool_id,
328            &mut self.current_tool_name_sent,
329            &mut self.streamed_args_for_tool,
330        );
331        self.array_closed = false;
332    }
333}
tool_parser/parsers/mistral.rs

tool_parser/parsers/
mistral.rs
