@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f9e684499f

[Rust Frontend] Migrate gemma4 to unified parser (#46602)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 26, 2026 at 07:59 UTC f9e684499f67071641bb2333d52dadd879231ac2
23 files changed +1000 -394
pyproject.toml
+2 -1
@@ -129,7 +129,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
129 "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
130 "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
131 "rust/src/chat/src/renderer/deepseek_v32/fixtures/*",
132 - "rust/src/parser/src/tool/gemma4.rs", "rust/src/text/src/output/decoded.rs",
132 + "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs",
133 + "rust/src/text/src/output/decoded.rs",
134 "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"]
135 ignore-hidden = false
136
rust/Cargo.toml
+1 -1
@@ -120,7 +120,7 @@ vllm-parser = { path = "src/parser" }
120 vllm-server = { path = "src/server" }
121 vllm-text = { path = "src/text" }
122 vllm-tokenizer = { path = "src/tokenizer" }
123 -winnow = "1.0.2"
123 +winnow = { version = "1.0.2", features = ["simd"] }
124 xgrammar-structural-tag = "0.1.0"
125 zeromq = { version = "0.6.0", default-features = false, features = [
126 "tokio-runtime",
rust/src/chat/src/output/default/mod.rs
+149 -15
@@ -18,7 +18,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent
18 use crate::parser::ParserSelection;
19 use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory};
20 use crate::parser::tool::{ToolParser, ToolParserFactory};
21 -use crate::request::ChatRequest;
21 +use crate::parser::unified::UnifiedParserFactory;
22 +use crate::request::{ChatRequest, ChatTool};
23 use crate::{Error, Result as ChatResult};
24
25 /// Default request-scoped output processor used by Hugging Face style chat
@@ -46,20 +47,31 @@ impl DefaultChatOutputProcessor {
47 tool_call_parser: &ParserSelection,
48 reasoning_parser: &ParserSelection,
49 ) -> ChatResult<Self> {
49 - let tool_parsing_enabled = request.tool_parsing_enabled();
50 - let tool_parser = if tool_parsing_enabled {
51 - Some(Self::resolve_tool_parser(
52 - request,
50 + let parser = if tool_call_parser == reasoning_parser
51 + && let Some(parser) = Self::resolve_optional_unified_parser(
52 + &request.tools,
53 model_id,
54 + tokenizer.clone(),
55 tool_call_parser,
55 - )?)
56 + )? {
57 + parser
58 } else {
57 - None
59 + let tool_parsing_enabled = request.tool_parsing_enabled();
60 + let tool_parser = if tool_parsing_enabled {
61 + Some(Self::resolve_tool_parser(
62 + &request.tools,
63 + model_id,
64 + tool_call_parser,
65 + )?)
66 + } else {
67 + None
68 + };
69 + let reasoning_parser =
70 + Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?;
71 + Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
72 };
59 - let reasoning_parser =
60 - Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?;
61 - let parser: Box<dyn UnifiedParser> =
62 - Box::new(CombinedParser::new(reasoning_parser, tool_parser));
73 +
74 + apply_structural_tag_constraint(request, parser.structural_tag_model())?;
75
76 if parser.preserve_special_tokens() {
77 request.decode_options.skip_special_tokens = false;
@@ -84,7 +96,7 @@ impl DefaultChatOutputProcessor {
96 }
97
98 fn resolve_tool_parser(
87 - request: &mut ChatRequest,
99 + tools: &[ChatTool],
100 model_id: &str,
101 selection: &ParserSelection,
102 ) -> ChatResult<Box<dyn ToolParser>> {
@@ -100,14 +112,36 @@ impl DefaultChatOutputProcessor {
112 ParserSelection::Explicit(name) => name.as_str(),
113 };
114
103 - let parser = factory.create(parser_name, &request.tools)?;
104 -
105 - apply_structural_tag_constraint(request, parser.as_ref())?;
115 + let parser = factory.create(parser_name, tools)?;
116
117 TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser"));
118 Ok(parser)
119 }
120
121 + fn resolve_optional_unified_parser(
122 + tools: &[ChatTool],
123 + model_id: &str,
124 + tokenizer: DynTokenizer,
125 + selection: &ParserSelection,
126 + ) -> ChatResult<Option<Box<dyn UnifiedParser>>> {
127 + let factory = UnifiedParserFactory::global();
128 + let parser_name = match selection {
129 + ParserSelection::Auto => factory.resolve_name_for_model(model_id),
130 + ParserSelection::None => None,
131 + ParserSelection::Explicit(name) if factory.contains(name) => Some(name.as_str()),
132 + ParserSelection::Explicit(_) => None,
133 + };
134 +
135 + let Some(parser_name) = parser_name else {
136 + return Ok(None);
137 + };
138 +
139 + let parser = factory.create(parser_name, tools, tokenizer)?;
140 +
141 + UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser"));
142 + Ok(Some(parser))
143 + }
144 +
145 fn resolve_optional_reasoning_parser(
146 model_id: &str,
147 tokenizer: DynTokenizer,
@@ -134,6 +168,7 @@ impl DefaultChatOutputProcessor {
168
169 static TOOL_PARSER_LOG_ONCE: Once = Once::new();
170 static REASONING_PARSER_LOG_ONCE: Once = Once::new();
171 +static UNIFIED_PARSER_LOG_ONCE: Once = Once::new();
172
173 impl ChatOutputProcessor for DefaultChatOutputProcessor {
174 /// Transforms a raw generate-output token stream into structured chat
@@ -149,3 +184,102 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
184 Ok(structured.boxed())
185 }
186 }
187 +
188 +#[cfg(test)]
189 +mod tests {
190 + use std::sync::Arc;
191 +
192 + use vllm_tokenizer::Tokenizer;
193 +
194 + use super::DefaultChatOutputProcessor;
195 + use crate::Error;
196 + use crate::parser::ParserSelection;
197 + use crate::request::ChatRequest;
198 +
199 + struct FakeTokenizer;
200 +
201 + impl Tokenizer for FakeTokenizer {
202 + fn encode(
203 + &self,
204 + text: &str,
205 + _add_special_tokens: bool,
206 + ) -> vllm_tokenizer::Result<Vec<u32>> {
207 + Ok(text.chars().map(u32::from).collect())
208 + }
209 +
210 + fn decode(
211 + &self,
212 + token_ids: &[u32],
213 + _skip_special_tokens: bool,
214 + ) -> vllm_tokenizer::Result<String> {
215 + Ok(token_ids
216 + .iter()
217 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
218 + .collect())
219 + }
220 +
221 + fn token_to_id(&self, token: &str) -> Option<u32> {
222 + match token {
223 + "<|channel>" => Some(1),
224 + "<channel|>" => Some(2),
225 + _ => None,
226 + }
227 + }
228 + }
229 +
230 + fn tokenizer() -> Arc<FakeTokenizer> {
231 + Arc::new(FakeTokenizer)
232 + }
233 +
234 + #[test]
235 + fn equal_explicit_gemma4_uses_unified_parser() {
236 + let mut request = ChatRequest::for_test();
237 + let selection = ParserSelection::Explicit("gemma4".to_string());
238 +
239 + DefaultChatOutputProcessor::new(
240 + &mut request,
241 + "other-model",
242 + tokenizer(),
243 + &selection,
244 + &selection,
245 + )
246 + .unwrap();
247 + }
248 +
249 + #[test]
250 + fn auto_auto_gemma4_model_uses_unified_parser() {
251 + let mut request = ChatRequest::for_test();
252 +
253 + DefaultChatOutputProcessor::new(
254 + &mut request,
255 + "google/gemma-4-27b-it",
256 + tokenizer(),
257 + &ParserSelection::Auto,
258 + &ParserSelection::Auto,
259 + )
260 + .unwrap();
261 + }
262 +
263 + #[test]
264 + fn mixed_gemma4_selection_uses_split_dummy_error() {
265 + let mut request = ChatRequest::for_test();
266 + let error = match DefaultChatOutputProcessor::new(
267 + &mut request,
268 + "other-model",
269 + tokenizer(),
270 + &ParserSelection::Auto,
271 + &ParserSelection::Explicit("gemma4".to_string()),
272 + ) {
273 + Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"),
274 + Err(error) => error,
275 + };
276 +
277 + let Error::ParserInitialization { error, .. } = error else {
278 + panic!("expected parser initialization error");
279 + };
280 + assert_eq!(
281 + error.to_string(),
282 + "`gemma4` only provides a unified parser; the same reasoning parser and tool parser should be specified together"
283 + );
284 + }
285 +}
rust/src/chat/src/output/default/structural_tag.rs
+12 -12
@@ -2,12 +2,12 @@
2
3 use thiserror_ext::AsReport;
4 use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
5 +use vllm_parser::tool::StructuralTagModel;
6 use xgrammar_structural_tag::{
7 FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
8 build_structural_tag,
9 };
10
10 -use crate::parser::tool::ToolParser;
11 use crate::request::{ChatRequest, ChatToolChoice};
12 use crate::{Error, Result as ChatResult};
13
@@ -15,9 +15,9 @@ use crate::{Error, Result as ChatResult};
15 /// support and the request's tool choice.
16 pub(super) fn apply_structural_tag_constraint(
17 request: &mut ChatRequest,
18 - parser: &dyn ToolParser,
18 + model: Option<StructuralTagModel>,
19 ) -> ChatResult<()> {
20 - let Some(model) = parser.structural_tag_model() else {
20 + let Some(model) = model else {
21 return Ok(());
22 };
23 let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -77,7 +77,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagTool
77 mod tests {
78 use serde_json::{Value, json};
79 use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
80 - use vllm_parser::tool::{Qwen3CoderToolParser, Tool};
80 + use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
81
82 use super::*;
83
@@ -134,7 +134,7 @@ mod tests {
134 let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
135 let parser = qwen3_coder_parser(&request.tools);
136
137 - apply_structural_tag_constraint(&mut request, parser.as_ref())
137 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
138 .expect("structural tag should build");
139
140 let tag = structural_tag_value(&request);
@@ -147,7 +147,7 @@ mod tests {
147 let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
148 let parser = qwen3_coder_parser(&request.tools);
149
150 - apply_structural_tag_constraint(&mut request, parser.as_ref())
150 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
151 .expect("structural tag decision should succeed");
152
153 assert!(request.sampling_params.structured_outputs.is_none());
@@ -163,7 +163,7 @@ mod tests {
163 });
164 let parser = qwen3_coder_parser(&request.tools);
165
166 - apply_structural_tag_constraint(&mut request, parser.as_ref())
166 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
167 .expect("structural tag should build");
168
169 let params = structured_outputs(&request);
@@ -179,7 +179,7 @@ mod tests {
179 let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
180 let parser = qwen3_coder_parser(&request.tools);
181
182 - apply_structural_tag_constraint(&mut request, parser.as_ref())
182 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
183 .expect("structural tag should build");
184
185 let tag = structural_tag_value(&request);
@@ -197,7 +197,7 @@ mod tests {
197 });
198 let parser = qwen3_coder_parser(&request.tools);
199
200 - apply_structural_tag_constraint(&mut request, parser.as_ref())
200 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
201 .expect("structural tag should build");
202
203 let params = structured_outputs(&request);
@@ -218,7 +218,7 @@ mod tests {
218 );
219 let parser = qwen3_coder_parser(&request.tools);
220
221 - apply_structural_tag_constraint(&mut request, parser.as_ref())
221 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
222 .expect("structural tag should build");
223
224 let tag = structural_tag_value(&request).to_string();
@@ -231,7 +231,7 @@ mod tests {
231 let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
232 let parser = qwen3_coder_parser(&request.tools);
233
234 - apply_structural_tag_constraint(&mut request, parser.as_ref())
234 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
235 .expect("structural tag decision should succeed");
236
237 assert!(request.sampling_params.structured_outputs.is_none());
@@ -247,7 +247,7 @@ mod tests {
247 });
248 let parser = qwen3_coder_parser(&request.tools);
249
250 - apply_structural_tag_constraint(&mut request, parser.as_ref())
250 + apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
251 .expect("structural tag decision should succeed");
252
253 let params = structured_outputs(&request);
rust/src/chat/src/output/default/unified.rs
+1 -1
@@ -317,7 +317,7 @@ mod tests {
317 where
318 Self: Sized + 'static,
319 {
320 - unreachable!("ScriptedParser is constructed directly in tests")
320 + Ok(Box::new(Self::new([])))
321 }
322
323 fn parse_into(
rust/src/chat/src/parser/mod.rs
+1
@@ -1,5 +1,6 @@
1 pub mod reasoning;
2 pub mod tool;
3 +pub mod unified;
4
5 use std::collections::HashMap;
6 use std::convert::Infallible;
rust/src/chat/src/parser/reasoning/mod.rs
+21 -10
@@ -1,13 +1,13 @@
1 //! Reasoning parser registration and selection boundary for `vllm-chat`.
2
3 -use std::sync::LazyLock;
3 +use std::sync::{Arc, LazyLock};
4
5 pub use vllm_parser::reasoning::{
6 CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
7 - DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
8 - KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
9 - NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
10 - ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser, Step3p5ReasoningParser,
7 + DeepSeekV4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser,
8 + MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, NemotronV3ReasoningParser,
9 + Qwen3ReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser,
10 + Step3ReasoningParser, Step3p5ReasoningParser,
11 };
12 use vllm_tokenizer::DynTokenizer;
13
@@ -33,8 +33,9 @@ pub mod names {
33 }
34
35 /// Constructor signature for one registered reasoning parser implementation.
36 -type ReasoningParserCreator =
37 - fn(DynTokenizer) -> vllm_parser::reasoning::Result<Box<dyn ReasoningParser>>;
36 +type ReasoningParserCreator = Arc<
37 + dyn Fn(DynTokenizer) -> vllm_parser::reasoning::Result<Box<dyn ReasoningParser>> + Send + Sync,
38 +>;
39
40 /// Registry and model matcher for reasoning parsers.
41 pub type ReasoningParserFactory = ParserFactory<ReasoningParserCreator>;
@@ -58,7 +59,7 @@ impl ReasoningParserFactory {
59 .register_parser::<DeepSeekR1ReasoningParser>(names::DEEPSEEK_R1)
60 .register_parser::<DeepSeekV3ReasoningParser>(names::DEEPSEEK_V3)
61 .register_parser::<DeepSeekV4ReasoningParser>(names::DEEPSEEK_V4)
61 - .register_parser::<Gemma4ReasoningParser>(names::GEMMA4)
62 + .register_unified_dummy(names::GEMMA4)
63 .register_parser::<Glm45ReasoningParser>(names::GLM45)
64 .register_parser::<KimiReasoningParser>(names::KIMI)
65 .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
@@ -109,7 +110,17 @@ impl ReasoningParserFactory {
110 where
111 T: ReasoningParser + 'static,
112 {
112 - self.register_creator(name, T::create)
113 + self.register_creator(name, Arc::new(T::create))
114 + }
115 +
116 + /// Register one unified-only parser name in the split reasoning registry.
117 + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self {
118 + let name = name.to_string();
119 + let registered_name = name.clone();
120 + self.register_creator(
121 + &registered_name,
122 + Arc::new(move |_| Err(ReasoningError::DummyUnifiedParser { name: name.clone() })),
123 + )
124 }
125
126 /// Construct a parser from an exact name.
@@ -124,7 +135,7 @@ impl ReasoningParserFactory {
135 available_names: self.list(),
136 })?;
137
127 - creator(tokenizer).map_err(|error| crate::Error::ParserInitialization {
138 + creator.as_ref()(tokenizer).map_err(|error| crate::Error::ParserInitialization {
139 kind: "reasoning",
140 name: name.to_string(),
141 error: error.into(),
rust/src/chat/src/parser/reasoning/tests.rs
+2
@@ -35,11 +35,13 @@ fn factory_contains_and_lists_registered_parsers() {
35 assert!(factory.contains(names::SEED_OSS));
36 assert!(factory.contains(names::STEP3P5));
37 assert!(factory.contains(names::MINIMAX_M3));
38 + assert!(factory.contains(names::GEMMA4));
39 assert!(factory.list().contains(&names::QWEN3.to_string()));
40 assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
41 assert!(factory.list().contains(&names::SEED_OSS.to_string()));
42 assert!(factory.list().contains(&names::STEP3P5.to_string()));
43 assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
44 + assert!(factory.list().contains(&names::GEMMA4.to_string()));
45 }
46
47 #[test]
rust/src/chat/src/parser/tool/mod.rs
+20 -9
@@ -1,13 +1,13 @@
1 //! Tool parser registration and selection boundary for `vllm-chat`.
2
3 -use std::sync::LazyLock;
3 +use std::sync::{Arc, LazyLock};
4
5 pub use vllm_parser::tool::{
6 DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
7 - Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser,
8 - HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
9 - MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
10 - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolParser, ToolParserError,
7 + Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser,
8 + Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
9 + MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser,
10 + Qwen3XmlToolParser, ToolParser, ToolParserError,
11 };
12
13 use crate::parser::ParserFactory;
@@ -40,7 +40,8 @@ pub mod names {
40 }
41
42 /// Constructor signature for one registered tool parser implementation.
43 -type ToolParserCreator = fn(&[ChatTool]) -> vllm_parser::tool::Result<Box<dyn ToolParser>>;
43 +type ToolParserCreator =
44 + Arc<dyn Fn(&[ChatTool]) -> vllm_parser::tool::Result<Box<dyn ToolParser>> + Send + Sync>;
45
46 /// Registry and model matcher for tool parsers.
47 pub type ToolParserFactory = ParserFactory<ToolParserCreator>;
@@ -65,7 +66,7 @@ impl ToolParserFactory {
66 .register_parser::<DeepSeekV4ToolParser>(names::DEEPSEEK_V4)
67 .register_parser::<Glm45MoeToolParser>(names::GLM45)
68 .register_parser::<Glm47MoeToolParser>(names::GLM47)
68 - .register_parser::<Gemma4ToolParser>(names::GEMMA4)
69 + .register_unified_dummy(names::GEMMA4)
70 .register_parser::<Granite4ToolParser>(names::GRANITE4)
71 .register_parser::<HermesToolParser>(names::HERMES)
72 .register_parser::<HyV3ToolParser>(names::HY_V3)
@@ -126,7 +127,17 @@ impl ToolParserFactory {
127 where
128 T: ToolParser + 'static,
129 {
129 - self.register_creator(name, T::create)
130 + self.register_creator(name, Arc::new(T::create))
131 + }
132 +
133 + /// Register one unified-only parser name in the split tool registry.
134 + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self {
135 + let name = name.to_string();
136 + let registered_name = name.clone();
137 + self.register_creator(
138 + &registered_name,
139 + Arc::new(move |_| Err(ToolParserError::DummyUnifiedParser { name: name.clone() })),
140 + )
141 }
142
143 /// Construct a parser from an exact name.
@@ -137,7 +148,7 @@ impl ToolParserFactory {
148 available_names: self.list(),
149 })?;
150
140 - creator(tools).map_err(|error| crate::Error::ParserInitialization {
151 + creator.as_ref()(tools).map_err(|error| crate::Error::ParserInitialization {
152 kind: "tool",
153 name: name.to_string(),
154 error: error.into(),
rust/src/chat/src/parser/unified.rs new
+124
@@ -0,0 +1,124 @@
1 +//! Unified parser registration and selection boundary for `vllm-chat`.
2 +
3 +use std::sync::LazyLock;
4 +
5 +pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser};
6 +use vllm_tokenizer::DynTokenizer;
7 +
8 +use crate::parser::ParserFactory;
9 +use crate::request::ChatTool;
10 +
11 +/// Canonical public names for registered unified parsers.
12 +pub mod names {
13 + pub const GEMMA4: &str = "gemma4";
14 +}
15 +
16 +/// Constructor signature for one registered unified parser implementation.
17 +type UnifiedParserCreator =
18 + fn(&[ChatTool], DynTokenizer) -> vllm_parser::unified::Result<Box<dyn UnifiedParser>>;
19 +
20 +/// Registry and model matcher for unified parsers.
21 +pub type UnifiedParserFactory = ParserFactory<UnifiedParserCreator>;
22 +
23 +impl UnifiedParserFactory {
24 + /// Get the global unified parser factory with built-in registrations and
25 + /// model mappings.
26 + pub fn global() -> &'static Self {
27 + static INSTANCE: LazyLock<UnifiedParserFactory> = LazyLock::new(UnifiedParserFactory::new);
28 + &INSTANCE
29 + }
30 +
31 + /// Create the default registry with built-in parser names and model
32 + /// mappings.
33 + pub fn new() -> Self {
34 + let mut factory = Self::default();
35 +
36 + factory.register_parser::<Gemma4UnifiedParser>(names::GEMMA4);
37 +
38 + factory
39 + .register_pattern("gemma-4", names::GEMMA4)
40 + .register_pattern("gemma4", names::GEMMA4);
41 +
42 + factory
43 + }
44 +
45 + /// Register one parser type that exposes a static `create()` constructor.
46 + pub fn register_parser<T>(&mut self, name: &str) -> &mut Self
47 + where
48 + T: UnifiedParser + 'static,
49 + {
50 + self.register_creator(name, T::create)
51 + }
52 +
53 + /// Construct a parser from an exact name.
54 + pub fn create(
55 + &self,
56 + name: &str,
57 + tools: &[ChatTool],
58 + tokenizer: DynTokenizer,
59 + ) -> crate::Result<Box<dyn UnifiedParser>> {
60 + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName {
61 + kind: "unified",
62 + name: name.to_string(),
63 + available_names: self.list(),
64 + })?;
65 +
66 + creator(tools, tokenizer).map_err(|error| crate::Error::ParserInitialization {
67 + kind: "unified",
68 + name: name.to_string(),
69 + error: error.into(),
70 + })
71 + }
72 +}
73 +
74 +#[cfg(test)]
75 +mod tests {
76 + use std::sync::Arc;
77 +
78 + use vllm_tokenizer::Tokenizer;
79 +
80 + use super::{UnifiedParserFactory, names};
81 +
82 + struct FakeTokenizer;
83 +
84 + impl Tokenizer for FakeTokenizer {
85 + fn encode(
86 + &self,
87 + text: &str,
88 + _add_special_tokens: bool,
89 + ) -> vllm_tokenizer::Result<Vec<u32>> {
90 + Ok(text.chars().map(u32::from).collect())
91 + }
92 +
93 + fn decode(
94 + &self,
95 + token_ids: &[u32],
96 + _skip_special_tokens: bool,
97 + ) -> vllm_tokenizer::Result<String> {
98 + Ok(token_ids
99 + .iter()
100 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
101 + .collect())
102 + }
103 +
104 + fn token_to_id(&self, token: &str) -> Option<u32> {
105 + match token {
106 + "<|channel>" => Some(1),
107 + "<channel|>" => Some(2),
108 + _ => None,
109 + }
110 + }
111 + }
112 +
113 + #[test]
114 + fn factory_registers_gemma4() {
115 + let factory = UnifiedParserFactory::new();
116 +
117 + assert!(factory.contains(names::GEMMA4));
118 + assert_eq!(
119 + factory.resolve_name_for_model("google/gemma-4-27b-it"),
120 + Some(names::GEMMA4)
121 + );
122 + factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap();
123 + }
124 +}
rust/src/chat/tests/roundtrip.rs
+50 -3
@@ -37,6 +37,8 @@ struct RoundtripCase {
37 /// JSON formatting expected after this model's template has materialized
38 /// tool-call arguments.
39 json_fmt: JsonFmt,
40 + /// Whether the template renders tool-call argument object keys in sorted order.
41 + sort_json_keys: bool,
42 }
43
44 #[derive(Clone, Copy)]
@@ -81,6 +83,7 @@ impl RoundtripCase {
83 reasoning_parser: ParserSelection::Auto,
84 thinking_behavior: ThinkingBehavior::Toggleable { default: true },
85 json_fmt: spaced_json_fmt(),
86 + sort_json_keys: false,
87 }
88 }
89
@@ -93,6 +96,7 @@ impl RoundtripCase {
96 reasoning_parser: ParserSelection::Auto,
97 thinking_behavior: ThinkingBehavior::Toggleable { default: true },
98 json_fmt: compact_json_fmt(),
99 + sort_json_keys: false,
100 }
101 }
102
@@ -105,6 +109,7 @@ impl RoundtripCase {
109 reasoning_parser: ParserSelection::Auto,
110 thinking_behavior: ThinkingBehavior::Always { value: true },
111 json_fmt: compact_json_fmt(),
112 + sort_json_keys: false,
113 }
114 }
115
@@ -117,6 +122,7 @@ impl RoundtripCase {
122 reasoning_parser: ParserSelection::Auto,
123 thinking_behavior: ThinkingBehavior::Toggleable { default: false },
124 json_fmt: compact_json_fmt(),
125 + sort_json_keys: false,
126 }
127 }
128
@@ -129,6 +135,20 @@ impl RoundtripCase {
135 reasoning_parser: ParserSelection::Auto,
136 thinking_behavior: ThinkingBehavior::Toggleable { default: true },
137 json_fmt: compact_json_fmt(),
138 + sort_json_keys: false,
139 + }
140 + }
141 +
142 + /// Gemma4 channel reasoning with custom function-call arguments.
143 + fn gemma4() -> Self {
144 + Self {
145 + model_id: "google/gemma-4-E4B-it",
146 + assistant_stop_suffix: "<|tool_response>",
147 + tool_call_parser: ParserSelection::Auto,
148 + reasoning_parser: ParserSelection::Auto,
149 + thinking_behavior: ThinkingBehavior::Always { value: true },
150 + json_fmt: compact_json_fmt(),
151 + sort_json_keys: true,
152 }
153 }
154
@@ -142,6 +162,7 @@ impl RoundtripCase {
162 reasoning_parser: ParserSelection::Auto,
163 thinking_behavior: ThinkingBehavior::Toggleable { default: true },
164 json_fmt: spaced_json_fmt(),
165 + sort_json_keys: false,
166 }
167 }
168
@@ -154,6 +175,7 @@ impl RoundtripCase {
175 reasoning_parser: ParserSelection::Auto,
176 thinking_behavior: ThinkingBehavior::Always { value: true },
177 json_fmt: compact_json_fmt(),
178 + sort_json_keys: false,
179 }
180 }
181
@@ -166,6 +188,7 @@ impl RoundtripCase {
188 reasoning_parser: ParserSelection::Auto,
189 thinking_behavior: ThinkingBehavior::Always { value: true },
190 json_fmt: compact_json_fmt(),
191 + sort_json_keys: false,
192 }
193 }
194 }
@@ -195,8 +218,8 @@ roundtrip_tests! {
218 seed_oss => [reasoning_and_content],
219 step3p5 => [reasoning_and_content],
220
198 - // Note: Kimi K2.5 strips the reasoning content in history.
199 - kimi_k25 => [tool_call_mix],
221 + gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
222 + kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
223 }
224
225 /// Run the fixed reasoning+content fixture for one model/parser case.
@@ -347,14 +370,38 @@ fn spaced_json_fmt() -> JsonFmt {
370 /// Pass in a raw JSON string instead of a structured value to ensure the exact precision and
371 /// formatting of numbers are preserved.
372 fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result<String> {
350 - let value: serde_json::Value =
373 + let mut value: serde_json::Value =
374 serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?;
375 + if case.sort_json_keys {
376 + sort_json_value(&mut value);
377 + }
378
379 case.json_fmt
380 .format_to_string(&value)
381 .context("failed to format expected tool-call arguments")
382 }
383
384 +/// Sort JSON object keys recursively to match templates that render mappings with `dictsort`.
385 +fn sort_json_value(value: &mut serde_json::Value) {
386 + match value {
387 + serde_json::Value::Object(map) => {
388 + for value in map.values_mut() {
389 + sort_json_value(value);
390 + }
391 +
392 + let mut entries = std::mem::take(map).into_iter().collect::<Vec<_>>();
393 + entries.sort_by(|(left, _), (right, _)| left.cmp(right));
394 + map.extend(entries);
395 + }
396 + serde_json::Value::Array(values) => {
397 + for value in values {
398 + sort_json_value(value);
399 + }
400 + }
401 + _ => {}
402 + }
403 +}
404 +
405 /// Load the real model chat/text backend for one roundtrip case.
406 async fn load_roundtrip_backends(case: &RoundtripCase) -> Result<vllm_chat::LoadedModelBackends> {
407 load_model_backends(
rust/src/parser/benches/gemma4.rs
+5 -3
@@ -2,10 +2,11 @@ use std::time::Duration;
2
3 use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4 use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
5 -use vllm_parser::tool::{Gemma4ToolParser, Tool, ToolParser};
5 +use vllm_parser::tool::{Tool, ToolParser};
6 +use vllm_parser::unified::Gemma4UnifiedParser;
7
8 mod utils;
8 -use utils::feed_parser;
9 +use utils::{UnifiedToolParserAdapter, feed_parser};
10
11 const CHUNK_CHARS: usize = 7;
12 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
@@ -68,7 +69,8 @@ fn long_tool_argument_fixture() -> String {
69 }
70
71 fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
71 - Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
72 + UnifiedToolParserAdapter::<Gemma4UnifiedParser>::create(tools)
73 + .expect("Gemma4 unified parser should initialize")
74 }
75
76 fn run_stream_group(
rust/src/parser/benches/utils/adapter.rs new
+106
@@ -0,0 +1,106 @@
1 +use std::sync::Arc;
2 +
3 +use vllm_parser::tool::{
4 + Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
5 +};
6 +use vllm_parser::unified::{
7 + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
8 +};
9 +use vllm_tokenizer::Tokenizer;
10 +
11 +/// Tokenizer stub used by unified-parser benchmarks.
12 +struct BenchTokenizer;
13 +
14 +impl Tokenizer for BenchTokenizer {
15 + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
16 + Ok(text.chars().map(|_| u32::MAX).collect())
17 + }
18 +
19 + fn decode(
20 + &self,
21 + token_ids: &[u32],
22 + _skip_special_tokens: bool,
23 + ) -> vllm_tokenizer::Result<String> {
24 + Ok("\u{FFFD}".repeat(token_ids.len()))
25 + }
26 +
27 + fn token_to_id(&self, _token: &str) -> Option<u32> {
28 + Some(u32::MAX)
29 + }
30 +}
31 +
32 +/// Bench-only adapter that exposes a unified parser through the tool-parser
33 +/// benchmark harness.
34 +///
35 +/// Returns error if the unified parser produces reasoning events.
36 +pub struct UnifiedToolParserAdapter<T> {
37 + inner: Box<dyn UnifiedParser>,
38 + _marker: std::marker::PhantomData<T>,
39 +}
40 +
41 +fn map_unified_error(error: UnifiedParserError) -> ToolParserError {
42 + ToolParserError::ParsingFailed {
43 + message: format!("unified parser failed: {error}"),
44 + }
45 +}
46 +
47 +fn append_unified_output(
48 + output: UnifiedParserOutput,
49 + tool_output: &mut ToolParserOutput,
50 +) -> Result<()> {
51 + for event in output.events {
52 + match event {
53 + UnifiedParserEvent::Text(text) => tool_output.push_text(text),
54 + UnifiedParserEvent::ToolCall(call) => tool_output.push_call(call),
55 + UnifiedParserEvent::Reasoning(_) => {
56 + return Err(ToolParserError::ParsingFailed {
57 + message: "unified parser emitted reasoning in tool-parser adapter".to_string(),
58 + });
59 + }
60 + }
61 + }
62 + Ok(())
63 +}
64 +
65 +impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
66 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
67 + where
68 + Self: Sized + 'static,
69 + {
70 + let inner = T::create(tools, Arc::new(BenchTokenizer)).map_err(map_unified_error)?;
71 + Ok(Box::new(Self {
72 + inner,
73 + _marker: std::marker::PhantomData,
74 + }))
75 + }
76 +
77 + fn preserve_special_tokens(&self) -> bool {
78 + self.inner.preserve_special_tokens()
79 + }
80 +
81 + fn structural_tag_model(&self) -> Option<StructuralTagModel> {
82 + self.inner.structural_tag_model()
83 + }
84 +
85 + fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
86 + self.inner.tool_call_id(tool_index)
87 + }
88 +
89 + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
90 + let mut unified_output = UnifiedParserOutput::default();
91 + let result = self.inner.parse_into(chunk, &mut unified_output).map_err(map_unified_error);
92 + append_unified_output(unified_output, output)?;
93 + result
94 + }
95 +
96 + fn finish(&mut self) -> Result<ToolParserOutput> {
97 + let unified_output = self.inner.finish().map_err(map_unified_error)?;
98 + let mut output = ToolParserOutput::default();
99 + append_unified_output(unified_output, &mut output)?;
100 + Ok(output)
101 + }
102 +
103 + fn reset(&mut self) -> String {
104 + self.inner.reset()
105 + }
106 +}
rust/src/parser/benches/utils/mod.rs
+7
@@ -1,4 +1,9 @@
1 +// This module is shared by multiple benchmark targets.
2 +// There could be false positives for unused code or imports, and fixing them would lead to some other benchmarks failing to compile.
3 #![allow(dead_code)]
4 +#![allow(unused_imports)]
5 +
6 +mod adapter;
7
8 use futures::FutureExt as _;
9 use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
@@ -6,6 +11,8 @@ use tool_parser::traits::ToolParser as ExternalToolParser;
11 use vllm_parser::tool::test_utils::collect_stream;
12 use vllm_parser::tool::{Tool, ToolParser};
13
14 +pub(super) use adapter::UnifiedToolParserAdapter;
15 +
16 pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
17 tools
18 .iter()
rust/src/parser/src/lib.rs
+1
@@ -3,3 +3,4 @@
3 pub mod reasoning;
4 pub mod tool;
5 pub mod unified;
6 +pub(crate) mod utils;
rust/src/parser/src/reasoning/delimited.rs
+5 -5
@@ -144,20 +144,20 @@ impl DelimitedReasoningParser {
144 }
145
146 /// Determine the reasoning state implied by the last prompt boundary, if any.
147 -fn last_reasoning_boundary(
147 +pub(crate) fn last_reasoning_boundary(
148 prompt_token_ids: &[u32],
149 start_token_id: u32,
150 end_token_id: u32,
151 tokenizer: &dyn Tokenizer,
152 ) -> Option<bool> {
153 - for token_id in prompt_token_ids.iter().rev() {
154 - if *token_id == start_token_id {
153 + for token_id in prompt_token_ids.iter().rev().copied() {
154 + if token_id == start_token_id {
155 return Some(true);
156 }
157 - if *token_id == end_token_id {
157 + if token_id == end_token_id {
158 return Some(false);
159 }
160 - if tokenizer.is_special_id(*token_id) {
160 + if tokenizer.is_special_id(token_id) {
161 return None;
162 }
163 }
rust/src/parser/src/reasoning/gemma4.rs deleted
-273
@@ -1,273 +0,0 @@
1 -use vllm_tokenizer::DynTokenizer;
2 -
3 -use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 -
5 -const THOUGHT_PREFIX: &str = "thought\n";
6 -
7 -/// Reasoning parser for Google Gemma4 thinking models.
8 -///
9 -/// Gemma4 emits reasoning inside `<|channel> ... <channel|>` spans and adds a
10 -/// structural `thought\n` label at the beginning of the reasoning channel.
11 -/// This parser keeps the delimiter handling in the shared delimited parser and
12 -/// only layers on Gemma4-specific request adjustment plus prefix stripping.
13 -///
14 -/// Original Python implementation:
15 -/// <https://github.com/vllm-project/vllm/blob/18b1c77211d8f6fe800bcfb89524d2b598708032/vllm/reasoning/gemma4_reasoning_parser.py#L23>
16 -pub struct Gemma4ReasoningParser {
17 - inner: DelimitedReasoningParser,
18 - reasoning_text: String,
19 - prefix_stripped: bool,
20 -}
21 -
22 -impl Gemma4ReasoningParser {
23 - /// Create a Gemma4 parser.
24 - pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
25 - Ok(Self {
26 - inner: DelimitedReasoningParser::new(tokenizer, "<|channel>", "<channel|>", false)?,
27 - reasoning_text: String::new(),
28 - prefix_stripped: false,
29 - })
30 - }
31 -
32 - /// Apply Gemma4's `thought\n` stripping rule to one reasoning delta.
33 - ///
34 - /// Early reasoning text is buffered until we can decide whether it begins
35 - /// with the structural channel label.
36 - fn strip_thought_prefix(&mut self, reasoning: &str) -> Option<String> {
37 - if self.prefix_stripped {
38 - return Some(reasoning.to_string());
39 - }
40 -
41 - self.reasoning_text.push_str(reasoning);
42 -
43 - if self.reasoning_text.starts_with(THOUGHT_PREFIX) {
44 - let prefix_len = THOUGHT_PREFIX.len();
45 - let previous_len = self.reasoning_text.len() - reasoning.len();
46 - if previous_len >= prefix_len {
47 - self.reasoning_text.clear();
48 - self.prefix_stripped = true;
49 - return Some(reasoning.to_string());
50 - }
51 -
52 - let prefix_chars_in_delta = prefix_len - previous_len;
53 - let stripped = &reasoning[prefix_chars_in_delta.min(reasoning.len())..];
54 - if stripped.is_empty() {
55 - if self.reasoning_text.len() >= prefix_len {
56 - self.reasoning_text.clear();
57 - self.prefix_stripped = true;
58 - }
59 - return None;
60 - }
61 -
62 - self.reasoning_text.clear();
63 - self.prefix_stripped = true;
64 - return Some(stripped.to_string());
65 - }
66 -
67 - if THOUGHT_PREFIX.starts_with(&self.reasoning_text) {
68 - return None;
69 - }
70 -
71 - self.prefix_stripped = true;
72 - Some(std::mem::take(&mut self.reasoning_text))
73 - }
74 -
75 - /// Apply Gemma4-specific reasoning post-processing to one parsed delta.
76 - fn post_process(&mut self, mut result: ReasoningDelta) -> ReasoningDelta {
77 - if let Some(reasoning) = result.reasoning.take() {
78 - result.reasoning =
79 - self.strip_thought_prefix(&reasoning).filter(|text| !text.is_empty());
80 - }
81 - result
82 - }
83 -}
84 -
85 -impl ReasoningParser for Gemma4ReasoningParser {
86 - fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
87 - where
88 - Self: Sized + 'static,
89 - {
90 - Ok(Box::new(Self::new(tokenizer)?))
91 - }
92 -
93 - fn preserve_special_tokens(&self) -> bool {
94 - true
95 - }
96 -
97 - fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
98 - self.inner.initialize(prompt_token_ids);
99 - self.reasoning_text.clear();
100 - self.prefix_stripped = false;
101 - Ok(())
102 - }
103 -
104 - fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
105 - let result = self.inner.push(delta);
106 - Ok(self.post_process(result))
107 - }
108 -
109 - fn finish(&mut self) -> Result<ReasoningDelta> {
110 - let result = self.inner.finish();
111 - Ok(self.post_process(result))
112 - }
113 -}
114 -
115 -#[cfg(test)]
116 -mod tests {
117 - use std::sync::Arc;
118 -
119 - use vllm_tokenizer::Tokenizer;
120 -
121 - use super::Gemma4ReasoningParser;
122 - use crate::reasoning::ReasoningParser;
123 -
124 - struct FakeTokenizer;
125 -
126 - impl Tokenizer for FakeTokenizer {
127 - fn encode(
128 - &self,
129 - text: &str,
130 - _add_special_tokens: bool,
131 - ) -> vllm_tokenizer::Result<Vec<u32>> {
132 - Ok(text.chars().map(u32::from).collect())
133 - }
134 -
135 - fn decode(
136 - &self,
137 - token_ids: &[u32],
138 - _skip_special_tokens: bool,
139 - ) -> vllm_tokenizer::Result<String> {
140 - Ok(token_ids
141 - .iter()
142 - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
143 - .collect())
144 - }
145 -
146 - fn token_to_id(&self, token: &str) -> Option<u32> {
147 - match token {
148 - "<|channel>" => Some(1000),
149 - "<channel|>" => Some(1001),
150 - _ => None,
151 - }
152 - }
153 - }
154 -
155 - fn run_streaming(output: &[&str]) -> (Option<String>, Option<String>) {
156 - let tokenizer = Arc::new(FakeTokenizer);
157 - let mut parser = Gemma4ReasoningParser::new(tokenizer).unwrap();
158 - let mut reasoning = String::new();
159 - let mut content = String::new();
160 -
161 - for delta in output {
162 - let result = parser.push(delta).unwrap();
163 - if let Some(next) = result.reasoning {
164 - reasoning.push_str(&next);
165 - }
166 - if let Some(next) = result.content {
167 - content.push_str(&next);
168 - }
169 - }
170 -
171 - let final_delta = parser.finish().unwrap();
172 - if let Some(next) = final_delta.reasoning {
173 - reasoning.push_str(&next);
174 - }
175 - if let Some(next) = final_delta.content {
176 - content.push_str(&next);
177 - }
178 -
179 - (
180 - (!reasoning.is_empty()).then_some(reasoning),
181 - (!content.is_empty()).then_some(content),
182 - )
183 - }
184 -
185 - #[test]
186 - fn gemma4_reasoning_streaming_handles_channel_delimited_outputs() {
187 - let cases = [
188 - (
189 - "no_reasoning",
190 - vec!["This is content"],
191 - None,
192 - Some("This is content"),
193 - ),
194 - (
195 - "reasoning_and_content",
196 - vec!["<|channel>This is a reasoning section<channel|>This is the rest"],
197 - Some("This is a reasoning section"),
198 - Some("This is the rest"),
199 - ),
200 - (
201 - "complete_reasoning",
202 - vec!["<|channel>This is a reasoning section<channel|>"],
203 - Some("This is a reasoning section"),
204 - None,
205 - ),
206 - (
207 - "multiple_lines",
208 - vec!["<|channel>This\nThat<channel|>This is the rest\nThat"],
209 - Some("This\nThat"),
210 - Some("This is the rest\nThat"),
211 - ),
212 - (
213 - "no_end",
214 - vec!["<|channel>This is a reasoning section"],
215 - Some("This is a reasoning section"),
216 - None,
217 - ),
218 - ("empty", vec![""], None, None),
219 - (
220 - "newline_around_reasoning",
221 - vec!["Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"],
222 - Some("This is a reasoning section"),
223 - Some("Before\n\nThis is the rest"),
224 - ),
225 - (
226 - "thought_prefix",
227 - vec!["<|channel>thought\nActual reasoning here<channel|>Final answer"],
228 - Some("Actual reasoning here"),
229 - Some("Final answer"),
230 - ),
231 - (
232 - "thought_prefix_only",
233 - vec!["<|channel>thought\n<channel|>"],
234 - None,
235 - None,
236 - ),
237 - (
238 - "thought_prefix_multiline",
239 - vec!["<|channel>thought\nLine1\nLine2<channel|>Answer"],
240 - Some("Line1\nLine2"),
241 - Some("Answer"),
242 - ),
243 - (
244 - "thought_prefix_diverge",
245 - vec!["<|channel>thousand reasons<channel|>Done"],
246 - Some("thousand reasons"),
247 - Some("Done"),
248 - ),
249 - ];
250 -
251 - for (name, output, expected_reasoning, expected_content) in cases {
252 - let (reasoning, content) = run_streaming(&output);
253 - assert_eq!(reasoning.as_deref(), expected_reasoning, "{name}");
254 - assert_eq!(content.as_deref(), expected_content, "{name}");
255 - }
256 - }
257 -
258 - #[test]
259 - fn gemma4_strips_thought_prefix_even_when_split_across_deltas() {
260 - let (reasoning, content) =
261 - run_streaming(&["<|channel>thou", "ght", "\nabc", "<channel|>done"]);
262 - assert_eq!(reasoning.as_deref(), Some("abc"));
263 - assert_eq!(content.as_deref(), Some("done"));
264 - }
265 -
266 - #[test]
267 - fn gemma4_preserves_special_tokens() {
268 - let tokenizer = Arc::new(FakeTokenizer);
269 - let parser = Gemma4ReasoningParser::new(tokenizer).unwrap();
270 -
271 - assert!(parser.preserve_special_tokens());
272 - }
273 -}
rust/src/parser/src/reasoning/mod.rs
+5 -3
@@ -17,7 +17,6 @@
17 mod cohere_cmd;
18 mod deepseek_r1;
19 mod delimited;
20 -mod gemma4;
20 mod kimi;
21 mod minimax_m3;
22 mod qwen3;
@@ -29,8 +28,7 @@ use vllm_tokenizer::DynTokenizer;
28
29 pub use self::cohere_cmd::CohereCmdReasoningParser;
30 pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
32 -pub(crate) use self::delimited::DelimitedReasoningParser;
33 -pub use self::gemma4::Gemma4ReasoningParser;
31 +pub(crate) use self::delimited::{DelimitedReasoningParser, last_reasoning_boundary};
32 pub use self::kimi::KimiReasoningParser;
33 pub use self::minimax_m3::MiniMaxM3ReasoningParser;
34 pub use self::qwen3::Qwen3ReasoningParser;
@@ -129,6 +127,10 @@ pub trait ReasoningParser: Send {
127 pub enum ReasoningError {
128 #[error("tokenizer is missing reasoning delimiter token `{token}`")]
129 MissingToken { token: String },
130 + #[error(
131 + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together"
132 + )]
133 + DummyUnifiedParser { name: String },
134 }
135
136 #[cfg(test)]
rust/src/parser/src/tool/error.rs
+4
@@ -10,4 +10,8 @@ pub type Result<T> = std::result::Result<T, ToolParserError>;
10 pub enum ToolParserError {
11 #[error("tool parser parsing failed: {message}")]
12 ParsingFailed { message: String },
13 + #[error(
14 + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together"
15 + )]
16 + DummyUnifiedParser { name: String },
17 }
rust/src/parser/src/tool/mod.rs
+1 -3
@@ -4,7 +4,6 @@
4 pub(crate) mod error;
5 mod deepseek_dsml;
6 pub(crate) mod deepseek_json;
7 -mod gemma4;
7 mod glm_xml;
8 mod hy_v3;
9 mod json;
@@ -15,14 +14,13 @@ mod parameters;
14 mod qwen_coder;
15 #[cfg(any(test, feature = "test-util"))]
16 pub mod test_utils;
18 -pub(crate) mod utils;
17 +use crate::utils;
18
19 use std::collections::{BTreeMap, btree_map};
20
21 pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser};
22 pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser};
23 pub use error::{Result, ToolParserError};
25 -pub use gemma4::Gemma4ToolParser;
24 pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
25 pub use hy_v3::HyV3ToolParser;
26 pub use json::{
rust/src/parser/src/unified/gemma4.rs renamed
+371 -38
@@ -6,10 +6,17 @@ use winnow::prelude::*;
6 use winnow::stream::{Partial, Stream};
7 use winnow::token::{literal, take_till, take_until};
8
9 -use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len};
10 -use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
11 -use crate::tool::Tool;
9 +use vllm_tokenizer::DynTokenizer;
10
11 +use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
12 +use crate::reasoning::last_reasoning_boundary;
13 +use crate::tool::{Tool, ToolCallDelta};
14 +use crate::unified::parsing_failed;
15 +use crate::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len_mul};
16 +
17 +const REASONING_START: &str = "<|channel>thought\n";
18 +const CHANNEL_START: &str = "<|channel>";
19 +const CHANNEL_END: &str = "<channel|>";
20 const TOOL_CALL_START: &str = "<|tool_call>";
21 const TOOL_CALL_END: &str = "<tool_call|>";
22 const STRING_DELIM: &str = "<|\"|>";
@@ -20,6 +27,9 @@ type Gemma4Input<'i> = Partial<&'i str>;
27 #[derive(Debug, Clone, PartialEq)]
28 enum Gemma4Event {
29 Text { len: usize },
30 + Reasoning { len: usize },
31 + ReasoningStart,
32 + ReasoningEnd,
33 ToolCallStart,
34 ToolCallHeader { name: String },
35 ToolCall { args: Map<String, Value> },
@@ -35,6 +45,7 @@ struct Gemma4ArgsScanState {
45 enum Gemma4Mode {
46 #[default]
47 Text,
48 + Reasoning,
49 Header,
50 ToolCall {
51 name: String,
@@ -42,36 +53,62 @@ enum Gemma4Mode {
53 },
54 }
55
45 -/// Tool parser for Google Gemma4 models.
56 +/// Unified parser for Google Gemma4 models.
57 ///
58 /// Original Python implementation:
48 -/// <https://github.com/vllm-project/vllm/blob/bf45e6d0a558da2b8d7b60efb07b4aa394f3b60b/vllm/tool_parsers/gemma4_tool_parser.py>
59 +/// <https://github.com/vllm-project/vllm/blob/main/vllm/parser/gemma4.py>
60 +///
61 +/// Handles Gemma4 reasoning and function-call formats:
62 ///
50 -/// Handles the Gemma4 function call format:
63 +/// `<|channel>thought\nreasoning<channel|>`
64 ///
65 /// `<|tool_call>call:func_name{key:<|"|>value<|"|>}<tool_call|>`
66 ///
67 /// Arguments are emitted only after a full Gemma4 tool call is parsed.
55 -pub struct Gemma4ToolParser {
68 +pub struct Gemma4UnifiedParser {
69 buffer: String,
70 mode: Gemma4Mode,
71 emitted_tool_count: usize,
72 + tokenizer: DynTokenizer,
73 + channel_start_token_id: u32,
74 + channel_end_token_id: u32,
75 }
76
61 -impl Gemma4ToolParser {
62 - fn new(_tools: &[Tool]) -> Self {
63 - Self {
77 +impl Gemma4UnifiedParser {
78 + /// Create a Gemma4 parser.
79 + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
80 + let channel_start_token_id = tokenizer.token_to_id(CHANNEL_START).ok_or_else(|| {
81 + UnifiedParserError::MissingToken {
82 + token: CHANNEL_START.to_string(),
83 + }
84 + })?;
85 + let channel_end_token_id =
86 + tokenizer
87 + .token_to_id(CHANNEL_END)
88 + .ok_or_else(|| UnifiedParserError::MissingToken {
89 + token: CHANNEL_END.to_string(),
90 + })?;
91 +
92 + Ok(Self {
93 buffer: String::new(),
94 mode: Gemma4Mode::default(),
95 emitted_tool_count: 0,
67 - }
96 + channel_start_token_id,
97 + channel_end_token_id,
98 + tokenizer,
99 + })
100 }
101
70 - fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> {
102 + fn apply_event(&mut self, event: Gemma4Event, output: &mut UnifiedParserOutput) -> Result<()> {
103 match event {
104 Gemma4Event::Text { len: consumed_len } => {
73 - output.push_text(&self.buffer[..consumed_len]);
105 + output.push_text(self.buffer[..consumed_len].to_string());
106 }
107 + Gemma4Event::Reasoning { len: consumed_len } => {
108 + output.push_reasoning(self.buffer[..consumed_len].to_string());
109 + }
110 + Gemma4Event::ReasoningStart => self.mode = Gemma4Mode::Reasoning,
111 + Gemma4Event::ReasoningEnd => self.mode = Gemma4Mode::Text,
112 Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,
113 Gemma4Event::ToolCallHeader { name } => {
114 self.mode = Gemma4Mode::ToolCall {
@@ -100,9 +137,24 @@ impl Gemma4ToolParser {
137 Ok(())
138 }
139
140 + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) {
141 + self.mode = match last_reasoning_boundary(
142 + prompt_token_ids,
143 + self.channel_start_token_id,
144 + self.channel_end_token_id,
145 + self.tokenizer.as_ref(),
146 + ) {
147 + Some(true) => Gemma4Mode::Reasoning,
148 + Some(false) | None => Gemma4Mode::Text,
149 + };
150 + }
151 +
152 fn reset(&mut self) -> String {
153 let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) {
154 Gemma4Mode::Text => std::mem::take(&mut self.buffer),
155 + Gemma4Mode::Reasoning => {
156 + format!("{}{}", REASONING_START, std::mem::take(&mut self.buffer))
157 + }
158 Gemma4Mode::Header => {
159 format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer))
160 }
@@ -122,19 +174,26 @@ impl Gemma4ToolParser {
174 }
175 }
176
125 -impl ToolParser for Gemma4ToolParser {
126 - fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
177 +impl UnifiedParser for Gemma4UnifiedParser {
178 + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>>
179 where
180 Self: Sized + 'static,
181 {
130 - Ok(Box::new(Self::new(tools)))
182 + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box<dyn UnifiedParser>)
183 + }
184 +
185 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
186 + self.buffer.clear();
187 + self.emitted_tool_count = 0;
188 + self.initialize_mode(prompt_token_ids);
189 + Ok(())
190 }
191
192 fn preserve_special_tokens(&self) -> bool {
193 true
194 }
195
137 - fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
196 + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {
197 self.buffer.push_str(chunk);
198
199 while let Some((event, consumed_len)) = {
@@ -149,11 +208,12 @@ impl ToolParser for Gemma4ToolParser {
208 Ok(())
209 }
210
152 - fn finish(&mut self) -> Result<ToolParserOutput> {
153 - let mut output = ToolParserOutput::default();
211 + fn finish(&mut self) -> Result<UnifiedParserOutput> {
212 + let mut output = UnifiedParserOutput::default();
213
214 match &self.mode {
156 - Gemma4Mode::Text => output.push_text(&self.buffer),
215 + Gemma4Mode::Text => output.push_text(std::mem::take(&mut self.buffer)),
216 + Gemma4Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),
217 Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {
218 return Err(parsing_failed!("incomplete Gemma4 tool call"));
219 }
@@ -164,7 +224,7 @@ impl ToolParser for Gemma4ToolParser {
224 }
225
226 fn reset(&mut self) -> String {
167 - Gemma4ToolParser::reset(self)
227 + Gemma4UnifiedParser::reset(self)
228 }
229 }
230
@@ -175,6 +235,7 @@ fn parse_next_gemma4_event(
235 ) -> ModalResult<Gemma4Event> {
236 match mode {
237 Gemma4Mode::Text => parse_text_event(input),
238 + Gemma4Mode::Reasoning => parse_reasoning_event(input),
239 Gemma4Mode::Header => tool_call_header_event(input),
240 Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan),
241 }
@@ -182,7 +243,32 @@ fn parse_next_gemma4_event(
243
244 /// Parse a Gemma4 text-mode event.
245 fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
185 - alt((tool_call_start_event, safe_text_event)).parse_next(input)
246 + alt((
247 + reasoning_start_event,
248 + tool_call_start_event,
249 + safe_text_event,
250 + ))
251 + .parse_next(input)
252 +}
253 +
254 +/// Parse a Gemma4 reasoning-mode event.
255 +fn parse_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
256 + alt((
257 + reasoning_end_event,
258 + tool_call_start_event,
259 + safe_reasoning_event,
260 + ))
261 + .parse_next(input)
262 +}
263 +
264 +/// Parse a Gemma4 reasoning start marker.
265 +fn reasoning_start_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
266 + literal(REASONING_START).value(Gemma4Event::ReasoningStart).parse_next(input)
267 +}
268 +
269 +/// Parse a Gemma4 reasoning end marker.
270 +fn reasoning_end_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
271 + literal(CHANNEL_END).value(Gemma4Event::ReasoningEnd).parse_next(input)
272 }
273
274 /// Parse a Gemma4 tool-call start marker.
@@ -226,7 +312,14 @@ fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
312
313 /// Parse a safe text run before the next Gemma4 marker.
314 fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
229 - safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len })
315 + safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START])
316 + .map(|len| Gemma4Event::Text { len })
317 +}
318 +
319 +/// Parse a safe reasoning run before the next Gemma4 marker.
320 +fn safe_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
321 + safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START])
322 + .map(|len| Gemma4Event::Reasoning { len })
323 }
324
325 /// Parse raw Gemma4 arguments through the first end marker outside a Gemma string.
@@ -418,17 +511,145 @@ fn parse_gemma4_scalar(value: &str) -> Value {
511
512 #[cfg(test)]
513 mod tests {
514 + use std::sync::Arc;
515 +
516 use serde_json::{Value, json};
517 use thiserror_ext::AsReport;
518 + use vllm_tokenizer::Tokenizer;
519 use winnow::combinator::{eof, terminated};
520 use winnow::error::ErrMode;
521 use winnow::prelude::*;
522
523 use super::{
428 - Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content,
429 - parse_gemma4_args,
524 + CHANNEL_END, CHANNEL_START, Gemma4UnifiedParser, ToolCallDelta, UnifiedParser,
525 + UnifiedParserError, UnifiedParserOutput, gemma4_array_content, parse_gemma4_args,
526 };
431 - use crate::tool::{Tool, ToolParserTestExt as _};
527 + use crate::tool::Tool;
528 + use crate::unified::{UnifiedParserEvent, parsing_failed};
529 +
530 + struct FakeTokenizer;
531 +
532 + impl Tokenizer for FakeTokenizer {
533 + fn encode(
534 + &self,
535 + text: &str,
536 + _add_special_tokens: bool,
537 + ) -> vllm_tokenizer::Result<Vec<u32>> {
538 + Ok(text.chars().map(u32::from).collect())
539 + }
540 +
541 + fn decode(
542 + &self,
543 + token_ids: &[u32],
544 + _skip_special_tokens: bool,
545 + ) -> vllm_tokenizer::Result<String> {
546 + Ok(token_ids
547 + .iter()
548 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
549 + .collect())
550 + }
551 +
552 + fn token_to_id(&self, token: &str) -> Option<u32> {
553 + match token {
554 + CHANNEL_START => Some(100),
555 + CHANNEL_END => Some(101),
556 + _ => None,
557 + }
558 + }
559 +
560 + fn is_special_id(&self, token_id: u32) -> bool {
561 + matches!(token_id, 100..=105)
562 + }
563 + }
564 +
565 + struct MissingTokenTokenizer;
566 +
567 + impl Tokenizer for MissingTokenTokenizer {
568 + fn encode(
569 + &self,
570 + text: &str,
571 + _add_special_tokens: bool,
572 + ) -> vllm_tokenizer::Result<Vec<u32>> {
573 + Ok(text.chars().map(u32::from).collect())
574 + }
575 +
576 + fn decode(
577 + &self,
578 + token_ids: &[u32],
579 + _skip_special_tokens: bool,
580 + ) -> vllm_tokenizer::Result<String> {
581 + Ok(token_ids
582 + .iter()
583 + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
584 + .collect())
585 + }
586 +
587 + fn token_to_id(&self, _token: &str) -> Option<u32> {
588 + None
589 + }
590 + }
591 +
592 + trait UnifiedParserTestExt {
593 + fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput>;
594 + fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput>;
595 + }
596 +
597 + impl UnifiedParserTestExt for Gemma4UnifiedParser {
598 + fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {
599 + let mut output = UnifiedParserOutput::default();
600 + self.parse_into(chunk, &mut output)?;
601 + Ok(output)
602 + }
603 +
604 + fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput> {
605 + let mut output = self.parse_chunk(text)?;
606 + output.append(self.finish()?);
607 + Ok(output)
608 + }
609 + }
610 +
611 + trait UnifiedOutputTestExt {
612 + fn normal_text(&self) -> String;
613 + fn reasoning_text(&self) -> String;
614 + fn calls(&self) -> Vec<&ToolCallDelta>;
615 + fn coalesce(self) -> Self;
616 + }
617 +
618 + impl UnifiedOutputTestExt for UnifiedParserOutput {
619 + fn normal_text(&self) -> String {
620 + self.events
621 + .iter()
622 + .filter_map(|event| match event {
623 + UnifiedParserEvent::Text(text) => Some(text.as_str()),
624 + UnifiedParserEvent::Reasoning(_) | UnifiedParserEvent::ToolCall(_) => None,
625 + })
626 + .collect()
627 + }
628 +
629 + fn reasoning_text(&self) -> String {
630 + self.events
631 + .iter()
632 + .filter_map(|event| match event {
633 + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),
634 + UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None,
635 + })
636 + .collect()
637 + }
638 +
639 + fn calls(&self) -> Vec<&ToolCallDelta> {
640 + self.events
641 + .iter()
642 + .filter_map(|event| match event {
643 + UnifiedParserEvent::Text(_) | UnifiedParserEvent::Reasoning(_) => None,
644 + UnifiedParserEvent::ToolCall(call) => Some(call),
645 + })
646 + .collect()
647 + }
648 +
649 + fn coalesce(self) -> Self {
650 + self
651 + }
652 + }
653
654 fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
655 let mut input = array;
@@ -494,9 +715,26 @@ mod tests {
715 ]
716 }
717
497 - fn collect_stream(chunks: &[&str]) -> ToolParserOutput {
498 - let mut parser = Gemma4ToolParser::new(&test_tools());
499 - let mut output = ToolParserOutput::default();
718 + fn test_parser() -> Gemma4UnifiedParser {
719 + Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap()
720 + }
721 +
722 + #[test]
723 + fn gemma4_create_requires_channel_start_token() {
724 + let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) {
725 + Ok(_) => panic!("expected missing token error"),
726 + Err(error) => error,
727 + };
728 +
729 + assert!(matches!(
730 + error,
731 + UnifiedParserError::MissingToken { token } if token == CHANNEL_START
732 + ));
733 + }
734 +
735 + fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput {
736 + let mut parser = test_parser();
737 + let mut output = UnifiedParserOutput::default();
738 for chunk in chunks {
739 output.append(parser.parse_chunk(chunk).unwrap());
740 }
@@ -504,7 +742,7 @@ mod tests {
742 output.coalesce()
743 }
744
507 - fn first_call(output: &ToolParserOutput) -> ToolCallDelta {
745 + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta {
746 (*output.calls().first().expect("expected one tool call")).clone()
747 }
748
@@ -542,7 +780,7 @@ mod tests {
780
781 #[test]
782 fn gemma4_parse_complete_extracts_single_tool_call() {
545 - let mut parser = Gemma4ToolParser::new(&test_tools());
783 + let mut parser = test_parser();
784 let output = parser
785 .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
786 .unwrap();
@@ -558,7 +796,7 @@ mod tests {
796
797 #[test]
798 fn gemma4_parse_complete_rejects_incomplete_tool_call() {
561 - let mut parser = Gemma4ToolParser::new(&test_tools());
799 + let mut parser = test_parser();
800 let error = parser
801 .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London")
802 .unwrap_err();
@@ -607,8 +845,8 @@ mod tests {
845
846 #[test]
847 fn gemma4_streaming_waits_for_complete_tool_call() {
610 - let mut parser = Gemma4ToolParser::new(&test_tools());
611 - let mut output = ToolParserOutput::default();
848 + let mut parser = test_parser();
849 + let mut output = UnifiedParserOutput::default();
850
851 for chunk in [
852 "<|tool_call>",
@@ -773,7 +1011,7 @@ mod tests {
1011
1012 #[test]
1013 fn gemma4_finish_flushes_partial_start_marker_as_text() {
776 - let mut parser = Gemma4ToolParser::new(&test_tools());
1014 + let mut parser = test_parser();
1015 let mut output = parser.parse_chunk("<").unwrap();
1016 output.append(parser.finish().unwrap());
1017
@@ -781,9 +1019,104 @@ mod tests {
1019 assert!(output.calls().is_empty());
1020 }
1021
1022 + #[test]
1023 + fn gemma4_streaming_emits_reasoning_then_text() {
1024 + let output = collect_stream(&["<|channel>thought\nreason<channel|>answer"]);
1025 +
1026 + assert_eq!(output.reasoning_text(), "reason");
1027 + assert_eq!(output.normal_text(), "answer");
1028 + assert!(output.calls().is_empty());
1029 + }
1030 +
1031 + #[test]
1032 + fn gemma4_streaming_holds_split_reasoning_start() {
1033 + let mut parser = test_parser();
1034 +
1035 + let first = parser.parse_chunk("<|channel>").unwrap();
1036 + assert!(first.events.is_empty());
1037 +
1038 + let mut output = parser.parse_chunk("thought\nrea").unwrap();
1039 + output.append(parser.parse_chunk("son<channel|>answer").unwrap());
1040 + output.append(parser.finish().unwrap());
1041 +
1042 + assert_eq!(output.reasoning_text(), "reason");
1043 + assert_eq!(output.normal_text(), "answer");
1044 + }
1045 +
1046 + #[test]
1047 + fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() {
1048 + let mut parser = test_parser();
1049 + parser.initialize(&[100, 3000, 3001]).unwrap();
1050 +
1051 + let output = parser.parse_complete("reason<channel|>answer").unwrap();
1052 +
1053 + assert_eq!(output.reasoning_text(), "reason");
1054 + assert_eq!(output.normal_text(), "answer");
1055 + }
1056 +
1057 + #[test]
1058 + fn gemma4_initialize_turn_prompt_starts_in_text() {
1059 + let mut parser = test_parser();
1060 + parser.initialize(&[104, 3000, 3001]).unwrap();
1061 +
1062 + let output = parser.parse_complete("<|channel>thought\nreason<channel|>answer").unwrap();
1063 +
1064 + assert_eq!(output.reasoning_text(), "reason");
1065 + assert_eq!(output.normal_text(), "answer");
1066 + }
1067 +
1068 + #[test]
1069 + fn gemma4_initialize_special_token_caps_boundary_scan() {
1070 + let mut parser = test_parser();
1071 + parser.initialize(&[100, 3000, 104, 3001]).unwrap();
1072 +
1073 + let output = parser.parse_complete("answer").unwrap();
1074 +
1075 + assert!(output.reasoning_text().is_empty());
1076 + assert_eq!(output.normal_text(), "answer");
1077 + }
1078 +
1079 + #[test]
1080 + fn gemma4_initialize_closed_channel_prompt_starts_in_text() {
1081 + let mut parser = test_parser();
1082 + parser.initialize(&[100, 3000, 3001, 101]).unwrap();
1083 +
1084 + let output = parser.parse_complete("answer").unwrap();
1085 +
1086 + assert!(output.reasoning_text().is_empty());
1087 + assert_eq!(output.normal_text(), "answer");
1088 + }
1089 +
1090 + #[test]
1091 + fn gemma4_reasoning_tool_call_implicitly_ends_reasoning() {
1092 + let output = collect_stream(&[
1093 + "<|channel>thought\nNeed weather.",
1094 + "<|tool_call>",
1095 + "call:get_weather{location:<|\"|>Paris<|\"|>}",
1096 + "<tool_call|>",
1097 + ]);
1098 +
1099 + assert_eq!(output.reasoning_text(), "Need weather.");
1100 + assert!(output.normal_text().is_empty());
1101 + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
1102 + assert_eq!(
1103 + serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
1104 + json!({ "location": "Paris" })
1105 + );
1106 + }
1107 +
1108 + #[test]
1109 + fn gemma4_bare_channel_start_is_plain_text() {
1110 + let output = collect_stream(&["<|channel>plain"]);
1111 +
1112 + assert_eq!(output.normal_text(), "<|channel>plain");
1113 + assert!(output.reasoning_text().is_empty());
1114 + assert!(output.calls().is_empty());
1115 + }
1116 +
1117 #[test]
1118 fn gemma4_finish_rejects_complete_args_without_end_marker() {
786 - let mut parser = Gemma4ToolParser::new(&test_tools());
1119 + let mut parser = test_parser();
1120 for chunk in ["<|tool_call>", "call:get_status{}"] {
1121 parser.parse_chunk(chunk).unwrap();
1122 }
@@ -795,7 +1128,7 @@ mod tests {
1128
1129 #[test]
1130 fn gemma4_reset_preserves_internally_buffered_arguments() {
798 - let mut parser = Gemma4ToolParser::new(&test_tools());
1131 + let mut parser = test_parser();
1132 for chunk in [
1133 "<|tool_call>",
1134 "call:write_file{",
@@ -815,7 +1148,7 @@ mod tests {
1148
1149 #[test]
1150 fn gemma4_reset_preserves_completed_arguments_after_parse_error() {
818 - let mut parser = Gemma4ToolParser::new(&test_tools());
1151 + let mut parser = test_parser();
1152 let input = "<|tool_call>call:set{broken}<tool_call|>";
1153
1154 let _error = parser.parse_chunk(input).unwrap_err();
rust/src/parser/src/unified/mod.rs
+9 -1
@@ -1,11 +1,14 @@
1 //! Unified parser interface for reasoning and tool-call deltas.
2
3 mod combined;
4 +mod gemma4;
5
6 use thiserror::Error;
7 +use thiserror_ext::Macro;
8 use vllm_tokenizer::DynTokenizer;
9
10 pub use combined::CombinedParser;
11 +pub use gemma4::Gemma4UnifiedParser;
12
13 use crate::reasoning::ReasoningError;
14 use crate::tool::{
@@ -189,10 +192,15 @@ pub trait UnifiedParser: Send {
192 }
193
194 /// Errors produced while creating or running unified parsers.
192 -#[derive(Debug, Error)]
195 +#[derive(Debug, Error, Macro)]
196 +#[thiserror_ext(macro(path = "crate::unified", mangle))]
197 pub enum UnifiedParserError {
198 #[error("combined parser is constructed from split parser instances")]
199 CombinedParserConstructor,
200 + #[error("tokenizer is missing unified parser token `{token}`")]
201 + MissingToken { token: String },
202 + #[error("unified parser parsing failed: {message}")]
203 + ParsingFailed { message: String },
204 #[error(transparent)]
205 Reasoning(#[from] ReasoningError),
206 #[error(transparent)]
rust/src/parser/src/utils.rs renamed
+103 -16
@@ -1,10 +1,10 @@
1 -//! Shared helpers for tool parsers.
1 +//! Shared helpers for streaming parsers.
2
3 use winnow::Parser;
4 use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue};
5 -use winnow::stream::{Offset, Partial, Stream};
5 +use winnow::stream::{FindSlice, Offset, Partial, Stream};
6
7 -use super::Result;
7 +use crate::tool::{Result, ToolParserError};
8
9 /// Return the byte length of the longest proper prefix of `token` that is also
10 /// a suffix of `buffer`.
@@ -15,7 +15,7 @@ use super::Result;
15 /// The returned length is always a valid UTF-8 boundary in `token`, so callers
16 /// can safely slice `&token[..len]` even when markers contain non-ASCII
17 /// characters such as DeepSeek's DSML delimiters.
18 -pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize {
18 +pub fn partial_prefix_len(buffer: &str, token: &str) -> usize {
19 let Some(first_byte) = token.as_bytes().first().copied() else {
20 return 0;
21 };
@@ -44,9 +44,10 @@ pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize {
44 }
45
46 /// Parse a safe text run before the next marker.
47 +/// This is the single-marker variant of [`safe_text_len_mul`].
48 ///
49 /// Returns the text length in bytes, and advances the input.
49 -pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult<usize> {
50 +pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult<usize> {
51 let text = **input;
52 if text.is_empty() {
53 return incomplete();
@@ -67,15 +68,53 @@ pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalRes
68 Ok(emit_len)
69 }
70
71 +/// Parse a safe text run before the earliest next marker.
72 +/// This is the multi-marker variant of [`safe_text_len`].
73 +///
74 +/// Returns the text length in bytes, and advances the input.
75 +pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalResult<usize> {
76 + let text = **input;
77 + if text.is_empty() {
78 + return incomplete();
79 + }
80 +
81 + if let Some(start_idx) = find_slice_mul(text, markers) {
82 + input.next_slice(start_idx);
83 + return Ok(start_idx);
84 + }
85 +
86 + let keep_len = markers.iter().map(|marker| partial_prefix_len(text, marker)).max().unwrap_or(0);
87 + let emit_len = text.len().saturating_sub(keep_len);
88 + if emit_len == 0 {
89 + return incomplete();
90 + }
91 +
92 + input.next_slice(emit_len);
93 + Ok(emit_len)
94 +}
95 +
96 +#[inline(always)]
97 +fn find_slice_mul(text: &str, markers: &[&str]) -> Option<usize> {
98 + let range = match markers {
99 + // Use the fast specialized `winnow::stream::FindSlice` impl for 1-3 markers.
100 + [first] => text.find_slice(*first),
101 + [first, second] => text.find_slice((*first, *second)),
102 + [first, second, third] => text.find_slice((*first, *second, *third)),
103 + // Fall back to a linear scan for 4+ markers.
104 + _ => return markers.iter().filter_map(|marker| text.find(marker)).min(),
105 + };
106 + range.map(|range| range.start)
107 +}
108 +
109 /// Streaming scan state for a buffered marker search [`take_until_marker`],
110 /// so that we don't have to rescan the whole buffered prefix when resuming.
111 #[derive(Debug, Clone, Default, PartialEq, Eq)]
73 -pub(super) struct MarkerScanState {
112 +pub struct MarkerScanState {
113 scan_start: usize,
114 }
115
116 impl MarkerScanState {
78 - pub(super) fn reset(&mut self) {
117 + pub fn reset(&mut self) {
118 self.scan_start = 0;
119 }
120 }
@@ -92,7 +131,7 @@ impl MarkerScanState {
131 /// chunks while waiting for a closing marker. Plain `take_until` is still a
132 /// better fit for one-shot parsers over a complete body, and for `1..` cases
133 /// where an empty slice before the marker should be rejected.
95 -pub(super) fn take_until_marker<'i, 'a>(
134 +pub fn take_until_marker<'i, 'a>(
135 marker: &'a str,
136 state: &'a mut MarkerScanState,
137 ) -> impl Parser<Partial<&'i str>, &'i str, ErrMode<ContextError>> + 'a {
@@ -137,7 +176,7 @@ fn floor_char_boundary(text: &str, index: usize) -> usize {
176
177 /// Streaming lexical state for a top-level JSON object.
178 #[derive(Debug, Clone, Default, PartialEq, Eq)]
140 -pub(super) struct JsonObjectScanState {
179 +pub struct JsonObjectScanState {
180 object_depth: usize,
181 array_depth: usize,
182 in_string: bool,
@@ -155,7 +194,7 @@ enum JsonObjectScanPhase {
194
195 impl JsonObjectScanState {
196 /// Returns whether the top-level JSON object has closed.
158 - pub(super) const fn complete(&self) -> bool {
197 + pub const fn complete(&self) -> bool {
198 matches!(self.phase, JsonObjectScanPhase::Complete)
199 }
200 }
@@ -165,7 +204,7 @@ impl JsonObjectScanState {
204 /// The returned length is safe to emit as raw argument text. This scans only
205 /// lexical boundaries from `{` through the matching `}`, preserving
206 /// malformed-but-balanced JSON without deserializing or normalizing it.
168 -pub(super) fn take_json_object(
207 +pub fn take_json_object(
208 input: &mut Partial<&str>,
209 state: &mut JsonObjectScanState,
210 ) -> ModalResult<usize> {
@@ -252,7 +291,7 @@ pub(super) fn take_json_object(
291 }
292
293 /// Parse a JSON string literal.
255 -pub(super) fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
294 +pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
295 let text = **input;
296 if text.is_empty() {
297 return incomplete();
@@ -311,7 +350,7 @@ fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode<Co
350 /// of bytes consumed from the buffer.
351 /// - `Ok(None)` if the buffer does not contain a full event yet, and more data is needed.
352 /// - `Err` if a parsing error occurred.
314 -pub(super) fn parse_buffered_event<E>(
353 +pub fn parse_buffered_event<E>(
354 buffer: &str,
355 parse: impl FnOnce(&mut Partial<&str>) -> ModalResult<E>,
356 ) -> Result<Option<(E, usize)>> {
@@ -322,7 +361,9 @@ pub(super) fn parse_buffered_event<E>(
361 Err(ErrMode::Incomplete(_)) => return Ok(None),
362 Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => {
363 // TODO: enrich context for error reporting
325 - return Err(parsing_failed!("{}", e));
364 + return Err(ToolParserError::ParsingFailed {
365 + message: e.to_string(),
366 + });
367 }
368 };
369 let consumed_len = input.offset_from(&checkpoint);
@@ -334,7 +375,7 @@ pub(super) fn parse_buffered_event<E>(
375 }
376
377 /// Returns an error indicating that we need more data to continue parsing.
337 -pub(super) fn incomplete<T>() -> ModalResult<T> {
378 +pub fn incomplete<T>() -> ModalResult<T> {
379 Err(ErrMode::Incomplete(Needed::Unknown))
380 }
381
@@ -348,7 +389,7 @@ mod tests {
389
390 use super::{
391 JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len,
351 - take_json_object, take_until_marker,
392 + safe_text_len_mul, take_json_object, take_until_marker,
393 };
394
395 #[test]
@@ -406,6 +447,52 @@ mod tests {
447 assert!(matches!(error, ErrMode::Incomplete(_)));
448 }
449
450 + #[test]
451 + fn safe_text_len_mul_stops_before_earliest_marker() {
452 + let mut input = Partial::new("hello<channel|><|tool_call>");
453 + let checkpoint = input.checkpoint();
454 +
455 + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<channel|>"]).unwrap();
456 +
457 + assert_eq!(len, "hello".len());
458 + assert_eq!(input.offset_from(&checkpoint), "hello".len());
459 + assert_eq!(*input, "<channel|><|tool_call>");
460 + }
461 +
462 + #[test]
463 + fn safe_text_len_mul_holds_back_longest_partial_marker() {
464 + let mut input = Partial::new("hello<|tool");
465 + let checkpoint = input.checkpoint();
466 +
467 + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap();
468 +
469 + assert_eq!(len, "hello".len());
470 + assert_eq!(input.offset_from(&checkpoint), "hello".len());
471 + assert_eq!(*input, "<|tool");
472 + }
473 +
474 + #[test]
475 + fn safe_text_len_mul_skips_false_same_prefix_candidate() {
476 + let mut input = Partial::new("hello<not_marker><|tool_call>");
477 + let checkpoint = input.checkpoint();
478 +
479 + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap();
480 +
481 + assert_eq!(len, "hello<not_marker>".len());
482 + assert_eq!(input.offset_from(&checkpoint), "hello<not_marker>".len());
483 + assert_eq!(*input, "<|tool_call>");
484 + }
485 +
486 + #[test]
487 + fn safe_text_len_mul_reports_incomplete_for_only_partial_marker() {
488 + let mut input = Partial::new("<|channel>thought");
489 +
490 + let error =
491 + safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap_err();
492 +
493 + assert!(matches!(error, ErrMode::Incomplete(_)));
494 + }
495 +
496 #[test]
497 fn take_until_marker_stops_before_marker() {
498 let mut state = MarkerScanState::default();