@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / de21863419
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 08:30 UTC

[Rust Frontend] Add InternLM2 tool parser (#43481)

Signed-off-by: Will.hou <1205157517@qq.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>

Will.hou committed Jun 1, 2026 at 16:58 UTC de218634194cd5ca4335eb478fbba5246cb54dbf
10 files changed +445 -14
rust/src/chat/src/lib.rs
+1 -1
@@ -233,7 +233,7 @@ mod tests {
233 )
234 .unwrap_err();
235
236 - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
236 + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
237 }
238
239 #[test]
rust/src/chat/src/parser/tool/mod.rs
+13 -3
@@ -5,9 +5,9 @@ use std::sync::LazyLock;
5 pub use vllm_tool_parser::{
6 DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
7 Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
8 - KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
9 - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
10 - ToolParserOutput,
8 + Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
9 + MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
10 + ToolParserError, ToolParserOutput,
11 };
12
13 use crate::parser::ParserFactory;
@@ -24,6 +24,9 @@ pub mod names {
24 pub const GEMMA4: &str = "gemma4";
25 pub const HERMES: &str = "hermes";
26 pub const HY_V3: &str = "hy_v3";
27 + // Matches the Python CLI name `--tool-call-parser internlm`, which Python
28 + // also routes to `Internlm2ToolParser` despite the version-agnostic name.
29 + pub const INTERNLM: &str = "internlm";
30 pub const KIMI_K2: &str = "kimi_k2";
31 pub const LLAMA3_JSON: &str = "llama3_json";
32 pub const LLAMA4_JSON: &str = "llama4_json";
@@ -62,6 +65,7 @@ impl ToolParserFactory {
65 .register_parser::<Gemma4ToolParser>(names::GEMMA4)
66 .register_parser::<HermesToolParser>(names::HERMES)
67 .register_parser::<HyV3ToolParser>(names::HY_V3)
68 + .register_parser::<Internlm2ToolParser>(names::INTERNLM)
69 .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
70 .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
71 .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
@@ -80,6 +84,12 @@ impl ToolParserFactory {
84 .register_pattern("hermes", names::HERMES)
85 .register_pattern("hy3", names::HY_V3)
86 .register_pattern("hy_v3", names::HY_V3)
87 + // Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
88 + // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
89 + // routes to Llama), `internlm3-*` (also Llama-architecture per
90 + // vllm/model_executor/models/registry.py:146), or `Intern-S1` /
91 + // `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
92 + .register_pattern("internlm2", names::INTERNLM)
93 .register_pattern("llama-4", names::LLAMA4_JSON)
94 .register_pattern("llama-3.2", names::LLAMA3_JSON)
95 .register_pattern("llama-3.1", names::LLAMA3_JSON)
rust/src/chat/src/parser/tool/tests.rs
+29
@@ -161,4 +161,33 @@ fn factory_new_resolves_default_patterns() {
161 factory.resolve_name_for_model("org/mm-m2-base"),
162 Some(names::MINIMAX_M2)
163 );
164 +
165 + // InternLM2 positive: both dashed and underscored versioned names route.
166 + assert_eq!(
167 + factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
168 + Some(names::INTERNLM)
169 + );
170 + assert_eq!(
171 + factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
172 + Some(names::INTERNLM)
173 + );
174 +
175 + // Negative: other internlm-org models do NOT route to the InternLM2 parser,
176 + // since they use unrelated prompt formats.
177 + // - InternLM v1 (`internlm-chat-7b`) routes to Llama
178 + // - InternLM3 (`internlm3-8b-instruct`) routes to Llama
179 + // - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
180 + assert_eq!(
181 + factory.resolve_name_for_model("internlm/internlm-chat-7b"),
182 + None
183 + );
184 + assert_eq!(
185 + factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
186 + None
187 + );
188 + assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
189 + assert_eq!(
190 + factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
191 + None
192 + );
193 }
rust/src/tool-parser/src/json/hermes.rs
+1 -1
@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
8 marker_whitespace: JsonToolCallWhitespace::Optional,
9 delimiter: None,
10 name_key: "name",
11 - arguments_key: "arguments",
11 + arguments_key: &["arguments"],
12 };
13
14 /// Tool parser for Hermes XML-wrapped JSON tool calls.
rust/src/tool-parser/src/json/internlm2.rs new
+351
@@ -0,0 +1,351 @@
1 +use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
2 +use crate::{Result, Tool, ToolParser, ToolParserOutput};
3 +
4 +const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
5 + parser_name: "InternLM2",
6 + start_marker: "<|action_start|><|plugin|>",
7 + end_marker: "<|action_end|>",
8 + marker_whitespace: JsonToolCallWhitespace::Optional,
9 + delimiter: None,
10 + name_key: "name",
11 + // The Python parser's `get_arguments()` accepts either `parameters` or
12 + // `arguments` and prefers `parameters` when both are present. This Rust
13 + // parser uses first-encountered semantics because the header parser only
14 + // permits one args key per tool-call object; if a future model emits
15 + // both keys in the same object, the Rust port will accept the first one
16 + // and reject the trailing one as a syntax error rather than silently
17 + // shadowing it.
18 + arguments_key: &["parameters", "arguments"],
19 +};
20 +
21 +/// Tool parser for InternLM2 special-token wrapped JSON tool calls.
22 +///
23 +/// Example tool call content:
24 +///
25 +/// ```text
26 +/// <|action_start|><|plugin|>{"name": "get_weather", "parameters": {"location":"Tokyo"}}<|action_end|>
27 +/// ```
28 +///
29 +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
30 +/// argument deltas without schema conversion or JSON normalization.
31 +///
32 +/// # Divergences from the Python reference
33 +///
34 +/// This Rust port intentionally diverges from
35 +/// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways:
36 +///
37 +/// - **Parallel tool calls are supported.** Python silently drops every
38 +/// `<|action_start|>` block after the first (`current_tool_id > 0` returns
39 +/// an empty delta); this parser emits every well-formed block with
40 +/// incrementing `tool_index`. Models that legitimately emit multiple action
41 +/// blocks therefore produce more tool calls under Rust than under Python.
42 +/// - **End-marker bytes inside JSON string values are preserved.** Python
43 +/// does `action.split("<|action_end|>")[0]` which truncates regardless of
44 +/// JSON context; this parser scans matched braces and quotes so a literal
45 +/// `<|action_end|>` inside an arguments string is forwarded intact.
46 +/// - **Only whitespace is allowed before the `{`.** Python's non-streaming
47 +/// `action[action.find("{"):]` drops any bytes before the first `{`, but
48 +/// its streaming path has no equivalent and the model format always emits
49 +/// `<|plugin|>{...`; this parser allows only whitespace there, matching the
50 +/// other JSON parsers in this crate.
51 +/// - **Truncated tool calls error rather than silently dropping.** Python's
52 +/// streaming wrapper swallows mid-stream errors with `except Exception:
53 +/// return None` (logging a traceback) while its non-streaming path raises
54 +/// `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool
55 +/// call` error from `finish()`, matching the other JSON parsers and Python's
56 +/// non-streaming behavior.
57 +///
58 +/// # Known unaddressed divergences (TODO)
59 +///
60 +/// The following Python behaviors are NOT yet matched. They are deferred to
61 +/// follow-up work because they require non-local changes to the shared
62 +/// `JsonToolCallParser` core that would affect Hermes / Llama / Mistral /
63 +/// Qwen as well. If a real-world InternLM2 deployment hits one of these,
64 +/// prioritize the corresponding fix.
65 +///
66 +/// - **Arguments value type.** The shared core requires the arguments value
67 +/// to be a JSON object (`take_json_object` rejects anything not starting
68 +/// with `{`). Python's `json.dumps(action_dict.get("parameters", ...))`
69 +/// accepts `null`, arrays, strings, and numbers and round-trips them
70 +/// verbatim. Models that legitimately emit `"parameters":null` will hard-
71 +/// fail under Rust.
72 +/// - **Unknown arguments key.** Python falls back to `{}` via
73 +/// `action_dict.get("parameters", action_dict.get("arguments", {}))` when
74 +/// neither key is present; the Rust header parser raises
75 +/// `parsing failed: invalid InternLM2` for any unrecognized key. A model
76 +/// that emits a typo (e.g. `"params"`) breaks the whole response.
77 +/// - **Field order independence.** The header parser requires the JSON keys
78 +/// to appear in the order `name` then arguments key. Python's
79 +/// `json.loads` + `dict.get` is order-independent, so a model emitting
80 +/// `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
81 +pub struct Internlm2ToolParser {
82 + inner: JsonToolCallParser,
83 +}
84 +
85 +impl Internlm2ToolParser {
86 + /// Create an InternLM2 tool parser.
87 + fn new(_tools: &[Tool]) -> Self {
88 + Self {
89 + inner: JsonToolCallParser::new(INTERNLM2_CONFIG),
90 + }
91 + }
92 +}
93 +
94 +impl ToolParser for Internlm2ToolParser {
95 + /// Create a boxed InternLM2 tool parser.
96 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
97 + where
98 + Self: Sized + 'static,
99 + {
100 + Ok(Box::new(Self::new(tools)))
101 + }
102 +
103 + /// Preserve special-token markers while decoding, since
104 + /// `<|action_start|>`, `<|plugin|>`, and `<|action_end|>` are tokenizer
105 + /// special tokens in InternLM2 models.
106 + fn preserve_special_tokens(&self) -> bool {
107 + true
108 + }
109 +
110 + /// Feed one decoded text chunk through the InternLM2 parser.
111 + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
112 + self.inner.parse_into(chunk, output)
113 + }
114 +
115 + /// Flush any buffered partial state at end of stream.
116 + fn finish(&mut self) -> Result<ToolParserOutput> {
117 + self.inner.finish()
118 + }
119 +
120 + /// Clear parser state and return currently uncommitted buffered text.
121 + fn reset(&mut self) -> String {
122 + self.inner.reset()
123 + }
124 +}
125 +
126 +#[cfg(test)]
127 +mod tests {
128 + use expect_test::expect;
129 + use thiserror_ext::AsReport;
130 +
131 + use super::Internlm2ToolParser;
132 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
133 + use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
134 +
135 + const ACTION_START: &str = "<|action_start|><|plugin|>";
136 + const ACTION_END: &str = "<|action_end|>";
137 +
138 + fn build_tool_call(function_name: &str, args_key: &str, arguments: &str) -> String {
139 + format!(
140 + r#"{ACTION_START}{{"name":"{function_name}","{args_key}":{arguments}}}{ACTION_END}"#
141 + )
142 + }
143 +
144 + #[test]
145 + fn internlm2_parse_complete_without_tool_call_keeps_text() {
146 + let mut parser = Internlm2ToolParser::new(&test_tools());
147 + let result = parser.parse_complete("Hello, world!").unwrap();
148 +
149 + assert_eq!(result.normal_text, "Hello, world!");
150 + assert!(result.calls.is_empty());
151 + }
152 +
153 + #[test]
154 + fn internlm2_parse_complete_extracts_parameters_key() {
155 + let mut parser = Internlm2ToolParser::new(&test_tools());
156 + let arguments = r#"{"location":"Tokyo","days":"3"}"#;
157 + let result = parser
158 + .parse_complete(&format!(
159 + "Let me check.\n{}",
160 + build_tool_call("get_weather", "parameters", arguments)
161 + ))
162 + .unwrap();
163 +
164 + assert_eq!(result.normal_text, "Let me check.\n");
165 + assert_eq!(result.calls.len(), 1);
166 + assert_eq!(result.calls[0].tool_index, 0);
167 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
168 + assert_eq!(result.calls[0].arguments, arguments);
169 + }
170 +
171 + #[test]
172 + fn internlm2_parse_complete_extracts_arguments_key_fallback() {
173 + let mut parser = Internlm2ToolParser::new(&test_tools());
174 + let arguments = r#"{"location":"Tokyo"}"#;
175 + let result = parser
176 + .parse_complete(&build_tool_call("get_weather", "arguments", arguments))
177 + .unwrap();
178 +
179 + assert_eq!(result.calls.len(), 1);
180 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
181 + assert_eq!(result.calls[0].arguments, arguments);
182 + }
183 +
184 + #[test]
185 + fn internlm2_accepts_whitespace_after_plugin_marker() {
186 + let mut parser = Internlm2ToolParser::new(&test_tools());
187 + let result = parser
188 + .parse_complete(&format!(
189 + r#"{ACTION_START}
190 +{{"name":"get_weather","parameters":{{}}}}{ACTION_END}"#
191 + ))
192 + .unwrap();
193 +
194 + assert_eq!(result.calls.len(), 1);
195 + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
196 + }
197 +
198 + #[test]
199 + fn internlm2_does_not_validate_or_normalize_arguments() {
200 + let mut parser = Internlm2ToolParser::new(&test_tools());
201 + let arguments = r#"{"location":"Tokyo",}"#;
202 + let result = parser
203 + .parse_complete(&build_tool_call("get_weather", "parameters", arguments))
204 + .unwrap();
205 +
206 + assert_eq!(result.calls[0].arguments, arguments);
207 + }
208 +
209 + #[test]
210 + fn internlm2_streaming_emits_argument_deltas() {
211 + let mut parser = Internlm2ToolParser::new(&test_tools());
212 + let chunks = [
213 + "preface <|action",
214 + "_start|><|plugin|>",
215 + r#"{"name":"get_weather","parameters":"#,
216 + r#"{"location":"#,
217 + r#""Beijing""#,
218 + r#"}"#,
219 + r#"}<|action_end|> suffix"#,
220 + ];
221 +
222 + let mut result = ToolParserOutput::default();
223 + let mut observed_arguments = Vec::new();
224 + for chunk in chunks {
225 + let next = parser.parse_chunk(chunk).unwrap();
226 + observed_arguments.extend(
227 + next.calls
228 + .iter()
229 + .filter(|call| call.name.is_none())
230 + .map(|call| call.arguments.clone()),
231 + );
232 + result.append(next);
233 + }
234 + result.append(parser.finish().unwrap());
235 +
236 + assert_eq!(
237 + observed_arguments,
238 + [r#"{"location":"#, r#""Beijing""#, r#"}"#]
239 + );
240 + assert_eq!(result.normal_text, "preface suffix");
241 + assert_eq!(
242 + result.coalesce_calls().calls[0].arguments,
243 + r#"{"location":"Beijing"}"#
244 + );
245 + }
246 +
247 + #[test]
248 + fn internlm2_streaming_handles_split_markers() {
249 + let input = format!(
250 + "hello {}",
251 + build_tool_call("get_weather", "parameters", r#"{"location":"Tokyo"}"#)
252 + );
253 + let chunks = split_by_chars(&input, 5);
254 + let mut parser = Internlm2ToolParser::new(&test_tools());
255 +
256 + let result = collect_stream(&mut parser, &chunks);
257 +
258 + assert_eq!(result.normal_text, "hello ");
259 + assert_eq!(result.calls.len(), 1);
260 + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
261 + }
262 +
263 + #[test]
264 + fn internlm2_streaming_extracts_multiple_blocks() {
265 + let input = format!(
266 + "{}{}",
267 + build_tool_call("get_weather", "parameters", r#"{"location":"Shanghai"}"#),
268 + build_tool_call("add", "arguments", r#"{"x":1,"y":2}"#),
269 + );
270 + let chunks = split_by_chars(&input, 7);
271 + let mut parser = Internlm2ToolParser::new(&test_tools());
272 +
273 + let result = collect_stream(&mut parser, &chunks);
274 +
275 + expect![[r#"
276 + ToolParserOutput {
277 + normal_text: "",
278 + calls: [
279 + ToolCallDelta {
280 + tool_index: 0,
281 + name: Some(
282 + "get_weather",
283 + ),
284 + arguments: "{\"location\":\"Shanghai\"}",
285 + },
286 + ToolCallDelta {
287 + tool_index: 1,
288 + name: Some(
289 + "add",
290 + ),
291 + arguments: "{\"x\":1,\"y\":2}",
292 + },
293 + ],
294 + }
295 + "#]]
296 + .assert_debug_eq(&result);
297 + }
298 +
299 + #[test]
300 + fn internlm2_keeps_end_marker_literal_inside_json_string() {
301 + let mut parser = Internlm2ToolParser::new(&test_tools());
302 + let arguments = format!(r#"{{"text":"literal {ACTION_END} inside"}}"#);
303 + let input = build_tool_call("echo", "parameters", &arguments);
304 +
305 + let result = parser.parse_complete(&input).unwrap();
306 +
307 + assert_eq!(result.calls.len(), 1);
308 + assert_eq!(result.calls[0].arguments, arguments);
309 + }
310 +
311 + #[test]
312 + fn internlm2_finish_errors_on_truncated_tool_call() {
313 + let mut parser = Internlm2ToolParser::new(&test_tools());
314 + let pre_finish = parser
315 + .parse_chunk(&format!(
316 + r#"{ACTION_START}{{"name":"get_weather","parameters":{{"location""#
317 + ))
318 + .unwrap();
319 + let error = parser.finish().unwrap_err();
320 +
321 + assert_eq!(
322 + pre_finish.calls[0].name.as_deref(),
323 + Some("get_weather"),
324 + "name delta is still emitted from parse_chunk() before truncation",
325 + );
326 + assert!(
327 + error.to_report_string().contains("incomplete InternLM2 tool call"),
328 + "finish() reports the truncated tool call as incomplete: {}",
329 + error.to_report_string(),
330 + );
331 + }
332 +
333 + #[test]
334 + fn internlm2_unknown_arguments_key_fails() {
335 + let mut parser = Internlm2ToolParser::new(&test_tools());
336 + let input = build_tool_call("get_weather", "params", r#"{"location":"Tokyo"}"#);
337 +
338 + let error = parser.parse_chunk(&input).unwrap_err();
339 +
340 + expect![[r#"
341 + tool parser parsing failed: invalid InternLM2
342 + expected `parameters`, `arguments`"#]]
343 + .assert_eq(&error.to_report_string());
344 + }
345 +
346 + #[test]
347 + fn internlm2_preserve_special_tokens_is_true() {
348 + let parser = Internlm2ToolParser::new(&test_tools());
349 + assert!(parser.preserve_special_tokens());
350 + }
351 +}
rust/src/tool-parser/src/json/llama.rs
+1 -1
@@ -203,7 +203,7 @@ fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<Ll
203 marker_whitespace: JsonToolCallWhitespace::Optional,
204 delimiter: Some(";"),
205 name_key: "name",
206 - arguments_key: "parameters",
206 + arguments_key: &["parameters"],
207 };
208
209 match tool_call_header_event(input, CONFIG)? {
rust/src/tool-parser/src/json/mistral.rs
+1 -1
@@ -8,7 +8,7 @@ const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
8 marker_whitespace: JsonToolCallWhitespace::Optional,
9 delimiter: Some(","),
10 name_key: "name",
11 - arguments_key: "arguments",
11 + arguments_key: &["arguments"],
12 };
13
14 /// Tool parser for Mistral JSON-array tool calls.
rust/src/tool-parser/src/json/mod.rs
+43 -5
@@ -1,20 +1,22 @@
1 //! Shared parser core for JSON tool calls wrapped by text markers.
2
3 pub use hermes::HermesToolParser;
4 +pub use internlm2::Internlm2ToolParser;
5 pub use llama::Llama3JsonToolParser;
6 pub use mistral::MistralToolParser;
7 pub use qwen::Qwen3XmlToolParser;
8
9 mod hermes;
10 +mod internlm2;
11 mod llama;
12 mod mistral;
13 mod qwen;
14
15 use winnow::ascii::multispace0 as ws0;
16 use winnow::combinator::{alt, seq};
15 -use winnow::error::{ModalResult, StrContext, StrContextValue};
17 +use winnow::error::{AddContext, ModalResult, StrContext, StrContextValue};
18 use winnow::prelude::*;
17 -use winnow::stream::Partial;
19 +use winnow::stream::{Partial, Stream};
20 use winnow::token::literal;
21
22 use super::utils::{
@@ -32,7 +34,10 @@ struct JsonToolCallConfig {
34 marker_whitespace: JsonToolCallWhitespace,
35 delimiter: Option<&'static str>,
36 name_key: &'static str,
35 - arguments_key: &'static str,
37 + /// Candidate JSON keys naming the arguments payload, tried in order.
38 + /// Most parsers use a single key like `["arguments"]`, but some accept
39 + /// multiple (e.g. InternLM2 accepts `parameters` or `arguments`).
40 + arguments_key: &'static [&'static str],
41 }
42
43 #[derive(Debug, Clone, Copy)]
@@ -224,7 +229,7 @@ fn tool_call_header_event(
229 _: ws0,
230 _: literal(","),
231 _: ws0,
227 - _: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key),
232 + _: |input: &mut JsonToolInput<'_>| json_arguments_key(input, config.arguments_key),
233 _: ws0,
234 _: literal(":"),
235 _: ws0,
@@ -246,6 +251,39 @@ fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()>
251 .parse_next(input)
252 }
253
254 +/// Parse a JSON object key accepting any of `candidates`.
255 +///
256 +/// The full quoted key is consumed and compared against the candidate list,
257 +/// so this works correctly under partial input regardless of key lengths.
258 +///
259 +/// On mismatch, each candidate is attached as its own `Expected` context so the
260 +/// error enumerates every valid key ("expected `a`, expected `b`"). Because
261 +/// `StrContextValue::StringLiteral` carries a single `&'static str`, the
262 +/// contexts are added in a loop over `candidates` rather than through chained
263 +/// `.context(...)` calls, which keeps the diagnostics complete for any number
264 +/// of candidates.
265 +fn json_arguments_key(
266 + input: &mut JsonToolInput<'_>,
267 + candidates: &'static [&'static str],
268 +) -> ModalResult<()> {
269 + let start = input.checkpoint();
270 + json_str
271 + .verify(|key: &String| candidates.contains(&key.as_str()))
272 + .void()
273 + .parse_next(input)
274 + .map_err(|err| {
275 + err.map(|context_error| {
276 + candidates.iter().fold(context_error, |context_error, candidate| {
277 + context_error.add_context(
278 + &*input,
279 + &start,
280 + StrContext::Expected(StrContextValue::StringLiteral(candidate)),
281 + )
282 + })
283 + })
284 + })
285 +}
286 +
287 /// Parse one event inside a marker-wrapped JSON tool-call arguments payload.
288 fn parse_arguments_event(
289 input: &mut JsonToolInput<'_>,
@@ -341,7 +379,7 @@ mod tests {
379 marker_whitespace: JsonToolCallWhitespace::Optional,
380 delimiter: Some("<"),
381 name_key: "function",
344 - arguments_key: "parameters",
382 + arguments_key: &["parameters"],
383 };
384
385 fn build_tool_call(function_name: &str, arguments: &str) -> String {
rust/src/tool-parser/src/json/qwen.rs
+1 -1
@@ -8,7 +8,7 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
8 marker_whitespace: JsonToolCallWhitespace::Exact("\n"),
9 delimiter: None,
10 name_key: "name",
11 - arguments_key: "arguments",
11 + arguments_key: &["arguments"],
12 };
13
14 /// Tool parser for Qwen XML-wrapped JSON tool calls.
rust/src/tool-parser/src/lib.rs
+4 -1
@@ -24,7 +24,10 @@ pub use error::{Result, ToolParserError};
24 pub use gemma4::Gemma4ToolParser;
25 pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
26 pub use hy_v3::HyV3ToolParser;
27 -pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
27 +pub use json::{
28 + HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser,
29 + Qwen3XmlToolParser,
30 +};
31 pub use kimi_k2::KimiK2ToolParser;
32 pub use minimax_m2::MinimaxM2ToolParser;
33 pub use qwen_coder::Qwen3CoderToolParser;