rlx_text/
tool_parse.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Parsing model-emitted tool / function calls into structured form.
17//!
18//! Models emit tool calls in a few conventions; this module recognizes the
19//! common ones and returns `{name, arguments}` (mirroring mlx-lm's
20//! `tool_parsers/`). All host-side string work — no model/runtime coupling.
21//!
22//! Formats:
23//!   - **Hermes / Qwen** — `<tool_call>{"name": …, "arguments": {…}}</tool_call>`
24//!     (possibly several blocks).
25//!   - **JSON** — the whole text is an object (or array of objects) with
26//!     `name` + `arguments`/`parameters`.
27//!   - **Pythonic** — `[fn(a="x", b=2)]` (optionally wrapped in tags).
28
29use serde_json::{Map, Value};
30
31/// A parsed tool call: function name + a JSON object of arguments.
32#[derive(Debug, Clone, PartialEq)]
33pub struct ToolCall {
34    pub name: String,
35    pub arguments: Value,
36}
37
38/// Known tool-call output formats.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum ToolFormat {
41    /// `<tool_call>{json}</tool_call>` blocks (Qwen3, Hermes, …).
42    Hermes,
43    /// Bare JSON object or array of `{name, arguments|parameters}`.
44    Json,
45    /// `[fn(arg="v", n=2)]` (Llama/`pythonic`).
46    Pythonic,
47}
48
49/// Parse tool calls from `text` in a known `fmt`. Returns an empty vec when
50/// nothing parses (a plain assistant message, not a tool call).
51pub fn parse(text: &str, fmt: ToolFormat) -> Vec<ToolCall> {
52    match fmt {
53        ToolFormat::Hermes => parse_hermes(text),
54        ToolFormat::Json => parse_json(text.trim()),
55        ToolFormat::Pythonic => parse_pythonic(text),
56    }
57}
58
59/// Try each known format, in order of specificity, and return the first that
60/// yields calls. Useful when the chat template's format isn't known up front.
61pub fn detect_and_parse(text: &str) -> Vec<ToolCall> {
62    for fmt in [ToolFormat::Hermes, ToolFormat::Pythonic, ToolFormat::Json] {
63        let calls = parse(text, fmt);
64        if !calls.is_empty() {
65            return calls;
66        }
67    }
68    Vec::new()
69}
70
71/// Normalize a JSON object that should contain `name` + arguments into a
72/// [`ToolCall`]. Accepts `arguments` or `parameters`; arguments may be an
73/// object or a JSON-encoded string.
74fn tool_call_from_obj(obj: &Map<String, Value>) -> Option<ToolCall> {
75    let name = obj.get("name")?.as_str()?.to_string();
76    let args = obj
77        .get("arguments")
78        .or_else(|| obj.get("parameters"))
79        .cloned()
80        .unwrap_or_else(|| Value::Object(Map::new()));
81    // Some models double-encode arguments as a JSON string.
82    let args = match args {
83        Value::String(s) => serde_json::from_str(&s).unwrap_or(Value::String(s)),
84        other => other,
85    };
86    Some(ToolCall {
87        name,
88        arguments: args,
89    })
90}
91
92fn parse_json(text: &str) -> Vec<ToolCall> {
93    match serde_json::from_str::<Value>(text) {
94        Ok(Value::Object(o)) => tool_call_from_obj(&o).into_iter().collect(),
95        Ok(Value::Array(a)) => a
96            .iter()
97            .filter_map(|v| v.as_object().and_then(tool_call_from_obj))
98            .collect(),
99        _ => Vec::new(),
100    }
101}
102
103fn parse_hermes(text: &str) -> Vec<ToolCall> {
104    const OPEN: &str = "<tool_call>";
105    const CLOSE: &str = "</tool_call>";
106    let mut out = Vec::new();
107    let mut rest = text;
108    while let Some(start) = rest.find(OPEN) {
109        let after = &rest[start + OPEN.len()..];
110        let Some(end) = after.find(CLOSE) else { break };
111        let inner = after[..end].trim();
112        out.extend(parse_json(inner));
113        rest = &after[end + CLOSE.len()..];
114    }
115    out
116}
117
118fn parse_pythonic(text: &str) -> Vec<ToolCall> {
119    // Locate the innermost `[ ... ]` payload (ignore any tag wrappers).
120    let Some(open) = text.find('[') else {
121        return Vec::new();
122    };
123    let Some(close_rel) = text[open..].rfind(']') else {
124        return Vec::new();
125    };
126    let body = &text[open + 1..open + close_rel];
127    let mut out = Vec::new();
128    for call in split_top_level(body, ',') {
129        let call = call.trim();
130        let Some(paren) = call.find('(') else {
131            continue;
132        };
133        if !call.ends_with(')') {
134            continue;
135        }
136        let name = call[..paren].trim();
137        if name.is_empty() || !name.chars().all(|c| c.is_alphanumeric() || c == '_') {
138            continue;
139        }
140        let args_str = &call[paren + 1..call.len() - 1];
141        let mut args = Map::new();
142        for pair in split_top_level(args_str, ',') {
143            let pair = pair.trim();
144            if pair.is_empty() {
145                continue;
146            }
147            if let Some(eq) = pair.find('=') {
148                let key = pair[..eq].trim().to_string();
149                let val = pair[eq + 1..].trim();
150                args.insert(key, literal_to_value(val));
151            }
152        }
153        out.push(ToolCall {
154            name: name.to_string(),
155            arguments: Value::Object(args),
156        });
157    }
158    out
159}
160
161/// Convert a Python-ish literal to a JSON value: try JSON first (numbers,
162/// `true`/`false`, quoted strings), then strip simple quotes, else keep raw.
163fn literal_to_value(s: &str) -> Value {
164    if let Ok(v) = serde_json::from_str::<Value>(s) {
165        return v;
166    }
167    // Python single-quoted string → strip quotes.
168    if s.len() >= 2 && s.starts_with('\'') && s.ends_with('\'') {
169        return Value::String(s[1..s.len() - 1].to_string());
170    }
171    match s {
172        "True" => Value::Bool(true),
173        "False" => Value::Bool(false),
174        "None" => Value::Null,
175        other => Value::String(other.to_string()),
176    }
177}
178
179/// Split `s` on `sep`, but only at the top nesting level (respecting `()`,
180/// `[]`, `{}`, and quoted strings). Avoids splitting inside argument values.
181fn split_top_level(s: &str, sep: char) -> Vec<String> {
182    let mut parts = Vec::new();
183    let mut depth = 0i32;
184    let mut quote: Option<char> = None;
185    let mut cur = String::new();
186    let mut chars = s.chars().peekable();
187    while let Some(c) = chars.next() {
188        match quote {
189            Some(q) => {
190                cur.push(c);
191                if c == '\\' {
192                    if let Some(&n) = chars.peek() {
193                        cur.push(n);
194                        chars.next();
195                    }
196                } else if c == q {
197                    quote = None;
198                }
199            }
200            None => match c {
201                '"' | '\'' => {
202                    quote = Some(c);
203                    cur.push(c);
204                }
205                '(' | '[' | '{' => {
206                    depth += 1;
207                    cur.push(c);
208                }
209                ')' | ']' | '}' => {
210                    depth -= 1;
211                    cur.push(c);
212                }
213                _ if c == sep && depth == 0 => {
214                    parts.push(std::mem::take(&mut cur));
215                }
216                _ => cur.push(c),
217            },
218        }
219    }
220    if !cur.trim().is_empty() || !parts.is_empty() {
221        parts.push(cur);
222    }
223    parts
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use serde_json::json;
230
231    #[test]
232    fn hermes_single_call() {
233        let text = r#"<tool_call>{"name":"get_weather","arguments":{"city":"Paris"}}</tool_call>"#;
234        let calls = parse(text, ToolFormat::Hermes);
235        assert_eq!(calls.len(), 1);
236        assert_eq!(calls[0].name, "get_weather");
237        assert_eq!(calls[0].arguments, json!({"city":"Paris"}));
238    }
239
240    #[test]
241    fn hermes_multiple_calls_with_surrounding_text() {
242        let text = "sure!\n<tool_call>{\"name\":\"a\",\"arguments\":{}}</tool_call>\n\
243                    <tool_call>{\"name\":\"b\",\"parameters\":{\"x\":1}}</tool_call>";
244        let calls = parse(text, ToolFormat::Hermes);
245        assert_eq!(calls.len(), 2);
246        assert_eq!(calls[0].name, "a");
247        assert_eq!(calls[1].name, "b");
248        assert_eq!(calls[1].arguments, json!({"x":1}));
249    }
250
251    #[test]
252    fn json_double_encoded_arguments() {
253        let text = r#"{"name":"f","arguments":"{\"k\":2}"}"#;
254        let calls = parse(text, ToolFormat::Json);
255        assert_eq!(calls.len(), 1);
256        assert_eq!(calls[0].arguments, json!({"k":2}));
257    }
258
259    #[test]
260    fn pythonic_call() {
261        let text = r#"[get_weather(city="Paris", days=2, metric=true)]"#;
262        let calls = parse(text, ToolFormat::Pythonic);
263        assert_eq!(calls.len(), 1);
264        assert_eq!(calls[0].name, "get_weather");
265        assert_eq!(
266            calls[0].arguments,
267            json!({"city":"Paris","days":2,"metric":true})
268        );
269    }
270
271    #[test]
272    fn pythonic_with_tag_wrapper_and_commas_in_strings() {
273        let text = r#"<|tool_call_start|>[note(text="a, b, c")]<|tool_call_end|>"#;
274        let calls = parse(text, ToolFormat::Pythonic);
275        assert_eq!(calls.len(), 1);
276        assert_eq!(calls[0].arguments, json!({"text":"a, b, c"}));
277    }
278
279    #[test]
280    fn detect_picks_the_right_format() {
281        assert_eq!(detect_and_parse("just text"), vec![]);
282        assert_eq!(
283            detect_and_parse(r#"<tool_call>{"name":"x","arguments":{}}</tool_call>"#)[0].name,
284            "x"
285        );
286        assert_eq!(detect_and_parse(r#"[y(a=1)]"#)[0].name, "y");
287    }
288}
rlx_text/tool_parse.rs

rlx_text/
tool_parse.rs
