1use 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.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 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 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,
185 "style": "glm_xml",
186 },
187 "end": "</tool_call>",
188 })
189 })
190 .collect();
191 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 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 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}