diff --git a/pyproject.toml b/pyproject.toml index 249832ff2e5..3819ad7fc8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/parser/src/tool/gemma4.rs", "rust/src/text/src/output/decoded.rs", + "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/text/src/output/decoded.rs", "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] ignore-hidden = false diff --git a/rust/Cargo.toml b/rust/Cargo.toml index dc3895c372d..601e009df07 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -120,7 +120,7 @@ vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } -winnow = "1.0.2" +winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.1.0" zeromq = { version = "0.6.0", default-features = false, features = [ "tokio-runtime", diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 8c9a4362d1c..c494df600f4 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -18,7 +18,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent use crate::parser::ParserSelection; use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; -use crate::request::ChatRequest; +use crate::parser::unified::UnifiedParserFactory; +use crate::request::{ChatRequest, ChatTool}; use crate::{Error, Result as ChatResult}; /// Default request-scoped output processor used by Hugging Face style chat @@ -46,20 +47,31 @@ impl DefaultChatOutputProcessor { tool_call_parser: &ParserSelection, reasoning_parser: &ParserSelection, ) -> ChatResult { - let tool_parsing_enabled = request.tool_parsing_enabled(); - let tool_parser = if tool_parsing_enabled { - Some(Self::resolve_tool_parser( - request, + let parser = if tool_call_parser == reasoning_parser + && let Some(parser) = Self::resolve_optional_unified_parser( + &request.tools, model_id, + tokenizer.clone(), tool_call_parser, - )?) + )? { + parser } else { - None + let tool_parsing_enabled = request.tool_parsing_enabled(); + let tool_parser = if tool_parsing_enabled { + Some(Self::resolve_tool_parser( + &request.tools, + model_id, + tool_call_parser, + )?) + } else { + None + }; + let reasoning_parser = + Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; + Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box }; - let reasoning_parser = - Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; - let parser: Box = - Box::new(CombinedParser::new(reasoning_parser, tool_parser)); + + apply_structural_tag_constraint(request, parser.structural_tag_model())?; if parser.preserve_special_tokens() { request.decode_options.skip_special_tokens = false; @@ -84,7 +96,7 @@ impl DefaultChatOutputProcessor { } fn resolve_tool_parser( - request: &mut ChatRequest, + tools: &[ChatTool], model_id: &str, selection: &ParserSelection, ) -> ChatResult> { @@ -100,14 +112,36 @@ impl DefaultChatOutputProcessor { ParserSelection::Explicit(name) => name.as_str(), }; - let parser = factory.create(parser_name, &request.tools)?; - - apply_structural_tag_constraint(request, parser.as_ref())?; + let parser = factory.create(parser_name, tools)?; TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser")); Ok(parser) } + fn resolve_optional_unified_parser( + tools: &[ChatTool], + model_id: &str, + tokenizer: DynTokenizer, + selection: &ParserSelection, + ) -> ChatResult>> { + let factory = UnifiedParserFactory::global(); + let parser_name = match selection { + ParserSelection::Auto => factory.resolve_name_for_model(model_id), + ParserSelection::None => None, + ParserSelection::Explicit(name) if factory.contains(name) => Some(name.as_str()), + ParserSelection::Explicit(_) => None, + }; + + let Some(parser_name) = parser_name else { + return Ok(None); + }; + + let parser = factory.create(parser_name, tools, tokenizer)?; + + UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser")); + Ok(Some(parser)) + } + fn resolve_optional_reasoning_parser( model_id: &str, tokenizer: DynTokenizer, @@ -134,6 +168,7 @@ impl DefaultChatOutputProcessor { static TOOL_PARSER_LOG_ONCE: Once = Once::new(); static REASONING_PARSER_LOG_ONCE: Once = Once::new(); +static UNIFIED_PARSER_LOG_ONCE: Once = Once::new(); impl ChatOutputProcessor for DefaultChatOutputProcessor { /// Transforms a raw generate-output token stream into structured chat @@ -149,3 +184,102 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor { Ok(structured.boxed()) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::Tokenizer; + + use super::DefaultChatOutputProcessor; + use crate::Error; + use crate::parser::ParserSelection; + use crate::request::ChatRequest; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "<|channel>" => Some(1), + "" => Some(2), + _ => None, + } + } + } + + fn tokenizer() -> Arc { + Arc::new(FakeTokenizer) + } + + #[test] + fn equal_explicit_gemma4_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + let selection = ParserSelection::Explicit("gemma4".to_string()); + + DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &selection, + &selection, + ) + .unwrap(); + } + + #[test] + fn auto_auto_gemma4_model_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + + DefaultChatOutputProcessor::new( + &mut request, + "google/gemma-4-27b-it", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Auto, + ) + .unwrap(); + } + + #[test] + fn mixed_gemma4_selection_uses_split_dummy_error() { + let mut request = ChatRequest::for_test(); + let error = match DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Explicit("gemma4".to_string()), + ) { + Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"), + Err(error) => error, + }; + + let Error::ParserInitialization { error, .. } = error else { + panic!("expected parser initialization error"); + }; + assert_eq!( + error.to_string(), + "`gemma4` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + ); + } +} diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index bdc5d8e14c2..6ba2458ca8d 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -2,12 +2,12 @@ use thiserror_ext::AsReport; use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams}; +use vllm_parser::tool::StructuralTagModel; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, build_structural_tag, }; -use crate::parser::tool::ToolParser; use crate::request::{ChatRequest, ChatToolChoice}; use crate::{Error, Result as ChatResult}; @@ -15,9 +15,9 @@ use crate::{Error, Result as ChatResult}; /// support and the request's tool choice. pub(super) fn apply_structural_tag_constraint( request: &mut ChatRequest, - parser: &dyn ToolParser, + model: Option, ) -> ChatResult<()> { - let Some(model) = parser.structural_tag_model() else { + let Some(model) = model else { return Ok(()); }; let Some(tool_choice) = structural_tag_tool_choice(request) else { @@ -77,7 +77,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option vllm_parser::reasoning::Result>; +type ReasoningParserCreator = Arc< + dyn Fn(DynTokenizer) -> vllm_parser::reasoning::Result> + Send + Sync, +>; /// Registry and model matcher for reasoning parsers. pub type ReasoningParserFactory = ParserFactory; @@ -58,7 +59,7 @@ impl ReasoningParserFactory { .register_parser::(names::DEEPSEEK_R1) .register_parser::(names::DEEPSEEK_V3) .register_parser::(names::DEEPSEEK_V4) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) @@ -109,7 +110,17 @@ impl ReasoningParserFactory { where T: ReasoningParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split reasoning registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ReasoningError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -124,7 +135,7 @@ impl ReasoningParserFactory { available_names: self.list(), })?; - creator(tokenizer).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tokenizer).map_err(|error| crate::Error::ParserInitialization { kind: "reasoning", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/reasoning/tests.rs b/rust/src/chat/src/parser/reasoning/tests.rs index 58d987770c6..e6255d14a00 100644 --- a/rust/src/chat/src/parser/reasoning/tests.rs +++ b/rust/src/chat/src/parser/reasoning/tests.rs @@ -35,11 +35,13 @@ fn factory_contains_and_lists_registered_parsers() { assert!(factory.contains(names::SEED_OSS)); assert!(factory.contains(names::STEP3P5)); assert!(factory.contains(names::MINIMAX_M3)); + assert!(factory.contains(names::GEMMA4)); assert!(factory.list().contains(&names::QWEN3.to_string())); assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string())); assert!(factory.list().contains(&names::SEED_OSS.to_string())); assert!(factory.list().contains(&names::STEP3P5.to_string())); assert!(factory.list().contains(&names::MINIMAX_M3.to_string())); + assert!(factory.list().contains(&names::GEMMA4.to_string())); } #[test] diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 9884d1aca2a..a156d670248 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -1,13 +1,13 @@ //! Tool parser registration and selection boundary for `vllm-chat`. -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; pub use vllm_parser::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, - Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, - HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, - MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolParser, ToolParserError, + Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, + Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, + MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, + Qwen3XmlToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -40,7 +40,8 @@ pub mod names { } /// Constructor signature for one registered tool parser implementation. -type ToolParserCreator = fn(&[ChatTool]) -> vllm_parser::tool::Result>; +type ToolParserCreator = + Arc vllm_parser::tool::Result> + Send + Sync>; /// Registry and model matcher for tool parsers. pub type ToolParserFactory = ParserFactory; @@ -65,7 +66,7 @@ impl ToolParserFactory { .register_parser::(names::DEEPSEEK_V4) .register_parser::(names::GLM45) .register_parser::(names::GLM47) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) .register_parser::(names::GRANITE4) .register_parser::(names::HERMES) .register_parser::(names::HY_V3) @@ -126,7 +127,17 @@ impl ToolParserFactory { where T: ToolParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split tool registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ToolParserError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -137,7 +148,7 @@ impl ToolParserFactory { available_names: self.list(), })?; - creator(tools).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tools).map_err(|error| crate::Error::ParserInitialization { kind: "tool", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs new file mode 100644 index 00000000000..6456cfda754 --- /dev/null +++ b/rust/src/chat/src/parser/unified.rs @@ -0,0 +1,124 @@ +//! Unified parser registration and selection boundary for `vllm-chat`. + +use std::sync::LazyLock; + +pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser}; +use vllm_tokenizer::DynTokenizer; + +use crate::parser::ParserFactory; +use crate::request::ChatTool; + +/// Canonical public names for registered unified parsers. +pub mod names { + pub const GEMMA4: &str = "gemma4"; +} + +/// Constructor signature for one registered unified parser implementation. +type UnifiedParserCreator = + fn(&[ChatTool], DynTokenizer) -> vllm_parser::unified::Result>; + +/// Registry and model matcher for unified parsers. +pub type UnifiedParserFactory = ParserFactory; + +impl UnifiedParserFactory { + /// Get the global unified parser factory with built-in registrations and + /// model mappings. + pub fn global() -> &'static Self { + static INSTANCE: LazyLock = LazyLock::new(UnifiedParserFactory::new); + &INSTANCE + } + + /// Create the default registry with built-in parser names and model + /// mappings. + pub fn new() -> Self { + let mut factory = Self::default(); + + factory.register_parser::(names::GEMMA4); + + factory + .register_pattern("gemma-4", names::GEMMA4) + .register_pattern("gemma4", names::GEMMA4); + + factory + } + + /// Register one parser type that exposes a static `create()` constructor. + pub fn register_parser(&mut self, name: &str) -> &mut Self + where + T: UnifiedParser + 'static, + { + self.register_creator(name, T::create) + } + + /// Construct a parser from an exact name. + pub fn create( + &self, + name: &str, + tools: &[ChatTool], + tokenizer: DynTokenizer, + ) -> crate::Result> { + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName { + kind: "unified", + name: name.to_string(), + available_names: self.list(), + })?; + + creator(tools, tokenizer).map_err(|error| crate::Error::ParserInitialization { + kind: "unified", + name: name.to_string(), + error: error.into(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::Tokenizer; + + use super::{UnifiedParserFactory, names}; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "<|channel>" => Some(1), + "" => Some(2), + _ => None, + } + } + } + + #[test] + fn factory_registers_gemma4() { + let factory = UnifiedParserFactory::new(); + + assert!(factory.contains(names::GEMMA4)); + assert_eq!( + factory.resolve_name_for_model("google/gemma-4-27b-it"), + Some(names::GEMMA4) + ); + factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap(); + } +} diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index b3d5d9eae34..15bd4aca23a 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -37,6 +37,8 @@ struct RoundtripCase { /// JSON formatting expected after this model's template has materialized /// tool-call arguments. json_fmt: JsonFmt, + /// Whether the template renders tool-call argument object keys in sorted order. + sort_json_keys: bool, } #[derive(Clone, Copy)] @@ -81,6 +83,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), + sort_json_keys: false, } } @@ -93,6 +96,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -105,6 +109,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -117,6 +122,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -129,6 +135,20 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + + /// Gemma4 channel reasoning with custom function-call arguments. + fn gemma4() -> Self { + Self { + model_id: "google/gemma-4-E4B-it", + assistant_stop_suffix: "<|tool_response>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Always { value: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: true, } } @@ -142,6 +162,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), + sort_json_keys: false, } } @@ -154,6 +175,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } @@ -166,6 +188,7 @@ impl RoundtripCase { reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), + sort_json_keys: false, } } } @@ -195,8 +218,8 @@ roundtrip_tests! { seed_oss => [reasoning_and_content], step3p5 => [reasoning_and_content], - // Note: Kimi K2.5 strips the reasoning content in history. - kimi_k25 => [tool_call_mix], + gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call + kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history } /// Run the fixed reasoning+content fixture for one model/parser case. @@ -347,14 +370,38 @@ fn spaced_json_fmt() -> JsonFmt { /// Pass in a raw JSON string instead of a structured value to ensure the exact precision and /// formatting of numbers are preserved. fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result { - let value: serde_json::Value = + let mut value: serde_json::Value = serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?; + if case.sort_json_keys { + sort_json_value(&mut value); + } case.json_fmt .format_to_string(&value) .context("failed to format expected tool-call arguments") } +/// Sort JSON object keys recursively to match templates that render mappings with `dictsort`. +fn sort_json_value(value: &mut serde_json::Value) { + match value { + serde_json::Value::Object(map) => { + for value in map.values_mut() { + sort_json_value(value); + } + + let mut entries = std::mem::take(map).into_iter().collect::>(); + entries.sort_by(|(left, _), (right, _)| left.cmp(right)); + map.extend(entries); + } + serde_json::Value::Array(values) => { + for value in values { + sort_json_value(value); + } + } + _ => {} + } +} + /// Load the real model chat/text backend for one roundtrip case. async fn load_roundtrip_backends(case: &RoundtripCase) -> Result { load_model_backends( diff --git a/rust/src/parser/benches/gemma4.rs b/rust/src/parser/benches/gemma4.rs index 761f8d4e235..fd29e77a9a2 100644 --- a/rust/src/parser/benches/gemma4.rs +++ b/rust/src/parser/benches/gemma4.rs @@ -2,10 +2,11 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; -use vllm_parser::tool::{Gemma4ToolParser, Tool, ToolParser}; +use vllm_parser::tool::{Tool, ToolParser}; +use vllm_parser::unified::Gemma4UnifiedParser; mod utils; -use utils::feed_parser; +use utils::{UnifiedToolParserAdapter, feed_parser}; const CHUNK_CHARS: usize = 7; const LONG_NORMAL_TEXT_REPEATS: usize = 2048; @@ -68,7 +69,8 @@ fn long_tool_argument_fixture() -> String { } fn parser(tools: &[Tool]) -> Box { - Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize") + UnifiedToolParserAdapter::::create(tools) + .expect("Gemma4 unified parser should initialize") } fn run_stream_group( diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs new file mode 100644 index 00000000000..103ed18d093 --- /dev/null +++ b/rust/src/parser/benches/utils/adapter.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use vllm_parser::tool::{ + Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput, +}; +use vllm_parser::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, +}; +use vllm_tokenizer::Tokenizer; + +/// Tokenizer stub used by unified-parser benchmarks. +struct BenchTokenizer; + +impl Tokenizer for BenchTokenizer { + fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result> { + Ok(text.chars().map(|_| u32::MAX).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok("\u{FFFD}".repeat(token_ids.len())) + } + + fn token_to_id(&self, _token: &str) -> Option { + Some(u32::MAX) + } +} + +/// Bench-only adapter that exposes a unified parser through the tool-parser +/// benchmark harness. +/// +/// Returns error if the unified parser produces reasoning events. +pub struct UnifiedToolParserAdapter { + inner: Box, + _marker: std::marker::PhantomData, +} + +fn map_unified_error(error: UnifiedParserError) -> ToolParserError { + ToolParserError::ParsingFailed { + message: format!("unified parser failed: {error}"), + } +} + +fn append_unified_output( + output: UnifiedParserOutput, + tool_output: &mut ToolParserOutput, +) -> Result<()> { + for event in output.events { + match event { + UnifiedParserEvent::Text(text) => tool_output.push_text(text), + UnifiedParserEvent::ToolCall(call) => tool_output.push_call(call), + UnifiedParserEvent::Reasoning(_) => { + return Err(ToolParserError::ParsingFailed { + message: "unified parser emitted reasoning in tool-parser adapter".to_string(), + }); + } + } + } + Ok(()) +} + +impl ToolParser for UnifiedToolParserAdapter { + fn create(tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + let inner = T::create(tools, Arc::new(BenchTokenizer)).map_err(map_unified_error)?; + Ok(Box::new(Self { + inner, + _marker: std::marker::PhantomData, + })) + } + + fn preserve_special_tokens(&self) -> bool { + self.inner.preserve_special_tokens() + } + + fn structural_tag_model(&self) -> Option { + self.inner.structural_tag_model() + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.inner.tool_call_id(tool_index) + } + + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + let mut unified_output = UnifiedParserOutput::default(); + let result = self.inner.parse_into(chunk, &mut unified_output).map_err(map_unified_error); + append_unified_output(unified_output, output)?; + result + } + + fn finish(&mut self) -> Result { + let unified_output = self.inner.finish().map_err(map_unified_error)?; + let mut output = ToolParserOutput::default(); + append_unified_output(unified_output, &mut output)?; + Ok(output) + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} diff --git a/rust/src/parser/benches/utils/mod.rs b/rust/src/parser/benches/utils/mod.rs index 914766a79aa..229f40a3681 100644 --- a/rust/src/parser/benches/utils/mod.rs +++ b/rust/src/parser/benches/utils/mod.rs @@ -1,4 +1,9 @@ +// This module is shared by multiple benchmark targets. +// There could be false positives for unused code or imports, and fixing them would lead to some other benchmarks failing to compile. #![allow(dead_code)] +#![allow(unused_imports)] + +mod adapter; use futures::FutureExt as _; use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool}; @@ -6,6 +11,8 @@ use tool_parser::traits::ToolParser as ExternalToolParser; use vllm_parser::tool::test_utils::collect_stream; use vllm_parser::tool::{Tool, ToolParser}; +pub(super) use adapter::UnifiedToolParserAdapter; + pub(super) fn openai_tools(tools: &[Tool]) -> Vec { tools .iter() diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs index 5ba2cf60edd..0b5c2b6d782 100644 --- a/rust/src/parser/src/lib.rs +++ b/rust/src/parser/src/lib.rs @@ -3,3 +3,4 @@ pub mod reasoning; pub mod tool; pub mod unified; +pub(crate) mod utils; diff --git a/rust/src/parser/src/reasoning/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs index 69b4db5f183..256e95fdde3 100644 --- a/rust/src/parser/src/reasoning/delimited.rs +++ b/rust/src/parser/src/reasoning/delimited.rs @@ -144,20 +144,20 @@ impl DelimitedReasoningParser { } /// Determine the reasoning state implied by the last prompt boundary, if any. -fn last_reasoning_boundary( +pub(crate) fn last_reasoning_boundary( prompt_token_ids: &[u32], start_token_id: u32, end_token_id: u32, tokenizer: &dyn Tokenizer, ) -> Option { - for token_id in prompt_token_ids.iter().rev() { - if *token_id == start_token_id { + for token_id in prompt_token_ids.iter().rev().copied() { + if token_id == start_token_id { return Some(true); } - if *token_id == end_token_id { + if token_id == end_token_id { return Some(false); } - if tokenizer.is_special_id(*token_id) { + if tokenizer.is_special_id(token_id) { return None; } } diff --git a/rust/src/parser/src/reasoning/gemma4.rs b/rust/src/parser/src/reasoning/gemma4.rs deleted file mode 100644 index ac5a6a17165..00000000000 --- a/rust/src/parser/src/reasoning/gemma4.rs +++ /dev/null @@ -1,273 +0,0 @@ -use vllm_tokenizer::DynTokenizer; - -use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; - -const THOUGHT_PREFIX: &str = "thought\n"; - -/// Reasoning parser for Google Gemma4 thinking models. -/// -/// Gemma4 emits reasoning inside `<|channel> ... ` spans and adds a -/// structural `thought\n` label at the beginning of the reasoning channel. -/// This parser keeps the delimiter handling in the shared delimited parser and -/// only layers on Gemma4-specific request adjustment plus prefix stripping. -/// -/// Original Python implementation: -/// -pub struct Gemma4ReasoningParser { - inner: DelimitedReasoningParser, - reasoning_text: String, - prefix_stripped: bool, -} - -impl Gemma4ReasoningParser { - /// Create a Gemma4 parser. - pub fn new(tokenizer: DynTokenizer) -> Result { - Ok(Self { - inner: DelimitedReasoningParser::new(tokenizer, "<|channel>", "", false)?, - reasoning_text: String::new(), - prefix_stripped: false, - }) - } - - /// Apply Gemma4's `thought\n` stripping rule to one reasoning delta. - /// - /// Early reasoning text is buffered until we can decide whether it begins - /// with the structural channel label. - fn strip_thought_prefix(&mut self, reasoning: &str) -> Option { - if self.prefix_stripped { - return Some(reasoning.to_string()); - } - - self.reasoning_text.push_str(reasoning); - - if self.reasoning_text.starts_with(THOUGHT_PREFIX) { - let prefix_len = THOUGHT_PREFIX.len(); - let previous_len = self.reasoning_text.len() - reasoning.len(); - if previous_len >= prefix_len { - self.reasoning_text.clear(); - self.prefix_stripped = true; - return Some(reasoning.to_string()); - } - - let prefix_chars_in_delta = prefix_len - previous_len; - let stripped = &reasoning[prefix_chars_in_delta.min(reasoning.len())..]; - if stripped.is_empty() { - if self.reasoning_text.len() >= prefix_len { - self.reasoning_text.clear(); - self.prefix_stripped = true; - } - return None; - } - - self.reasoning_text.clear(); - self.prefix_stripped = true; - return Some(stripped.to_string()); - } - - if THOUGHT_PREFIX.starts_with(&self.reasoning_text) { - return None; - } - - self.prefix_stripped = true; - Some(std::mem::take(&mut self.reasoning_text)) - } - - /// Apply Gemma4-specific reasoning post-processing to one parsed delta. - fn post_process(&mut self, mut result: ReasoningDelta) -> ReasoningDelta { - if let Some(reasoning) = result.reasoning.take() { - result.reasoning = - self.strip_thought_prefix(&reasoning).filter(|text| !text.is_empty()); - } - result - } -} - -impl ReasoningParser for Gemma4ReasoningParser { - fn create(tokenizer: DynTokenizer) -> Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self::new(tokenizer)?)) - } - - fn preserve_special_tokens(&self) -> bool { - true - } - - fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { - self.inner.initialize(prompt_token_ids); - self.reasoning_text.clear(); - self.prefix_stripped = false; - Ok(()) - } - - fn push(&mut self, delta: &str) -> Result { - let result = self.inner.push(delta); - Ok(self.post_process(result)) - } - - fn finish(&mut self) -> Result { - let result = self.inner.finish(); - Ok(self.post_process(result)) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use vllm_tokenizer::Tokenizer; - - use super::Gemma4ReasoningParser; - use crate::reasoning::ReasoningParser; - - struct FakeTokenizer; - - impl Tokenizer for FakeTokenizer { - fn encode( - &self, - text: &str, - _add_special_tokens: bool, - ) -> vllm_tokenizer::Result> { - Ok(text.chars().map(u32::from).collect()) - } - - fn decode( - &self, - token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_tokenizer::Result { - Ok(token_ids - .iter() - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) - .collect()) - } - - fn token_to_id(&self, token: &str) -> Option { - match token { - "<|channel>" => Some(1000), - "" => Some(1001), - _ => None, - } - } - } - - fn run_streaming(output: &[&str]) -> (Option, Option) { - let tokenizer = Arc::new(FakeTokenizer); - let mut parser = Gemma4ReasoningParser::new(tokenizer).unwrap(); - let mut reasoning = String::new(); - let mut content = String::new(); - - for delta in output { - let result = parser.push(delta).unwrap(); - if let Some(next) = result.reasoning { - reasoning.push_str(&next); - } - if let Some(next) = result.content { - content.push_str(&next); - } - } - - let final_delta = parser.finish().unwrap(); - if let Some(next) = final_delta.reasoning { - reasoning.push_str(&next); - } - if let Some(next) = final_delta.content { - content.push_str(&next); - } - - ( - (!reasoning.is_empty()).then_some(reasoning), - (!content.is_empty()).then_some(content), - ) - } - - #[test] - fn gemma4_reasoning_streaming_handles_channel_delimited_outputs() { - let cases = [ - ( - "no_reasoning", - vec!["This is content"], - None, - Some("This is content"), - ), - ( - "reasoning_and_content", - vec!["<|channel>This is a reasoning sectionThis is the rest"], - Some("This is a reasoning section"), - Some("This is the rest"), - ), - ( - "complete_reasoning", - vec!["<|channel>This is a reasoning section"], - Some("This is a reasoning section"), - None, - ), - ( - "multiple_lines", - vec!["<|channel>This\nThatThis is the rest\nThat"], - Some("This\nThat"), - Some("This is the rest\nThat"), - ), - ( - "no_end", - vec!["<|channel>This is a reasoning section"], - Some("This is a reasoning section"), - None, - ), - ("empty", vec![""], None, None), - ( - "newline_around_reasoning", - vec!["Before\n<|channel>This is a reasoning section\nThis is the rest"], - Some("This is a reasoning section"), - Some("Before\n\nThis is the rest"), - ), - ( - "thought_prefix", - vec!["<|channel>thought\nActual reasoning hereFinal answer"], - Some("Actual reasoning here"), - Some("Final answer"), - ), - ( - "thought_prefix_only", - vec!["<|channel>thought\n"], - None, - None, - ), - ( - "thought_prefix_multiline", - vec!["<|channel>thought\nLine1\nLine2Answer"], - Some("Line1\nLine2"), - Some("Answer"), - ), - ( - "thought_prefix_diverge", - vec!["<|channel>thousand reasonsDone"], - Some("thousand reasons"), - Some("Done"), - ), - ]; - - for (name, output, expected_reasoning, expected_content) in cases { - let (reasoning, content) = run_streaming(&output); - assert_eq!(reasoning.as_deref(), expected_reasoning, "{name}"); - assert_eq!(content.as_deref(), expected_content, "{name}"); - } - } - - #[test] - fn gemma4_strips_thought_prefix_even_when_split_across_deltas() { - let (reasoning, content) = - run_streaming(&["<|channel>thou", "ght", "\nabc", "done"]); - assert_eq!(reasoning.as_deref(), Some("abc")); - assert_eq!(content.as_deref(), Some("done")); - } - - #[test] - fn gemma4_preserves_special_tokens() { - let tokenizer = Arc::new(FakeTokenizer); - let parser = Gemma4ReasoningParser::new(tokenizer).unwrap(); - - assert!(parser.preserve_special_tokens()); - } -} diff --git a/rust/src/parser/src/reasoning/mod.rs b/rust/src/parser/src/reasoning/mod.rs index 1f71e14cef7..fcb0f96792a 100644 --- a/rust/src/parser/src/reasoning/mod.rs +++ b/rust/src/parser/src/reasoning/mod.rs @@ -17,7 +17,6 @@ mod cohere_cmd; mod deepseek_r1; mod delimited; -mod gemma4; mod kimi; mod minimax_m3; mod qwen3; @@ -29,8 +28,7 @@ use vllm_tokenizer::DynTokenizer; pub use self::cohere_cmd::CohereCmdReasoningParser; pub use self::deepseek_r1::DeepSeekR1ReasoningParser; -pub(crate) use self::delimited::DelimitedReasoningParser; -pub use self::gemma4::Gemma4ReasoningParser; +pub(crate) use self::delimited::{DelimitedReasoningParser, last_reasoning_boundary}; pub use self::kimi::KimiReasoningParser; pub use self::minimax_m3::MiniMaxM3ReasoningParser; pub use self::qwen3::Qwen3ReasoningParser; @@ -129,6 +127,10 @@ pub trait ReasoningParser: Send { pub enum ReasoningError { #[error("tokenizer is missing reasoning delimiter token `{token}`")] MissingToken { token: String }, + #[error( + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + )] + DummyUnifiedParser { name: String }, } #[cfg(test)] diff --git a/rust/src/parser/src/tool/error.rs b/rust/src/parser/src/tool/error.rs index 6a64c257d8c..4b2b4efcb45 100644 --- a/rust/src/parser/src/tool/error.rs +++ b/rust/src/parser/src/tool/error.rs @@ -10,4 +10,8 @@ pub type Result = std::result::Result; pub enum ToolParserError { #[error("tool parser parsing failed: {message}")] ParsingFailed { message: String }, + #[error( + "`{name}` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + )] + DummyUnifiedParser { name: String }, } diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index a27e202e660..dd4630b1c6b 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -4,7 +4,6 @@ pub(crate) mod error; mod deepseek_dsml; pub(crate) mod deepseek_json; -mod gemma4; mod glm_xml; mod hy_v3; mod json; @@ -15,14 +14,13 @@ mod parameters; mod qwen_coder; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; -pub(crate) mod utils; +use crate::utils; use std::collections::{BTreeMap, btree_map}; pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser}; pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser}; pub use error::{Result, ToolParserError}; -pub use gemma4::Gemma4ToolParser; pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; pub use hy_v3::HyV3ToolParser; pub use json::{ diff --git a/rust/src/parser/src/tool/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs similarity index 64% rename from rust/src/parser/src/tool/gemma4.rs rename to rust/src/parser/src/unified/gemma4.rs index e5a95485ce8..3276088913d 100644 --- a/rust/src/parser/src/tool/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -6,10 +6,17 @@ use winnow::prelude::*; use winnow::stream::{Partial, Stream}; use winnow::token::{literal, take_till, take_until}; -use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len}; -use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::tool::Tool; +use vllm_tokenizer::DynTokenizer; +use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; +use crate::reasoning::last_reasoning_boundary; +use crate::tool::{Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len_mul}; + +const REASONING_START: &str = "<|channel>thought\n"; +const CHANNEL_START: &str = "<|channel>"; +const CHANNEL_END: &str = ""; const TOOL_CALL_START: &str = "<|tool_call>"; const TOOL_CALL_END: &str = ""; const STRING_DELIM: &str = "<|\"|>"; @@ -20,6 +27,9 @@ type Gemma4Input<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq)] enum Gemma4Event { Text { len: usize }, + Reasoning { len: usize }, + ReasoningStart, + ReasoningEnd, ToolCallStart, ToolCallHeader { name: String }, ToolCall { args: Map }, @@ -35,6 +45,7 @@ struct Gemma4ArgsScanState { enum Gemma4Mode { #[default] Text, + Reasoning, Header, ToolCall { name: String, @@ -42,36 +53,62 @@ enum Gemma4Mode { }, } -/// Tool parser for Google Gemma4 models. +/// Unified parser for Google Gemma4 models. /// /// Original Python implementation: -/// +/// /// -/// Handles the Gemma4 function call format: +/// Handles Gemma4 reasoning and function-call formats: +/// +/// `<|channel>thought\nreasoning` /// /// `<|tool_call>call:func_name{key:<|"|>value<|"|>}` /// /// Arguments are emitted only after a full Gemma4 tool call is parsed. -pub struct Gemma4ToolParser { +pub struct Gemma4UnifiedParser { buffer: String, mode: Gemma4Mode, emitted_tool_count: usize, + tokenizer: DynTokenizer, + channel_start_token_id: u32, + channel_end_token_id: u32, } -impl Gemma4ToolParser { - fn new(_tools: &[Tool]) -> Self { - Self { +impl Gemma4UnifiedParser { + /// Create a Gemma4 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let channel_start_token_id = tokenizer.token_to_id(CHANNEL_START).ok_or_else(|| { + UnifiedParserError::MissingToken { + token: CHANNEL_START.to_string(), + } + })?; + let channel_end_token_id = + tokenizer + .token_to_id(CHANNEL_END) + .ok_or_else(|| UnifiedParserError::MissingToken { + token: CHANNEL_END.to_string(), + })?; + + Ok(Self { buffer: String::new(), mode: Gemma4Mode::default(), emitted_tool_count: 0, - } + channel_start_token_id, + channel_end_token_id, + tokenizer, + }) } - fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> { + fn apply_event(&mut self, event: Gemma4Event, output: &mut UnifiedParserOutput) -> Result<()> { match event { Gemma4Event::Text { len: consumed_len } => { - output.push_text(&self.buffer[..consumed_len]); + output.push_text(self.buffer[..consumed_len].to_string()); } + Gemma4Event::Reasoning { len: consumed_len } => { + output.push_reasoning(self.buffer[..consumed_len].to_string()); + } + Gemma4Event::ReasoningStart => self.mode = Gemma4Mode::Reasoning, + Gemma4Event::ReasoningEnd => self.mode = Gemma4Mode::Text, Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header, Gemma4Event::ToolCallHeader { name } => { self.mode = Gemma4Mode::ToolCall { @@ -100,9 +137,24 @@ impl Gemma4ToolParser { Ok(()) } + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = match last_reasoning_boundary( + prompt_token_ids, + self.channel_start_token_id, + self.channel_end_token_id, + self.tokenizer.as_ref(), + ) { + Some(true) => Gemma4Mode::Reasoning, + Some(false) | None => Gemma4Mode::Text, + }; + } + fn reset(&mut self) -> String { let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) { Gemma4Mode::Text => std::mem::take(&mut self.buffer), + Gemma4Mode::Reasoning => { + format!("{}{}", REASONING_START, std::mem::take(&mut self.buffer)) + } Gemma4Mode::Header => { format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer)) } @@ -122,19 +174,26 @@ impl Gemma4ToolParser { } } -impl ToolParser for Gemma4ToolParser { - fn create(tools: &[Tool]) -> Result> +impl UnifiedParser for Gemma4UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> where Self: Sized + 'static, { - Ok(Box::new(Self::new(tools))) + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.emitted_tool_count = 0; + self.initialize_mode(prompt_token_ids); + Ok(()) } fn preserve_special_tokens(&self) -> bool { true } - fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { self.buffer.push_str(chunk); while let Some((event, consumed_len)) = { @@ -149,11 +208,12 @@ impl ToolParser for Gemma4ToolParser { Ok(()) } - fn finish(&mut self) -> Result { - let mut output = ToolParserOutput::default(); + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); match &self.mode { - Gemma4Mode::Text => output.push_text(&self.buffer), + Gemma4Mode::Text => output.push_text(std::mem::take(&mut self.buffer)), + Gemma4Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => { return Err(parsing_failed!("incomplete Gemma4 tool call")); } @@ -164,7 +224,7 @@ impl ToolParser for Gemma4ToolParser { } fn reset(&mut self) -> String { - Gemma4ToolParser::reset(self) + Gemma4UnifiedParser::reset(self) } } @@ -175,6 +235,7 @@ fn parse_next_gemma4_event( ) -> ModalResult { match mode { Gemma4Mode::Text => parse_text_event(input), + Gemma4Mode::Reasoning => parse_reasoning_event(input), Gemma4Mode::Header => tool_call_header_event(input), Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan), } @@ -182,7 +243,32 @@ fn parse_next_gemma4_event( /// Parse a Gemma4 text-mode event. fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult { - alt((tool_call_start_event, safe_text_event)).parse_next(input) + alt(( + reasoning_start_event, + tool_call_start_event, + safe_text_event, + )) + .parse_next(input) +} + +/// Parse a Gemma4 reasoning-mode event. +fn parse_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult { + alt(( + reasoning_end_event, + tool_call_start_event, + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse a Gemma4 reasoning start marker. +fn reasoning_start_event(input: &mut Gemma4Input<'_>) -> ModalResult { + literal(REASONING_START).value(Gemma4Event::ReasoningStart).parse_next(input) +} + +/// Parse a Gemma4 reasoning end marker. +fn reasoning_end_event(input: &mut Gemma4Input<'_>) -> ModalResult { + literal(CHANNEL_END).value(Gemma4Event::ReasoningEnd).parse_next(input) } /// Parse a Gemma4 tool-call start marker. @@ -226,7 +312,14 @@ fn gemma4_tool_name(input: &mut Gemma4Input<'_>) -> ModalResult { /// Parse a safe text run before the next Gemma4 marker. fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult { - safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len }) + safe_text_len_mul(input, &[REASONING_START, TOOL_CALL_START]) + .map(|len| Gemma4Event::Text { len }) +} + +/// Parse a safe reasoning run before the next Gemma4 marker. +fn safe_reasoning_event(input: &mut Gemma4Input<'_>) -> ModalResult { + safe_text_len_mul(input, &[CHANNEL_END, TOOL_CALL_START]) + .map(|len| Gemma4Event::Reasoning { len }) } /// Parse raw Gemma4 arguments through the first end marker outside a Gemma string. @@ -418,17 +511,145 @@ fn parse_gemma4_scalar(value: &str) -> Value { #[cfg(test)] mod tests { + use std::sync::Arc; + use serde_json::{Value, json}; use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer; use winnow::combinator::{eof, terminated}; use winnow::error::ErrMode; use winnow::prelude::*; use super::{ - Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content, - parse_gemma4_args, + CHANNEL_END, CHANNEL_START, Gemma4UnifiedParser, ToolCallDelta, UnifiedParser, + UnifiedParserError, UnifiedParserOutput, gemma4_array_content, parse_gemma4_args, }; - use crate::tool::{Tool, ToolParserTestExt as _}; + use crate::tool::Tool; + use crate::unified::{UnifiedParserEvent, parsing_failed}; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + CHANNEL_START => Some(100), + CHANNEL_END => Some(101), + _ => None, + } + } + + fn is_special_id(&self, token_id: u32) -> bool { + matches!(token_id, 100..=105) + } + } + + struct MissingTokenTokenizer; + + impl Tokenizer for MissingTokenTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, _token: &str) -> Option { + None + } + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for Gemma4UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec<&ToolCallDelta>; + fn coalesce(self) -> Self; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + UnifiedParserEvent::Reasoning(_) | UnifiedParserEvent::ToolCall(_) => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None, + }) + .collect() + } + + fn calls(&self) -> Vec<&ToolCallDelta> { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(_) | UnifiedParserEvent::Reasoning(_) => None, + UnifiedParserEvent::ToolCall(call) => Some(call), + }) + .collect() + } + + fn coalesce(self) -> Self { + self + } + } fn parse_gemma4_array(array: &str) -> super::Result> { let mut input = array; @@ -494,9 +715,26 @@ mod tests { ] } - fn collect_stream(chunks: &[&str]) -> ToolParserOutput { - let mut parser = Gemma4ToolParser::new(&test_tools()); - let mut output = ToolParserOutput::default(); + fn test_parser() -> Gemma4UnifiedParser { + Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap() + } + + #[test] + fn gemma4_create_requires_channel_start_token() { + let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == CHANNEL_START + )); + } + + fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); for chunk in chunks { output.append(parser.parse_chunk(chunk).unwrap()); } @@ -504,7 +742,7 @@ mod tests { output.coalesce() } - fn first_call(output: &ToolParserOutput) -> ToolCallDelta { + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { (*output.calls().first().expect("expected one tool call")).clone() } @@ -542,7 +780,7 @@ mod tests { #[test] fn gemma4_parse_complete_extracts_single_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let output = parser .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}") .unwrap(); @@ -558,7 +796,7 @@ mod tests { #[test] fn gemma4_parse_complete_rejects_incomplete_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let error = parser .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London") .unwrap_err(); @@ -607,8 +845,8 @@ mod tests { #[test] fn gemma4_streaming_waits_for_complete_tool_call() { - let mut parser = Gemma4ToolParser::new(&test_tools()); - let mut output = ToolParserOutput::default(); + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); for chunk in [ "<|tool_call>", @@ -773,7 +1011,7 @@ mod tests { #[test] fn gemma4_finish_flushes_partial_start_marker_as_text() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let mut output = parser.parse_chunk("<").unwrap(); output.append(parser.finish().unwrap()); @@ -781,9 +1019,104 @@ mod tests { assert!(output.calls().is_empty()); } + #[test] + fn gemma4_streaming_emits_reasoning_then_text() { + let output = collect_stream(&["<|channel>thought\nreasonanswer"]); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + assert!(output.calls().is_empty()); + } + + #[test] + fn gemma4_streaming_holds_split_reasoning_start() { + let mut parser = test_parser(); + + let first = parser.parse_chunk("<|channel>").unwrap(); + assert!(first.events.is_empty()); + + let mut output = parser.parse_chunk("thought\nrea").unwrap(); + output.append(parser.parse_chunk("sonanswer").unwrap()); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 3001]).unwrap(); + + let output = parser.parse_complete("reasonanswer").unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_turn_prompt_starts_in_text() { + let mut parser = test_parser(); + parser.initialize(&[104, 3000, 3001]).unwrap(); + + let output = parser.parse_complete("<|channel>thought\nreasonanswer").unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_special_token_caps_boundary_scan() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 104, 3001]).unwrap(); + + let output = parser.parse_complete("answer").unwrap(); + + assert!(output.reasoning_text().is_empty()); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_initialize_closed_channel_prompt_starts_in_text() { + let mut parser = test_parser(); + parser.initialize(&[100, 3000, 3001, 101]).unwrap(); + + let output = parser.parse_complete("answer").unwrap(); + + assert!(output.reasoning_text().is_empty()); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn gemma4_reasoning_tool_call_implicitly_ends_reasoning() { + let output = collect_stream(&[ + "<|channel>thought\nNeed weather.", + "<|tool_call>", + "call:get_weather{location:<|\"|>Paris<|\"|>}", + "", + ]); + + assert_eq!(output.reasoning_text(), "Need weather."); + assert!(output.normal_text().is_empty()); + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "location": "Paris" }) + ); + } + + #[test] + fn gemma4_bare_channel_start_is_plain_text() { + let output = collect_stream(&["<|channel>plain"]); + + assert_eq!(output.normal_text(), "<|channel>plain"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + #[test] fn gemma4_finish_rejects_complete_args_without_end_marker() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); for chunk in ["<|tool_call>", "call:get_status{}"] { parser.parse_chunk(chunk).unwrap(); } @@ -795,7 +1128,7 @@ mod tests { #[test] fn gemma4_reset_preserves_internally_buffered_arguments() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); for chunk in [ "<|tool_call>", "call:write_file{", @@ -815,7 +1148,7 @@ mod tests { #[test] fn gemma4_reset_preserves_completed_arguments_after_parse_error() { - let mut parser = Gemma4ToolParser::new(&test_tools()); + let mut parser = test_parser(); let input = "<|tool_call>call:set{broken}"; let _error = parser.parse_chunk(input).unwrap_err(); diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 49955b4d818..6fe7d29b879 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -1,11 +1,14 @@ //! Unified parser interface for reasoning and tool-call deltas. mod combined; +mod gemma4; use thiserror::Error; +use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; pub use combined::CombinedParser; +pub use gemma4::Gemma4UnifiedParser; use crate::reasoning::ReasoningError; use crate::tool::{ @@ -189,10 +192,15 @@ pub trait UnifiedParser: Send { } /// Errors produced while creating or running unified parsers. -#[derive(Debug, Error)] +#[derive(Debug, Error, Macro)] +#[thiserror_ext(macro(path = "crate::unified", mangle))] pub enum UnifiedParserError { #[error("combined parser is constructed from split parser instances")] CombinedParserConstructor, + #[error("tokenizer is missing unified parser token `{token}`")] + MissingToken { token: String }, + #[error("unified parser parsing failed: {message}")] + ParsingFailed { message: String }, #[error(transparent)] Reasoning(#[from] ReasoningError), #[error(transparent)] diff --git a/rust/src/parser/src/tool/utils.rs b/rust/src/parser/src/utils.rs similarity index 83% rename from rust/src/parser/src/tool/utils.rs rename to rust/src/parser/src/utils.rs index b545c5881c1..70255215393 100644 --- a/rust/src/parser/src/tool/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -1,10 +1,10 @@ -//! Shared helpers for tool parsers. +//! Shared helpers for streaming parsers. use winnow::Parser; use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue}; -use winnow::stream::{Offset, Partial, Stream}; +use winnow::stream::{FindSlice, Offset, Partial, Stream}; -use super::Result; +use crate::tool::{Result, ToolParserError}; /// Return the byte length of the longest proper prefix of `token` that is also /// a suffix of `buffer`. @@ -15,7 +15,7 @@ use super::Result; /// The returned length is always a valid UTF-8 boundary in `token`, so callers /// can safely slice `&token[..len]` even when markers contain non-ASCII /// characters such as DeepSeek's DSML delimiters. -pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize { +pub fn partial_prefix_len(buffer: &str, token: &str) -> usize { let Some(first_byte) = token.as_bytes().first().copied() else { return 0; }; @@ -44,9 +44,10 @@ pub(super) fn partial_prefix_len(buffer: &str, token: &str) -> usize { } /// Parse a safe text run before the next marker. +/// This is the single-marker variant of [`safe_text_len_mul`]. /// /// Returns the text length in bytes, and advances the input. -pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult { +pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -67,15 +68,53 @@ pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalRes Ok(emit_len) } +/// Parse a safe text run before the earliest next marker. +/// This is the multi-marker variant of [`safe_text_len`]. +/// +/// Returns the text length in bytes, and advances the input. +pub fn safe_text_len_mul(input: &mut Partial<&str>, markers: &[&str]) -> ModalResult { + let text = **input; + if text.is_empty() { + return incomplete(); + } + + if let Some(start_idx) = find_slice_mul(text, markers) { + input.next_slice(start_idx); + return Ok(start_idx); + } + + let keep_len = markers.iter().map(|marker| partial_prefix_len(text, marker)).max().unwrap_or(0); + let emit_len = text.len().saturating_sub(keep_len); + if emit_len == 0 { + return incomplete(); + } + + input.next_slice(emit_len); + Ok(emit_len) +} + +#[inline(always)] +fn find_slice_mul(text: &str, markers: &[&str]) -> Option { + let range = match markers { + // Use the fast specialized `winnow::stream::FindSlice` impl for 1-3 markers. + [first] => text.find_slice(*first), + [first, second] => text.find_slice((*first, *second)), + [first, second, third] => text.find_slice((*first, *second, *third)), + // Fall back to a linear scan for 4+ markers. + _ => return markers.iter().filter_map(|marker| text.find(marker)).min(), + }; + range.map(|range| range.start) +} + /// Streaming scan state for a buffered marker search [`take_until_marker`], /// so that we don't have to rescan the whole buffered prefix when resuming. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct MarkerScanState { +pub struct MarkerScanState { scan_start: usize, } impl MarkerScanState { - pub(super) fn reset(&mut self) { + pub fn reset(&mut self) { self.scan_start = 0; } } @@ -92,7 +131,7 @@ impl MarkerScanState { /// chunks while waiting for a closing marker. Plain `take_until` is still a /// better fit for one-shot parsers over a complete body, and for `1..` cases /// where an empty slice before the marker should be rejected. -pub(super) fn take_until_marker<'i, 'a>( +pub fn take_until_marker<'i, 'a>( marker: &'a str, state: &'a mut MarkerScanState, ) -> impl Parser, &'i str, ErrMode> + 'a { @@ -137,7 +176,7 @@ fn floor_char_boundary(text: &str, index: usize) -> usize { /// Streaming lexical state for a top-level JSON object. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct JsonObjectScanState { +pub struct JsonObjectScanState { object_depth: usize, array_depth: usize, in_string: bool, @@ -155,7 +194,7 @@ enum JsonObjectScanPhase { impl JsonObjectScanState { /// Returns whether the top-level JSON object has closed. - pub(super) const fn complete(&self) -> bool { + pub const fn complete(&self) -> bool { matches!(self.phase, JsonObjectScanPhase::Complete) } } @@ -165,7 +204,7 @@ impl JsonObjectScanState { /// The returned length is safe to emit as raw argument text. This scans only /// lexical boundaries from `{` through the matching `}`, preserving /// malformed-but-balanced JSON without deserializing or normalizing it. -pub(super) fn take_json_object( +pub fn take_json_object( input: &mut Partial<&str>, state: &mut JsonObjectScanState, ) -> ModalResult { @@ -252,7 +291,7 @@ pub(super) fn take_json_object( } /// Parse a JSON string literal. -pub(super) fn json_str(input: &mut Partial<&str>) -> ModalResult { +pub fn json_str(input: &mut Partial<&str>) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -311,7 +350,7 @@ fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode( +pub fn parse_buffered_event( buffer: &str, parse: impl FnOnce(&mut Partial<&str>) -> ModalResult, ) -> Result> { @@ -322,7 +361,9 @@ pub(super) fn parse_buffered_event( Err(ErrMode::Incomplete(_)) => return Ok(None), Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => { // TODO: enrich context for error reporting - return Err(parsing_failed!("{}", e)); + return Err(ToolParserError::ParsingFailed { + message: e.to_string(), + }); } }; let consumed_len = input.offset_from(&checkpoint); @@ -334,7 +375,7 @@ pub(super) fn parse_buffered_event( } /// Returns an error indicating that we need more data to continue parsing. -pub(super) fn incomplete() -> ModalResult { +pub fn incomplete() -> ModalResult { Err(ErrMode::Incomplete(Needed::Unknown)) } @@ -348,7 +389,7 @@ mod tests { use super::{ JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len, - take_json_object, take_until_marker, + safe_text_len_mul, take_json_object, take_until_marker, }; #[test] @@ -406,6 +447,52 @@ mod tests { assert!(matches!(error, ErrMode::Incomplete(_))); } + #[test] + fn safe_text_len_mul_stops_before_earliest_marker() { + let mut input = Partial::new("hello<|tool_call>"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", ""]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool_call>"); + } + + #[test] + fn safe_text_len_mul_holds_back_longest_partial_marker() { + let mut input = Partial::new("hello<|tool"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool"); + } + + #[test] + fn safe_text_len_mul_skips_false_same_prefix_candidate() { + let mut input = Partial::new("hello<|tool_call>"); + let checkpoint = input.checkpoint(); + + let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap(); + + assert_eq!(len, "hello".len()); + assert_eq!(input.offset_from(&checkpoint), "hello".len()); + assert_eq!(*input, "<|tool_call>"); + } + + #[test] + fn safe_text_len_mul_reports_incomplete_for_only_partial_marker() { + let mut input = Partial::new("<|channel>thought"); + + let error = + safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + } + #[test] fn take_until_marker_stops_before_marker() { let mut state = MarkerScanState::default();