Skip to main content

tool_parser/parsers/
hy_v4.rs

1//! Hunyuan v4 tagged tool calls. Checkpoint suffixes are learned from the
2//! outer marker rather than hard-coded (e.g. `:6124c78e`).
3use async_trait::async_trait;
4use openai_protocol::common::Tool;
5use serde_json::{json, Map, Value};
6
7use crate::{
8    errors::{ParserError, ParserResult},
9    factory::GrammarStart,
10    traits::ToolParser,
11    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
12};
13
14mod streaming;
15
16const LIMIT: usize = 4 * 1024 * 1024;
17const START: &str = "<tool_calls";
18
19#[derive(Default)]
20pub struct HyV4Parser {
21    buffer: String,
22    suffix: Option<String>,
23    index: usize,
24    opening: String,
25    stream: streaming::StreamParser,
26}
27
28impl HyV4Parser {
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    fn call(body: &str, suffix: &str, tools: &[Tool]) -> Option<ToolCall> {
34        let key_start = format!("<arg_key{suffix}>");
35        let key_end = format!("</arg_key{suffix}>");
36        let val_start = format!("<arg_value{suffix}>");
37        let val_end = format!("</arg_value{suffix}>");
38        let (name, mut rest) = body
39            .find(&key_start)
40            .map_or_else(|| (body.trim(), ""), |i| (body[..i].trim(), &body[i..]));
41        if name.is_empty() || name.contains(['<', '>']) {
42            return None;
43        }
44        let props = tools
45            .iter()
46            .find(|t| t.function.name == name)
47            .and_then(|t| t.function.parameters_schema().get("properties"));
48        let mut args = Map::new();
49        while !rest.trim().is_empty() {
50            rest = rest.trim_start().strip_prefix(&key_start)?;
51            let (key, tail) = rest.split_once(&key_end)?;
52            let key = key.trim();
53            if key.is_empty() || args.contains_key(key) {
54                return None;
55            }
56            rest = tail.trim_start().strip_prefix(&val_start)?;
57            let (value, tail) = rest.split_once(&val_end)?;
58            args.insert(
59                key.to_owned(),
60                coerce(value, props.and_then(|p| p.get(key))),
61            );
62            rest = tail;
63        }
64        Some(ToolCall {
65            function: FunctionCall {
66                name: name.to_owned(),
67                arguments: Value::Object(args).to_string(),
68            },
69        })
70    }
71
72    fn drain(&mut self, tools: &[Tool]) -> StreamingParseResult {
73        let mut out = StreamingParseResult::default();
74        loop {
75            let Some(suffix) = self.suffix.clone() else {
76                if let Some(i) = self.buffer.find(START) {
77                    out.normal_text.push_str(&self.buffer[..i]);
78                    self.buffer.drain(..i);
79                    let Some(end) = self.buffer.find('>') else {
80                        break;
81                    };
82                    let suffix = &self.buffer[START.len()..end];
83                    if !suffix.is_empty()
84                        && (!suffix.starts_with(':')
85                            || suffix.len() == 1
86                            || suffix.chars().any(|c| c.is_whitespace() || c == '<'))
87                    {
88                        out.normal_text.push('<');
89                        self.buffer.drain(..1);
90                        continue;
91                    }
92                    self.suffix = Some(suffix.to_owned());
93                    self.opening = self.buffer[..end + 1].to_owned();
94                    self.buffer.drain(..end + 1);
95                    continue;
96                }
97                let keep = (1..START.len())
98                    .rev()
99                    .find(|&n| self.buffer.ends_with(&START[..n]))
100                    .unwrap_or(0);
101                let end = self.buffer.len() - keep;
102                out.normal_text.push_str(&self.buffer[..end]);
103                self.buffer.drain(..end);
104                break;
105            };
106            let close = format!("</tool_calls{suffix}>");
107            let begin = format!("<tool_call{suffix}>");
108            let end = format!("</tool_call{suffix}>");
109            let leading = self.buffer.len() - self.buffer.trim_start().len();
110            if self.buffer[leading..].starts_with(&close) {
111                self.buffer.drain(..leading + close.len());
112                self.opening.clear();
113                self.suffix = None;
114                continue;
115            }
116            if !self.buffer[leading..].starts_with(&begin) {
117                if begin.starts_with(&self.buffer[leading..])
118                    || close.starts_with(&self.buffer[leading..])
119                {
120                    break;
121                }
122                out.normal_text.push_str(&std::mem::take(&mut self.opening));
123                out.normal_text.push_str(&self.buffer);
124                self.buffer.clear();
125                self.suffix = None;
126                break;
127            }
128            let body_start = leading + begin.len();
129            // A closing-call literal inside an argument value is data, not framing.
130            let value_begin = format!("<arg_value{suffix}>");
131            let value_end = format!("</arg_value{suffix}>");
132            let mut cursor = body_start;
133            let end_index = loop {
134                let Some(rel) = self.buffer[cursor..].find(&end) else {
135                    break None;
136                };
137                let candidate = cursor + rel;
138                if let Some(v) = self.buffer[cursor..candidate].find(&value_begin) {
139                    let value_pos = cursor + v + value_begin.len();
140                    let Some(vend) = self.buffer[value_pos..].find(&value_end) else {
141                        break None;
142                    };
143                    cursor = value_pos + vend + value_end.len();
144                } else {
145                    break Some(candidate);
146                }
147            };
148            let Some(i) = end_index else {
149                break;
150            };
151            if let Some(call) = Self::call(&self.buffer[body_start..i], &suffix, tools) {
152                self.opening.clear();
153                out.calls.push(ToolCallItem {
154                    tool_index: self.index,
155                    name: Some(call.function.name),
156                    parameters: call.function.arguments,
157                });
158                self.index += 1;
159            } else {
160                out.normal_text.push_str(&std::mem::take(&mut self.opening));
161                out.normal_text.push_str(&self.buffer[..i + end.len()]);
162                self.suffix = None;
163            }
164            self.buffer.drain(..i + end.len());
165        }
166        out
167    }
168
169    /// Constrain a `<tool_calls>` block containing one or more calls. Arguments
170    /// use xgrammar's `glm_xml` style: raw strings and schema-typed JSON values,
171    /// with required/optional properties handled by xgrammar as for GLM-4.7.
172    /// Uses suffix-less markers; this builder has no tokenizer information
173    /// from which to select checkpoint-specific suffixes.
174    pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
175        let calls: Vec<Value> = tools
176            .iter()
177            .filter(|tool| !tool.function.name.is_empty())
178            .map(|tool| {
179                json!({
180                    "type": "tag",
181                    "begin": format!("<tool_call>{}", tool.function.name),
182                    "content": {
183                        "type": "json_schema",
184                        "json_schema": tool.function.parameters_schema(),
185                        "style": "glm_xml",
186                    },
187                    "end": "</tool_call>",
188                })
189            })
190            .collect();
191        // An empty `or` is unsatisfiable; with no named tool, admit only an
192        // empty block, which parses to no calls.
193        let content = if calls.is_empty() {
194            json!({ "type": "const_string", "value": "" })
195        } else {
196            json!({ "type": "plus", "content": { "type": "or", "elements": calls } })
197        };
198
199        json!({
200            "format": {
201                "type": "triggered_tags",
202                "triggers": ["<tool_calls>"],
203                "tags": [{
204                    "begin": "<tool_calls>",
205                    "content": content,
206                    "end": "</tool_calls>",
207                }],
208                "at_least_one": at_least_one,
209            }
210        })
211    }
212
213    /// What precedes a forced call on a thinking prompt, by where the engine
214    /// starts the grammar, as `Glm4MoeParser::reasoning_prefix`: nothing
215    /// behind a reasoning parser's gate, the thought block closed by
216    /// `</think>` where the grammar runs from the first token, and free text
217    /// with nothing owed at its end where the worker does not say. The
218    /// excluded markers are spelled without their closing bracket so the
219    /// checkpoint-suffixed spellings are excluded with them.
220    pub fn reasoning_prefix(start: GrammarStart) -> Option<Value> {
221        const CALL_MARKERS: [&str; 6] = [
222            "<tool_call",
223            "</tool_call",
224            "<arg_key",
225            "</arg_key",
226            "<arg_value",
227            "</arg_value",
228        ];
229        match start {
230            GrammarStart::AfterReasoning => None,
231            GrammarStart::Unknown => Some(json!({
232                "type": "any_text",
233                "excludes": CALL_MARKERS,
234            })),
235            GrammarStart::FirstToken => {
236                let excludes = [&["<think", "</think"][..], &CALL_MARKERS[..]].concat();
237                Some(json!({
238                    "type": "tag",
239                    "begin": "",
240                    "content": { "type": "any_text", "excludes": excludes },
241                    "end": "</think>",
242                }))
243            }
244        }
245    }
246}
247
248fn types(schema: Option<&Value>) -> Vec<&str> {
249    let Some(s) = schema else {
250        return vec!["string"];
251    };
252    if let Some(t) = s.get("type") {
253        if let Some(t) = t.as_str() {
254            return vec![t];
255        }
256        if let Some(ts) = t.as_array() {
257            return ts.iter().filter_map(Value::as_str).collect();
258        }
259    }
260    if let Some(v) = s
261        .get("anyOf")
262        .or_else(|| s.get("oneOf"))
263        .and_then(Value::as_array)
264    {
265        return v.iter().flat_map(|s| types(Some(s))).collect();
266    }
267    vec!["string"]
268}
269
270fn coerce(raw: &str, schema: Option<&Value>) -> Value {
271    let ts = types(schema);
272    if ts.contains(&"boolean") {
273        match raw.trim().to_ascii_lowercase().as_str() {
274            "true" => return Value::Bool(true),
275            "false" => return Value::Bool(false),
276            _ => {}
277        }
278    }
279    if let Ok(v) = serde_json::from_str::<Value>(raw) {
280        let matches = ts.iter().any(|t| match *t {
281            "integer" | "int" => v.is_i64() || v.is_u64(),
282            "number" | "float" | "double" => v.is_number(),
283            "object" | "dict" | "map" => v.is_object(),
284            "array" | "list" => v.is_array(),
285            "null" => v.is_null(),
286            _ => false,
287        });
288        if matches {
289            return v;
290        }
291    }
292    // Strings remain literal: preserve whitespace, quotes and XML-like text.
293    Value::String(raw.to_owned())
294}
295
296#[async_trait]
297impl ToolParser for HyV4Parser {
298    async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec<ToolCall>)> {
299        self.parse_complete_with_tools(output, &[]).await
300    }
301    async fn parse_complete_with_tools(
302        &self,
303        output: &str,
304        tools: &[Tool],
305    ) -> ParserResult<(String, Vec<ToolCall>)> {
306        let mut p = Self::new();
307        if output.len() > LIMIT {
308            return Err(ParserError::ParsingFailed(
309                "Hy4 tool buffer exceeds 4 MiB".into(),
310            ));
311        }
312        p.buffer.push_str(output);
313        let mut result = p.drain(tools);
314        result
315            .normal_text
316            .push_str(&format!("{}{}", p.opening, p.buffer));
317        Ok((
318            result.normal_text,
319            result
320                .calls
321                .into_iter()
322                .map(|c| ToolCall {
323                    function: FunctionCall {
324                        name: c.name.unwrap_or_default(),
325                        arguments: c.parameters,
326                    },
327                })
328                .collect(),
329        ))
330    }
331    async fn parse_incremental(
332        &mut self,
333        chunk: &str,
334        tools: &[Tool],
335    ) -> ParserResult<StreamingParseResult> {
336        self.stream.parse(chunk, tools)
337    }
338
339    fn has_tool_markers(&self, text: &str) -> bool {
340        text.contains(START)
341    }
342    fn take_unstreamed_normal_text(&mut self) -> String {
343        self.stream.flush_text()
344    }
345
346    fn reset(&mut self) {
347        *self = Self::new();
348    }
349}