tool_parser/parsers/
inkling.rs

1use std::collections::HashSet;
2
3use async_trait::async_trait;
4use openai_protocol::common::Tool;
5use serde_json::Value;
6
7use crate::{
8    errors::ParserResult,
9    partial_json::PartialJson,
10    traits::ToolParser,
11    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
12};
13
14const TOOL_CALL_JSON_START: &str = "<|content_invoke_tool_json|>";
15const TOOL_CALL_TEXT_START: &str = "<|content_invoke_tool_text|>";
16const END_MESSAGE: &str = "<|end_message|>";
17const MESSAGE_MODEL: &str = "<|message_model|>";
18const CONTENT_TEXT: &str = "<|content_text|>";
19const CONTENT_THINKING: &str = "<|content_thinking|>";
20const MODEL_END_SAMPLING: &str = "<|content_model_end_sampling|>";
21
22const STREAM_CONTROL_TOKENS: [&str; 7] = [
23    TOOL_CALL_JSON_START,
24    TOOL_CALL_TEXT_START,
25    MESSAGE_MODEL,
26    CONTENT_TEXT,
27    CONTENT_THINKING,
28    END_MESSAGE,
29    MODEL_END_SAMPLING,
30];
31
32const HEADER_CONTROL_TOKENS: [&str; 7] = [
33    TOOL_CALL_JSON_START,
34    TOOL_CALL_TEXT_START,
35    MESSAGE_MODEL,
36    CONTENT_TEXT,
37    CONTENT_THINKING,
38    END_MESSAGE,
39    MODEL_END_SAMPLING,
40];
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43enum StreamingState {
44    Text,
45    MessageHeader,
46    ToolCallJson,
47    ToolCallEnd,
48    DiscardToolCall,
49}
50
51/// Parser for Inkling's TML tool-call format.
52///
53/// A call is emitted as:
54/// `<|message_model|>name<|content_invoke_tool_json|>{"name":"...",`
55/// `"args":{...}}<|end_message|>`.
56/// TML response framing tokens around ordinary assistant text are removed before
57/// that text is returned to the OpenAI-compatible API. Headerless text-mode
58/// invocations cannot be represented losslessly by the OpenAI function-call
59/// shape, so their frame is safely suppressed instead of becoming answer text.
60pub struct InklingParser {
61    partial_json: PartialJson,
62    buffer: String,
63    state: StreamingState,
64    current_tool_index: usize,
65    current_tool_name: Option<String>,
66}
67
68impl InklingParser {
69    pub fn new() -> Self {
70        Self {
71            partial_json: PartialJson::default(),
72            buffer: String::new(),
73            state: StreamingState::Text,
74            current_tool_index: 0,
75            current_tool_name: None,
76        }
77    }
78
79    /// Build an xgrammar structural tag that constrains the JSON arguments for
80    /// each declared tool while retaining Inkling's native TML framing.
81    pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
82        let tags = tools
83            .iter()
84            .filter(|tool| !tool.function.name.is_empty())
85            .map(|tool| {
86                let name = serde_json::to_string(&tool.function.name).unwrap_or_default();
87                serde_json::json!({
88                    "begin": format!(
89                        "{TOOL_CALL_JSON_START}{{\"name\":{name},\"args\":"
90                    ),
91                    "content": {
92                        "type": "json_schema",
93                        "json_schema": &tool.function.parameters,
94                    },
95                    "end": format!("}}{END_MESSAGE}"),
96                })
97            })
98            .collect::<Vec<_>>();
99
100        serde_json::json!({
101            "format": {
102                "type": "triggered_tags",
103                "triggers": [TOOL_CALL_JSON_START],
104                "tags": tags,
105                "at_least_one": at_least_one,
106            }
107        })
108    }
109
110    fn clean_normal_text(text: &str) -> String {
111        let controls = [
112            MESSAGE_MODEL,
113            CONTENT_TEXT,
114            CONTENT_THINKING,
115            END_MESSAGE,
116            MODEL_END_SAMPLING,
117        ];
118        let mut output = String::new();
119        let mut cursor = 0;
120        let mut in_header = false;
121
122        while cursor < text.len() {
123            let remaining = &text[cursor..];
124            let Some((start, token)) = find_earliest_token(remaining, &controls) else {
125                if !in_header {
126                    output.push_str(remaining);
127                }
128                break;
129            };
130
131            if !in_header {
132                output.push_str(&remaining[..start]);
133            }
134            cursor += start + token.len();
135            in_header = token == MESSAGE_MODEL;
136        }
137
138        output
139    }
140
141    fn parse_complete_impl(
142        text: &str,
143        allowed_tools: Option<&HashSet<&str>>,
144    ) -> (String, Vec<ToolCall>) {
145        let tool_markers = [TOOL_CALL_JSON_START, TOOL_CALL_TEXT_START];
146        let mut normal_text = String::new();
147        let mut calls = Vec::new();
148        let mut cursor = 0;
149
150        while let Some((relative_start, marker)) =
151            find_earliest_token(&text[cursor..], &tool_markers)
152        {
153            let marker_start = cursor + relative_start;
154            normal_text.push_str(&Self::clean_normal_text(&text[cursor..marker_start]));
155            let payload_start = marker_start + marker.len();
156
157            if marker == TOOL_CALL_TEXT_START {
158                tracing::warn!(
159                    "Ignoring TML text-mode tool invocation; OpenAI tool calls require structured JSON"
160                );
161                cursor = after_discarded_tool_frame(text, payload_start);
162                continue;
163            }
164
165            let payload = &text[payload_start..];
166            let whitespace = payload.len() - payload.trim_start().len();
167            let json_start = payload_start + whitespace;
168            let Some(json_len) = complete_json_object_len(&text[json_start..]) else {
169                tracing::warn!("Ignoring malformed TML JSON tool invocation");
170                cursor = after_discarded_tool_frame(text, payload_start);
171                continue;
172            };
173            let json_end = json_start + json_len;
174            let after_json = &text[json_end..];
175            let whitespace = after_json.len() - after_json.trim_start().len();
176            let end_start = json_end + whitespace;
177
178            let Some(call) = parse_tool_call_json(&text[json_start..json_end], allowed_tools)
179            else {
180                // Malformed and undefined calls are protocol data. Suppress the
181                // whole frame, then resume at the next typed message.
182                cursor = after_discarded_tool_frame(text, json_end);
183                continue;
184            };
185            calls.push(call);
186
187            if text[end_start..].starts_with(END_MESSAGE) {
188                cursor = end_start + END_MESSAGE.len();
189            } else if text[end_start..].starts_with(MODEL_END_SAMPLING) {
190                cursor = end_start + MODEL_END_SAMPLING.len();
191            } else {
192                // Match the streaming parser's permissive ToolCallEnd state:
193                // after a complete valid JSON object, a missing end marker does
194                // not turn subsequent bytes into tool payload.
195                cursor = end_start;
196            }
197        }
198
199        normal_text.push_str(&Self::clean_normal_text(&text[cursor..]));
200        (normal_text, calls)
201    }
202
203    fn process_text(&mut self, result: &mut StreamingParseResult) -> bool {
204        let state_markers = [MESSAGE_MODEL, TOOL_CALL_JSON_START, TOOL_CALL_TEXT_START];
205        if let Some((start, marker)) = find_earliest_token(&self.buffer, &state_markers) {
206            let normal_text = Self::clean_normal_text(&self.buffer[..start]);
207            result.normal_text.push_str(&normal_text);
208            self.buffer.drain(..start + marker.len());
209            self.state = if marker == MESSAGE_MODEL {
210                StreamingState::MessageHeader
211            } else if marker == TOOL_CALL_JSON_START {
212                StreamingState::ToolCallJson
213            } else {
214                tracing::warn!(
215                    "Ignoring TML text-mode tool invocation; OpenAI tool calls require structured JSON"
216                );
217                StreamingState::DiscardToolCall
218            };
219            return true;
220        }
221
222        let keep = longest_partial_token_suffix(&self.buffer, &STREAM_CONTROL_TOKENS);
223        let safe_len = self.buffer.len() - keep;
224        if safe_len > 0 {
225            let safe_text = self.buffer[..safe_len].to_string();
226            self.buffer.drain(..safe_len);
227            result
228                .normal_text
229                .push_str(&Self::clean_normal_text(&safe_text));
230        }
231        false
232    }
233
234    fn process_message_header(&mut self) -> bool {
235        if let Some((start, token)) = find_earliest_token(&self.buffer, &HEADER_CONTROL_TOKENS) {
236            // Everything before the content-type token is an author, tool
237            // name, or channel header. It is metadata and must not be emitted
238            // as assistant content.
239            self.buffer.drain(..start + token.len());
240            self.state = if token == MESSAGE_MODEL {
241                StreamingState::MessageHeader
242            } else if token == TOOL_CALL_JSON_START {
243                StreamingState::ToolCallJson
244            } else if token == TOOL_CALL_TEXT_START {
245                tracing::warn!(
246                    "Ignoring TML text-mode tool invocation; OpenAI tool calls require structured JSON"
247                );
248                StreamingState::DiscardToolCall
249            } else {
250                StreamingState::Text
251            };
252            return true;
253        }
254
255        // Header bytes are never user-visible. Drop author/channel bytes while
256        // retaining only a possible control-token prefix across chunk splits.
257        let keep = longest_partial_token_suffix(&self.buffer, &HEADER_CONTROL_TOKENS);
258        let safe_len = self.buffer.len() - keep;
259        if safe_len > 0 {
260            self.buffer.drain(..safe_len);
261        }
262        false
263    }
264
265    fn process_tool_call(
266        &mut self,
267        allowed_tools: &HashSet<&str>,
268        result: &mut StreamingParseResult,
269    ) -> bool {
270        let whitespace = self.buffer.len() - self.buffer.trim_start().len();
271        if whitespace > 0 {
272            self.buffer.drain(..whitespace);
273        }
274        if self.buffer.is_empty() {
275            return false;
276        }
277
278        if !self.buffer.starts_with('{') {
279            self.current_tool_name = None;
280            self.state = StreamingState::DiscardToolCall;
281            return true;
282        }
283
284        if self.current_tool_name.is_none() {
285            if let Ok((Value::Object(object), _)) =
286                self.partial_json.parse_value(&self.buffer, false)
287            {
288                if let Some(name) = object.get("name").and_then(Value::as_str) {
289                    if allowed_tools.contains(name) {
290                        self.current_tool_name = Some(name.to_string());
291                        result.calls.push(ToolCallItem {
292                            tool_index: self.current_tool_index,
293                            name: Some(name.to_string()),
294                            parameters: String::new(),
295                        });
296                    } else {
297                        tracing::debug!("Inkling attempted to call undefined tool: {}", name);
298                        self.state = StreamingState::DiscardToolCall;
299                        return true;
300                    }
301                }
302            }
303        }
304
305        let Some(json_len) = complete_json_object_len(&self.buffer) else {
306            return false;
307        };
308        let json = &self.buffer[..json_len];
309        let Some(call) = parse_tool_call_json(json, Some(allowed_tools)) else {
310            self.current_tool_name = None;
311            self.state = StreamingState::DiscardToolCall;
312            return true;
313        };
314
315        // A complete object may expose the name and arguments in the same
316        // chunk. Emit the name first in that case, as required by stream APIs.
317        if self.current_tool_name.is_none() {
318            result.calls.push(ToolCallItem {
319                tool_index: self.current_tool_index,
320                name: Some(call.function.name.clone()),
321                parameters: String::new(),
322            });
323        }
324        result.calls.push(ToolCallItem {
325            tool_index: self.current_tool_index,
326            name: None,
327            parameters: call.function.arguments,
328        });
329
330        self.buffer.drain(..json_len);
331        self.current_tool_index += 1;
332        self.current_tool_name = None;
333        self.state = StreamingState::ToolCallEnd;
334        true
335    }
336
337    fn process_tool_call_end(&mut self) -> bool {
338        let whitespace = self.buffer.len() - self.buffer.trim_start().len();
339        if whitespace > 0 {
340            self.buffer.drain(..whitespace);
341        }
342        if self.buffer.is_empty() {
343            return false;
344        }
345        if self.buffer.starts_with(END_MESSAGE) {
346            self.buffer.drain(..END_MESSAGE.len());
347            self.state = StreamingState::Text;
348            return true;
349        }
350        if END_MESSAGE.starts_with(&self.buffer) {
351            return false;
352        }
353
354        // Be permissive when a backend omits the end token after an otherwise
355        // complete call. The remaining bytes still need normal text handling.
356        self.state = StreamingState::Text;
357        true
358    }
359
360    fn discard_tool_call(&mut self) -> bool {
361        let end_tokens = [END_MESSAGE, MODEL_END_SAMPLING];
362        if let Some((end, token)) = find_earliest_token(&self.buffer, &end_tokens) {
363            self.buffer.drain(..end + token.len());
364            self.current_tool_name = None;
365            self.state = StreamingState::Text;
366            return true;
367        }
368
369        let keep = longest_partial_token_suffix(&self.buffer, &end_tokens);
370        let safe_len = self.buffer.len() - keep;
371        if safe_len > 0 {
372            self.buffer.drain(..safe_len);
373        }
374        false
375    }
376}
377
378impl Default for InklingParser {
379    fn default() -> Self {
380        Self::new()
381    }
382}
383
384#[async_trait]
385impl ToolParser for InklingParser {
386    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
387        Ok(Self::parse_complete_impl(text, None))
388    }
389
390    async fn parse_complete_with_tools(
391        &self,
392        text: &str,
393        tools: &[Tool],
394    ) -> ParserResult<(String, Vec<ToolCall>)> {
395        let allowed_tools = tools
396            .iter()
397            .map(|tool| tool.function.name.as_str())
398            .collect::<HashSet<_>>();
399        Ok(Self::parse_complete_impl(text, Some(&allowed_tools)))
400    }
401
402    async fn parse_incremental(
403        &mut self,
404        chunk: &str,
405        tools: &[Tool],
406    ) -> ParserResult<StreamingParseResult> {
407        self.buffer.push_str(chunk);
408        let allowed_tools = tools
409            .iter()
410            .map(|tool| tool.function.name.as_str())
411            .collect::<HashSet<_>>();
412        let mut result = StreamingParseResult::default();
413
414        loop {
415            let progressed = match self.state {
416                StreamingState::Text => self.process_text(&mut result),
417                StreamingState::MessageHeader => self.process_message_header(),
418                StreamingState::ToolCallJson => self.process_tool_call(&allowed_tools, &mut result),
419                StreamingState::ToolCallEnd => self.process_tool_call_end(),
420                StreamingState::DiscardToolCall => self.discard_tool_call(),
421            };
422            if !progressed {
423                break;
424            }
425        }
426
427        Ok(result)
428    }
429
430    fn has_tool_markers(&self, text: &str) -> bool {
431        text.contains(TOOL_CALL_JSON_START) || text.contains(TOOL_CALL_TEXT_START)
432    }
433
434    fn reset(&mut self) {
435        self.buffer.clear();
436        self.state = StreamingState::Text;
437        self.current_tool_index = 0;
438        self.current_tool_name = None;
439    }
440}
441
442fn parse_tool_call_json(json: &str, allowed_tools: Option<&HashSet<&str>>) -> Option<ToolCall> {
443    let value = serde_json::from_str::<Value>(json).ok()?;
444    let object = value.as_object()?;
445    let name = object.get("name")?.as_str()?;
446    let args = object.get("args")?.as_object()?;
447
448    if allowed_tools.is_some_and(|tools| !tools.contains(name)) {
449        tracing::debug!("Inkling attempted to call undefined tool: {}", name);
450        return None;
451    }
452
453    Some(ToolCall {
454        function: FunctionCall {
455            name: name.to_string(),
456            arguments: serde_json::to_string(args).ok()?,
457        },
458    })
459}
460
461/// Find the earliest complete protocol token, preferring the token list order
462/// when two entries start at the same byte offset.
463fn find_earliest_token<'a>(text: &str, tokens: &'a [&'a str]) -> Option<(usize, &'a str)> {
464    tokens
465        .iter()
466        .filter_map(|token| text.find(token).map(|start| (start, *token)))
467        .min_by_key(|(start, _)| *start)
468}
469
470/// Return the first byte after a discarded tool frame. Both end markers match
471/// the streaming parser's `DiscardToolCall` recovery behavior. If a malformed
472/// frame is unterminated, all remaining bytes belong to that protocol frame.
473fn after_discarded_tool_frame(text: &str, payload_start: usize) -> usize {
474    let end_tokens = [END_MESSAGE, MODEL_END_SAMPLING];
475    find_earliest_token(&text[payload_start..], &end_tokens)
476        .map_or(text.len(), |(relative_end, token)| {
477            payload_start + relative_end + token.len()
478        })
479}
480
481/// Return the byte length of one complete top-level JSON object. Braces inside
482/// quoted strings are ignored, and trailing data is left for the caller.
483fn complete_json_object_len(text: &str) -> Option<usize> {
484    if !text.starts_with('{') {
485        return None;
486    }
487
488    let mut depth = 0usize;
489    let mut in_string = false;
490    let mut escaped = false;
491
492    for (index, byte) in text.bytes().enumerate() {
493        if in_string {
494            if escaped {
495                escaped = false;
496            } else if byte == b'\\' {
497                escaped = true;
498            } else if byte == b'"' {
499                in_string = false;
500            }
501            continue;
502        }
503
504        match byte {
505            b'"' => in_string = true,
506            b'{' => depth += 1,
507            b'}' => {
508                depth = depth.checked_sub(1)?;
509                if depth == 0 {
510                    return Some(index + 1);
511                }
512            }
513            _ => {}
514        }
515    }
516    None
517}
518
519fn longest_partial_token_suffix(buffer: &str, tokens: &[&str]) -> usize {
520    tokens
521        .iter()
522        .flat_map(|token| {
523            token
524                .char_indices()
525                .skip(1)
526                .map(move |(index, _)| &token[..index])
527        })
528        .filter(|prefix| buffer.ends_with(prefix))
529        .map(str::len)
530        .max()
531        .unwrap_or(0)
532}
tool_parser/parsers/inkling.rs

tool_parser/parsers/
inkling.rs
