diff --git a/pyproject.toml b/pyproject.toml index 3819ad7fc8e..a45db6a167c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,9 +129,9 @@ 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/parser/src/unified/gemma4.rs", + "rust/src/parser/*", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] + "rust/src/tokenizer/src/incremental.rs",] ignore-hidden = false [tool.typos.default] diff --git a/rust/src/parser/benches/gemma4.rs b/rust/src/parser/benches/gemma4.rs index fd29e77a9a2..b6820550574 100644 --- a/rust/src/parser/benches/gemma4.rs +++ b/rust/src/parser/benches/gemma4.rs @@ -3,7 +3,7 @@ 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::{Tool, ToolParser}; -use vllm_parser::unified::Gemma4UnifiedParser; +use vllm_parser::unified::{Gemma4ConfigDrivenParser, Gemma4UnifiedParser}; mod utils; use utils::{UnifiedToolParserAdapter, feed_parser}; @@ -68,14 +68,22 @@ fn long_tool_argument_fixture() -> String { ) } +type MakeParser = fn(&[Tool]) -> Box; + fn parser(tools: &[Tool]) -> Box { UnifiedToolParserAdapter::::create(tools) .expect("Gemma4 unified parser should initialize") } +fn config_driven_parser(tools: &[Tool]) -> Box { + UnifiedToolParserAdapter::::create(tools) + .expect("Gemma4 config-driven parser should initialize") +} + fn run_stream_group( c: &mut Criterion, name: &str, + make_parser: MakeParser, tools: &[Tool], text: &str, chunk_chars: usize, @@ -91,7 +99,7 @@ fn run_stream_group( group.throughput(Throughput::Bytes(text.len() as u64)); group.bench_function("reuse_parser", |b| { - let mut parser = parser(tools); + let mut parser = make_parser(tools); b.iter(|| { let result = feed_parser(&mut *parser, black_box(&chunks)); debug_assert_eq!(result.0, expected_normal_text); @@ -102,7 +110,7 @@ fn run_stream_group( group.bench_function("create_parser", |b| { b.iter_batched( - || parser(tools), + || make_parser(tools), |mut parser| { let result = feed_parser(&mut *parser, black_box(&chunks)); debug_assert_eq!(result.0, expected_normal_text); @@ -122,35 +130,45 @@ fn bench_gemma4(c: &mut Criterion) { let long_tool_argument = long_tool_argument_fixture(); let long_normal_text = long_normal_text_fixture(); - run_stream_group( - c, - "gemma4/mixed_complex_tool_call", - &tools, - &mixed_text, - CHUNK_CHARS, - "I will inspect the data before answering.\n Finished.", - 2, - ); + let variants: [(&str, MakeParser); 2] = [ + ("gemma4", parser), + ("gemma4_config_driven", config_driven_parser), + ]; - run_stream_group( - c, - "gemma4/long_tool_argument", - &tools, - &long_tool_argument, - CHUNK_CHARS, - "I will write the file.\nDone.", - 1, - ); + for (prefix, make_parser) in variants { + run_stream_group( + c, + &format!("{prefix}/mixed_complex_tool_call"), + make_parser, + &tools, + &mixed_text, + CHUNK_CHARS, + "I will inspect the data before answering.\n Finished.", + 2, + ); - run_stream_group( - c, - "gemma4/long_normal_text", - &tools, - &long_normal_text, - CHUNK_CHARS, - &long_normal_text, - 0, - ); + run_stream_group( + c, + &format!("{prefix}/long_tool_argument"), + make_parser, + &tools, + &long_tool_argument, + CHUNK_CHARS, + "I will write the file.\nDone.", + 1, + ); + + run_stream_group( + c, + &format!("{prefix}/long_normal_text"), + make_parser, + &tools, + &long_normal_text, + CHUNK_CHARS, + &long_normal_text, + 0, + ); + } } criterion_group!(benches, bench_gemma4); diff --git a/rust/src/parser/src/unified/config_driven/gemma4.rs b/rust/src/parser/src/unified/config_driven/gemma4.rs new file mode 100644 index 00000000000..817296b8071 --- /dev/null +++ b/rust/src/parser/src/unified/config_driven/gemma4.rs @@ -0,0 +1,870 @@ +//! Gemma4 format for the config-driven unified parser. +//! +//! Original Python implementation: +//! +//! +//! Handles Gemma4 reasoning and function-call formats: +//! +//! `<|channel>thought\nreasoning` +//! +//! `<|tool_call>call:func_name{key:<|"|>value<|"|>}` + +use serde_json::{Map, Number, Value}; +use winnow::ascii::multispace0 as ws0; +use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated}; +use winnow::error::{ContextError, ErrMode, ModalResult}; +use winnow::prelude::*; +use winnow::stream::Stream; +use winnow::token::{literal, take_till, take_until}; + +use super::{ArgsEndScan, ConfigDrivenParser, Input, ParserFormat}; +use crate::utils::{incomplete, partial_prefix_len}; + +const CHANNEL_START: &str = "<|channel>"; +const CHANNEL_END: &str = ""; +const REASONING_START: &str = "<|channel>thought\n"; +const TOOL_CALL_START: &str = "<|tool_call>"; +const TOOL_CALL_END: &str = ""; +const STRING_DELIM: &str = "<|\"|>"; +const CALL_PREFIX: &str = "call:"; + +/// Zero-sized marker type providing the Google Gemma4 format. +pub struct Gemma4Format; + +impl ParserFormat for Gemma4Format { + const NAME: &'static str = "Gemma4"; + const REASONING_START: Option<&'static str> = Some(REASONING_START); + const REASONING_END: Option<&'static str> = Some(CHANNEL_END); + const CALL_START: &'static str = TOOL_CALL_START; + const CALL_END: &'static str = TOOL_CALL_END; + const REASONING_BOUNDARY_TOKENS: Option<(&'static str, &'static str)> = + Some((CHANNEL_START, CHANNEL_END)); + + type ArgsScan = Gemma4ArgsScan; + + fn tool_header(input: &mut Input<'_>) -> ModalResult { + let (name,) = seq!( + _: literal(CALL_PREFIX), + gemma4_tool_name, + _: literal("{"), + ) + .parse_next(input)?; + Ok(name) + } + + fn tool_args(body: &str) -> ModalResult> { + let Some(args_input) = body.strip_suffix('}') else { + return Err(ErrMode::Cut(ContextError::new())); + }; + parse_gemma4_args(args_input) + } +} + +/// Config-driven unified parser for Google Gemma4 models. +/// +/// Arguments are emitted only after a full Gemma4 tool call is parsed. +pub type Gemma4ConfigDrivenParser = ConfigDrivenParser; + +/// Resumable Gemma4 argument scanner: locates the call end marker outside +/// `<|"|>`-delimited strings, so marker-shaped literals inside string values +/// do not terminate the call early. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct Gemma4ArgsScan { + scanned_len: usize, + in_string: bool, +} + +impl ArgsEndScan for Gemma4ArgsScan { + fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<&'i str> { + let text = **input; + if self.scanned_len > text.len() { + return incomplete(); + } + + loop { + let rest = &text[self.scanned_len..]; + if self.in_string { + let Some(string_delim) = rest.find(STRING_DELIM) else { + self.scanned_len = safe_scan_len(text, self.scanned_len, &[STRING_DELIM]); + return incomplete(); + }; + + self.scanned_len += string_delim + STRING_DELIM.len(); + self.in_string = false; + continue; + } + + let next_string_delim = rest.find(STRING_DELIM); + let next_call_end = rest.find(end); + match (next_string_delim, next_call_end) { + (Some(string_delim), Some(call_end)) if call_end < string_delim => { + let body_end = self.scanned_len + call_end; + self.scanned_len = body_end + end.len(); + input.next_slice(self.scanned_len); + return Ok(&text[..body_end]); + } + (Some(string_delim), _) => { + self.scanned_len += string_delim + STRING_DELIM.len(); + self.in_string = true; + } + (None, Some(call_end)) => { + let body_end = self.scanned_len + call_end; + self.scanned_len = body_end + end.len(); + input.next_slice(self.scanned_len); + return Ok(&text[..body_end]); + } + (None, None) => { + self.scanned_len = safe_scan_len(text, self.scanned_len, &[STRING_DELIM, end]); + return incomplete(); + } + } + } + } +} + +/// Return the scan length while holding back a split marker prefix. +fn safe_scan_len(text: &str, start: usize, markers: &[&str]) -> usize { + let max_partial = markers + .iter() + .map(|marker| partial_prefix_len(&text[start..], marker)) + .max() + .unwrap_or(0); + text.len() - max_partial +} + +/// Parse a Gemma4 tool name. +fn gemma4_tool_name(input: &mut Input<'_>) -> ModalResult { + let name = take_until(1.., "{").parse_next(input)?.trim(); + if name.is_empty() { + return Err(ErrMode::Cut(ContextError::new())); + } + Ok(name.to_string()) +} + +/// Parse complete Gemma4 custom key-value arguments. +fn parse_gemma4_args(args: &str) -> ModalResult> { + let mut input = args; + terminated(gemma4_args, eof).parse_next(&mut input) +} + +/// Parse Gemma4's custom key-value argument object content. +fn gemma4_args(input: &mut &str) -> ModalResult> { + let pairs: Vec<(String, Value)> = delimited( + ws0, + terminated( + separated(0.., gemma4_pair, comma_separator), + opt(comma_separator), + ), + ws0, + ) + .parse_next(input)?; + Ok(pairs.into_iter().collect()) +} + +/// Parse a Gemma4 key-value pair. +fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> { + let (key, value) = seq!( + _: ws0, + gemma4_key, + _: ws0, + _: literal(":"), + _: ws0, + gemma4_value, + ) + .parse_next(input)?; + Ok((key, value)) +} + +/// Parse a Gemma4 bare key. +fn gemma4_key(input: &mut &str) -> ModalResult { + let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim(); + if key.is_empty() { + return Err(ErrMode::Cut(ContextError::new())); + } + Ok(key.to_string()) +} + +/// Parse a Gemma4 value. +fn gemma4_value(input: &mut &str) -> ModalResult { + alt(( + gemma4_string.map(|value: &str| Value::String(value.to_string())), + gemma4_object.map(Value::Object), + gemma4_array_value.map(Value::Array), + gemma4_bare_value, + )) + .parse_next(input) +} + +/// Parse a Gemma4 string delimited by `<|"|>`. +fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> { + delimited( + literal(STRING_DELIM), + take_until(0.., STRING_DELIM), + literal(STRING_DELIM), + ) + .parse_next(input) +} + +/// Parse a nested Gemma4 object. +fn gemma4_object(input: &mut &str) -> ModalResult> { + delimited(literal("{"), gemma4_args, literal("}")).parse_next(input) +} + +/// Parse a Gemma4 array value. +fn gemma4_array_value(input: &mut &str) -> ModalResult> { + delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input) +} + +/// Parse Gemma4 array content. +fn gemma4_array_content(input: &mut &str) -> ModalResult> { + delimited( + ws0, + terminated( + separated(0.., gemma4_value, comma_separator), + opt(comma_separator), + ), + ws0, + ) + .parse_next(input) +} + +/// Parse a Gemma4 bare scalar. +fn gemma4_bare_value(input: &mut &str) -> ModalResult { + take_till(1.., |char: char| matches!(char, ',' | '}' | ']')) + .map(parse_gemma4_scalar) + .parse_next(input) +} + +/// Parse a Gemma4 comma separator. +fn comma_separator(input: &mut &str) -> ModalResult<()> { + delimited(ws0, literal(","), ws0).void().parse_next(input) +} + +fn parse_gemma4_scalar(value: &str) -> Value { + let value = value.trim(); + if value.is_empty() { + return Value::String(String::new()); + } + if value == "true" { + return Value::Bool(true); + } + if value == "false" { + return Value::Bool(false); + } + if matches!(value, "null" | "none" | "nil" | "NULL" | "None" | "NIL") { + return Value::Null; + } + if value.contains('.') { + if let Ok(parsed) = value.parse::() + && let Some(number) = Number::from_f64(parsed) + { + return Value::Number(number); + } + } else if let Ok(parsed) = value.parse::() { + return Value::Number(Number::from(parsed)); + } + + Value::String(value.to_string()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use vllm_tokenizer::test_utils::TestTokenizer; + use winnow::combinator::{eof, terminated}; + use winnow::error::ErrMode; + use winnow::prelude::*; + + use super::{ + CHANNEL_END, CHANNEL_START, Gemma4ConfigDrivenParser, gemma4_array_content, + parse_gemma4_args, + }; + use crate::tool::{Tool, ToolCallDelta}; + use crate::unified::{ + Result, UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + parsing_failed, + }; + + const CHANNEL_START_ID: u32 = 256; + const CHANNEL_END_ID: u32 = 257; + const TURN_BOUNDARY_ID: u32 = 258; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(CHANNEL_START, CHANNEL_START_ID) + .with_special_token(CHANNEL_END, CHANNEL_END_ID) + .with_special_token("", TURN_BOUNDARY_ID) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> Result; + fn parse_complete(&mut self, text: &str) -> Result; + } + + impl UnifiedParserTestExt for Gemma4ConfigDrivenParser { + fn parse_chunk(&mut self, chunk: &str) -> Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> 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) -> Result> { + let mut input = array; + match terminated(gemma4_array_content, eof).parse_next(&mut input) { + Ok(value) => Ok(value), + Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")), + Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => { + Err(parsing_failed!("{}", error)) + } + } + } + + fn test_tools() -> Vec { + vec![ + Tool { + name: "get_weather".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "get_time".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "write_file".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "Edit".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "search".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "set".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "get_status".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + Tool { + name: "todowrite".to_string(), + description: None, + parameters: json!({ "type": "object" }), + strict: None, + }, + ] + } + + fn test_parser() -> Gemma4ConfigDrivenParser { + Gemma4ConfigDrivenParser::new(&test_tools(), Arc::new(tokenizer())).unwrap() + } + + #[test] + fn gemma4_create_requires_channel_start_token() { + let error = + match Gemma4ConfigDrivenParser::new(&test_tools(), Arc::new(TestTokenizer::new())) { + 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()); + } + output.append(parser.finish().unwrap()); + output.coalesce() + } + + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { + (*output.calls().first().expect("expected one tool call")).clone() + } + + #[test] + fn gemma4_parse_args_handles_scalars_and_nested_values() { + let parsed = parse_gemma4_args( + "name:<|\"|>test<|\"|>,count:42,active:true,score:114.514,nested:{inner:<|\"|>value<|\"|>},items:[<|\"|>a<|\"|>,<|\"|>b<|\"|>]", + ) + .unwrap(); + + assert_eq!( + Value::Object(parsed), + json!({ + "name": "test", + "count": 42, + "active": true, + "score": 114.514, + "nested": { "inner": "value" }, + "items": ["a", "b"], + }) + ); + } + + #[test] + fn gemma4_parse_args_handles_empty_arguments() { + let parsed = parse_gemma4_args("").unwrap(); + assert_eq!(Value::Object(parsed), json!({})); + } + + #[test] + fn gemma4_parse_array_handles_bare_values() { + let parsed = parse_gemma4_array("42,true,114.514").unwrap(); + assert_eq!(Value::Array(parsed), json!([42, true, 114.514])); + } + + #[test] + fn gemma4_parse_complete_extracts_single_tool_call() { + let mut parser = test_parser(); + let output = parser + .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}") + .unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "location": "London" }) + ); + } + + #[test] + fn gemma4_parse_complete_rejects_incomplete_tool_call() { + let mut parser = test_parser(); + let error = parser + .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London") + .unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Gemma4 tool call")); + } + + #[test] + fn gemma4_streaming_basic_single_tool_call() { + let output = collect_stream(&[ + "<|tool_call>", + "call:get_weather{", + "location:<|\"|>Paris", + ", France", + "<|\"|>}", + "", + ]); + + 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, France" }) + ); + } + + #[test] + fn gemma4_streaming_text_before_and_after_tool_call() { + let output = collect_stream(&[ + "Let me check ", + "the weather. ", + "<|tool_call>", + "call:get_weather{", + "location:<|\"|>London<|\"|>}", + "<", + "div>", + ]); + + assert_eq!(output.normal_text(), "Let me check the weather.
"); + assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "location": "London" }) + ); + } + + #[test] + fn gemma4_streaming_waits_for_complete_tool_call() { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + + for chunk in [ + "<|tool_call>", + "call:get_weather{", + "location:<|\"|>Paris<|\"|>}", + ] { + output.append(parser.parse_chunk(chunk).unwrap()); + assert!(output.calls().is_empty()); + } + + output.append(parser.parse_chunk("").unwrap()); + let output = output.coalesce(); + + 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_streaming_handles_boolean_split_across_chunks() { + let output = collect_stream(&[ + "<|tool_call>", + "call:search{input:{all:tru", + "e}}", + "", + ]); + + assert_eq!(first_call(&output).name.as_deref(), Some("search")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "input": { "all": true } }) + ); + } + + #[test] + fn gemma4_streaming_handles_false_split_across_chunks() { + let output = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", ""]); + + assert_eq!(first_call(&output).name.as_deref(), Some("set")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "flag": false }) + ); + } + + #[test] + fn gemma4_streaming_handles_number_split_across_chunks() { + let output = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", ""]); + + assert_eq!(first_call(&output).name.as_deref(), Some("set")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "count": 42 }) + ); + } + + #[test] + fn gemma4_streaming_handles_split_string_delimiter() { + let output = collect_stream(&[ + "<|tool_call>", + "call:todowrite{", + "content:<|\"|>Buy milk<|", + "\"|>}", + "", + ]); + + assert_eq!(first_call(&output).name.as_deref(), Some("todowrite")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": "Buy milk" }) + ); + assert!(!first_call(&output).arguments.contains("<|")); + } + + #[test] + fn gemma4_streaming_handles_split_tool_call_end_marker() { + let output = collect_stream(&[ + "<|tool_call>", + "call:get_weather{location:<|\"|>Paris<|\"|>}", + ]); + + 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_streaming_handles_end_marker_literal_inside_string() { + let output = collect_stream(&[ + "<|tool_call>", + "call:todowrite{", + "content:<|\"|>literal } inside", + "<|\"|>}", + "", + ]); + + assert_eq!(first_call(&output).name.as_deref(), Some("todowrite")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": "literal } inside" }) + ); + } + + #[test] + fn gemma4_streaming_handles_html_argument_without_duplication() { + let output = collect_stream(&[ + "<|tool_call>", + "call:write_file{", + "path:<|\"|>index.html<|\"|>,", + "content:<|\"|>\n<", + "html lang=\"zh-CN\">\n<", + "head>\n <", + "meta charset=\"UTF-8\">\n <", + "meta name=\"viewport\" content=\"width=device-width\">\n", + "<|\"|>}", + "", + ]); + + assert_eq!(first_call(&output).name.as_deref(), Some("write_file")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ + "path": "index.html", + "content": "\n\n\n \n \n", + }) + ); + } + + #[test] + fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() { + let output = collect_stream(&[ + "<|tool_call>", + "call:Edit{", + "file_path:<|\"|>src/env.py<|\"|>,", + "old_string:<|\"|>old_val<|\"|>,", + "new_string:<|\"|>new_val<|\"|>,", + "replace_all:", + "false}", + "", + ]); + + assert_eq!(first_call(&output).name.as_deref(), Some("Edit")); + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ + "file_path": "src/env.py", + "old_string": "old_val", + "new_string": "new_val", + "replace_all": false, + }) + ); + assert_eq!( + first_call(&output).arguments.matches("replace_all").count(), + 1 + ); + } + + #[test] + fn gemma4_finish_flushes_partial_start_marker_as_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk("<").unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "<"); + 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(&[CHANNEL_START_ID, 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(&[TURN_BOUNDARY_ID, 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(&[CHANNEL_START_ID, 3000, TURN_BOUNDARY_ID, 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(&[CHANNEL_START_ID, 3000, 3001, CHANNEL_END_ID]).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 = test_parser(); + for chunk in ["<|tool_call>", "call:get_status{}"] { + parser.parse_chunk(chunk).unwrap(); + } + + let error = parser.finish().unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Gemma4 tool call")); + } + + #[test] + fn gemma4_reset_preserves_internally_buffered_arguments() { + let mut parser = test_parser(); + for chunk in [ + "<|tool_call>", + "call:write_file{", + "content:<|\"|>hello ", + "world<|\"|>", + ] { + parser.parse_chunk(chunk).unwrap(); + } + + let raw = parser.reset(); + + assert_eq!( + raw, + "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>" + ); + } + + #[test] + fn gemma4_reset_preserves_completed_arguments_after_parse_error() { + let mut parser = test_parser(); + let input = "<|tool_call>call:set{broken}"; + + let _error = parser.parse_chunk(input).unwrap_err(); + let raw = parser.reset(); + + assert_eq!(raw, input); + } +} diff --git a/rust/src/parser/src/unified/config_driven/mod.rs b/rust/src/parser/src/unified/config_driven/mod.rs new file mode 100644 index 00000000000..9a2501842f8 --- /dev/null +++ b/rust/src/parser/src/unified/config_driven/mod.rs @@ -0,0 +1,474 @@ +//! Config-driven unified parser: a shared boundary state machine parameterized +//! by a per-model [`ParserFormat`] implemented on zero-sized marker types. +//! +//! The engine owns a fixed five-state boundary FSM (plus a terminal `Done` +//! state) and all streaming mechanics: buffering, partial-marker holdback, +//! event application, raw-text reconstruction for [`UnifiedParser::reset`], +//! and end-of-stream semantics. A format contributes only: +//! +//! - boundary marker literals (with `Option` encoding which constructs exist), +//! - a winnow grammar for the tool-call header and the complete argument body, +//! - a resumable scanner locating the end of the raw argument body. +//! +//! Formats never define states or transitions. A format whose behavior does +//! not fit the fixed graph plus the small policy surface below should be +//! written as a standalone parser instead of growing this engine. +//! +//! Because every [`ParserFormat`] item is an associated constant or a static +//! method, a model parser is a zero-sized marker type plus a type alias in its +//! own format module (see [`gemma4`]): +//! +//! ```ignore +//! pub struct Gemma4Format; +//! impl ParserFormat for Gemma4Format { /* consts + two grammar fns */ } +//! pub type Gemma4ConfigDrivenParser = ConfigDrivenParser; +//! ``` +//! +//! Monomorphization then compiles `ConfigDrivenParser` down to the +//! equivalent of a hand-written per-model parser: literals are inlined and +//! there is no dynamic dispatch on the format. + +mod gemma4; + +use std::marker::PhantomData; + +pub use gemma4::{Gemma4ConfigDrivenParser, Gemma4Format}; +use serde_json::{Map, Value}; +use vllm_tokenizer::DynTokenizer; +use winnow::error::{ErrMode, ModalResult}; +use winnow::prelude::*; +use winnow::stream::{Partial, Stream}; +use winnow::token::{literal, rest}; + +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::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +/// Partial streaming input over buffered decoded text. +pub type Input<'i> = Partial<&'i str>; + +/// Per-model format description consumed by [`ConfigDrivenParser`]. +/// +/// All items are associated constants or static methods, so implementors are +/// zero-sized marker types and the engine monomorphizes per format. +pub trait ParserFormat: 'static { + /// Human-readable format name used in error messages. + const NAME: &'static str; + + /// Reasoning block opener (e.g. `<|channel>thought\n`), or `None` when the + /// format has no reasoning markers. Free to span multiple tokens. + const REASONING_START: Option<&'static str>; + + /// Reasoning block closer, or `None` when the format has no reasoning. + const REASONING_END: Option<&'static str>; + + /// Start marker of one tool call. Consumed by the engine before + /// [`ParserFormat::tool_header`] runs. + const CALL_START: &'static str; + + /// End marker of one tool call. Scan target of [`ParserFormat::ArgsScan`]. + const CALL_END: &'static str; + + /// Wrapper markers around consecutive tool calls (e.g. Qwen3's + /// ``/`` around `` blocks), or `None` + /// when each `CALL_START`..`CALL_END` stands alone. + const WRAPPER: Option<(&'static str, &'static str)> = None; + + /// Swallow all output after the wrapper closes, like the `IgnoredRest` + /// event in standalone parsers. Only meaningful with `WRAPPER`. + const SUPPRESS_TEXT_AFTER_WRAPPER: bool = false; + + /// Token texts of the reasoning boundary markers, used only to derive the + /// initial parser state from prompt token IDs (never for stream matching). + const REASONING_BOUNDARY_TOKENS: Option<(&'static str, &'static str)> = None; + + /// Whether decoded output must keep tokenizer special tokens. + const PRESERVE_SPECIAL_TOKENS: bool = true; + + /// Resumable scanner locating `CALL_END` for the raw argument body. + type ArgsScan: ArgsEndScan; + + /// Parse the tool-call header after `CALL_START`: the tool name plus + /// everything up to the start of the raw argument body. + fn tool_header(input: &mut Input<'_>) -> ModalResult; + + /// Parse one complete raw argument body (end marker already stripped) + /// into a JSON object. + fn tool_args(body: &str) -> ModalResult>; +} + +/// Resumable scanner that locates a tool-call end marker in buffered input. +/// +/// On success the scanner consumes through the end marker and returns the +/// body before it. On incomplete input it must record a resume position so +/// repeated scans stay linear in the argument length. +pub trait ArgsEndScan: Default + Send { + /// Scan buffered `input` for `end`, resuming from prior progress. + fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<&'i str>; +} + +/// Default argument scanner: a plain resumable marker scan with no lexical +/// structure awareness. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct PlainArgsScan(MarkerScanState); + +impl ArgsEndScan for PlainArgsScan { + fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<&'i str> { + let body = take_until_marker(end, &mut self.0).parse_next(input)?; + input.next_slice(end.len()); + Ok(body) + } +} + +/// Boundary marker recognized by the fixed transition graph. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Boundary { + ReasoningStart, + ReasoningEnd, + WrapperStart, + WrapperEnd, + CallStart, +} + +/// One parsed engine event over buffered streaming input. +#[derive(Debug, Clone, PartialEq)] +enum Event { + /// A safe run of visible text (value is the consumed buffer prefix). + Text, + /// A safe run of reasoning text (value is the consumed buffer prefix). + Reasoning, + /// Consumed input that produces no output (between-call noise, suppressed + /// trailing text). + Ignored, + /// A boundary marker driving a fixed state transition. + Boundary(Boundary), + /// A complete tool-call header. + CallHeader { name: String }, + /// A complete tool call with parsed arguments. + CallComplete { args: Map }, +} + +/// Engine parsing state. The state set and transition graph are fixed; +/// `Option`-typed format constants decide which edges are reachable. +enum Mode { + /// Visible assistant text. + Text, + /// Inside a reasoning block. + Reasoning, + /// Inside the wrapper, between tool calls. Unreachable without `WRAPPER`. + ToolBetween, + /// After `CALL_START`, parsing the tool-call header. + ToolHeader, + /// Inside the raw argument body, scanning for `CALL_END`. + ToolArgs { + name: String, + /// Raw header text consumed by `tool_header`, kept for `reset()`. + raw_header: String, + scan: F::ArgsScan, + }, + /// Wrapper closed with `SUPPRESS_TEXT_AFTER_WRAPPER`: swallow the rest. + Done, +} + +/// Boundary marker candidates watched in one engine mode, precomputed from +/// the format constants so per-event parsing builds nothing. +struct ModeMarkers { + /// Marker texts tried in order; also the safe-content fallback scan set. + texts: Vec<&'static str>, + /// Transition triggered by each marker, index-aligned with `texts`. + boundaries: Vec, +} + +impl ModeMarkers { + fn new(candidates: impl IntoIterator) -> Self { + let (texts, boundaries) = candidates.into_iter().unzip(); + Self { texts, boundaries } + } +} + +/// Per-mode marker candidates derived once from a [`ParserFormat`]. +struct ModeMarkerSet { + text: ModeMarkers, + reasoning: ModeMarkers, + between: ModeMarkers, +} + +impl ModeMarkerSet { + fn of() -> Self { + // The tool-entry marker watched in text and reasoning modes: the + // wrapper opener when the format has one, the call start otherwise. + let tool_entry = match F::WRAPPER { + Some((wrapper_start, _)) => (wrapper_start, Boundary::WrapperStart), + None => (F::CALL_START, Boundary::CallStart), + }; + let reasoning_start = F::REASONING_START.map(|marker| (marker, Boundary::ReasoningStart)); + // Tool entry inside reasoning implicitly ends it: the transition to a + // tool state simply stops routing text to the reasoning channel. + let reasoning_end = F::REASONING_END.map(|marker| (marker, Boundary::ReasoningEnd)); + let between = F::WRAPPER.map(|(_, wrapper_end)| { + [ + (F::CALL_START, Boundary::CallStart), + (wrapper_end, Boundary::WrapperEnd), + ] + }); + + Self { + text: ModeMarkers::new(reasoning_start.into_iter().chain([tool_entry])), + reasoning: ModeMarkers::new(reasoning_end.into_iter().chain([tool_entry])), + between: ModeMarkers::new(between.into_iter().flatten()), + } + } +} + +/// Unified parser driven by a [`ParserFormat`] marker type. +/// +/// Owns all streaming state; the format contributes constants and grammar. +pub struct ConfigDrivenParser { + buffer: String, + mode: Mode, + markers: ModeMarkerSet, + emitted_tool_count: usize, + tokenizer: DynTokenizer, + /// Resolved reasoning boundary token IDs for prompt-state derivation. + boundary_ids: Option<(u32, u32)>, + _format: PhantomData F>, +} + +impl ConfigDrivenParser { + /// Create a parser for one request stream. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let boundary_ids = match F::REASONING_BOUNDARY_TOKENS { + Some((start, end)) => { + let start_id = tokenizer.token_to_id(start).ok_or_else(|| { + UnifiedParserError::MissingToken { + token: start.to_string(), + } + })?; + let end_id = + tokenizer.token_to_id(end).ok_or_else(|| UnifiedParserError::MissingToken { + token: end.to_string(), + })?; + Some((start_id, end_id)) + } + None => None, + }; + + Ok(Self { + buffer: String::new(), + mode: Mode::Text, + markers: ModeMarkerSet::of::(), + emitted_tool_count: 0, + tokenizer, + boundary_ids, + _format: PhantomData, + }) + } + + fn apply_event( + &mut self, + event: Event, + consumed_len: usize, + output: &mut UnifiedParserOutput, + ) -> Result<()> { + match event { + Event::Text => output.push_text(self.buffer[..consumed_len].to_string()), + Event::Reasoning => output.push_reasoning(self.buffer[..consumed_len].to_string()), + Event::Ignored => {} + Event::Boundary(boundary) => match boundary { + Boundary::ReasoningStart => self.mode = Mode::Reasoning, + Boundary::ReasoningEnd => self.mode = Mode::Text, + Boundary::WrapperStart => self.mode = Mode::ToolBetween, + Boundary::WrapperEnd => { + self.mode = if F::SUPPRESS_TEXT_AFTER_WRAPPER { + Mode::Done + } else { + Mode::Text + }; + } + Boundary::CallStart => self.mode = Mode::ToolHeader, + }, + Event::CallHeader { name } => { + self.mode = Mode::ToolArgs { + name, + raw_header: self.buffer[..consumed_len].to_string(), + scan: F::ArgsScan::default(), + }; + } + Event::CallComplete { args } => { + let mode = std::mem::replace(&mut self.mode, Mode::Text); + let Mode::ToolArgs { name, .. } = mode else { + return Err(parsing_failed!( + "{} arguments without an active tool call", + F::NAME + )); + }; + let arguments = serde_json::to_string(&args) + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; + + output.push_call(ToolCallDelta { + tool_index: self.emitted_tool_count, + name: Some(name), + arguments, + }); + self.emitted_tool_count += 1; + self.mode = if F::WRAPPER.is_some() { + Mode::ToolBetween + } else { + Mode::Text + }; + } + } + Ok(()) + } + + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + let in_reasoning = self.boundary_ids.and_then(|(start_id, end_id)| { + last_reasoning_boundary(prompt_token_ids, start_id, end_id, self.tokenizer.as_ref()) + }); + self.mode = match in_reasoning { + Some(true) => Mode::Reasoning, + Some(false) | None => Mode::Text, + }; + } + + fn reset(&mut self) -> String { + let raw = match std::mem::replace(&mut self.mode, Mode::Text) { + Mode::Text | Mode::Done => std::mem::take(&mut self.buffer), + Mode::Reasoning => { + // Reasoning mode is only reachable through this marker (or a + // prompt already inside reasoning, where there is no marker + // to restore; the opener default keeps reset lossless for the + // common in-stream case). + let start = F::REASONING_START.unwrap_or(""); + format!("{}{}", start, std::mem::take(&mut self.buffer)) + } + // The wrapper opener was already consumed alongside previously + // emitted calls, so only the unconsumed buffer remains. + Mode::ToolBetween => std::mem::take(&mut self.buffer), + Mode::ToolHeader => { + format!("{}{}", F::CALL_START, std::mem::take(&mut self.buffer)) + } + Mode::ToolArgs { raw_header, .. } => { + format!( + "{}{}{}", + F::CALL_START, + raw_header, + std::mem::take(&mut self.buffer) + ) + } + }; + self.mode = Mode::Text; + self.emitted_tool_count = 0; + raw + } +} + +impl UnifiedParser for ConfigDrivenParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + 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 { + F::PRESERVE_SPECIAL_TOKENS + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = { + parse_buffered_event(&self.buffer, |input| { + parse_next_event::(input, &mut self.mode, &self.markers) + })? + } { + self.apply_event(event, consumed_len, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + Mode::Text => output.push_text(std::mem::take(&mut self.buffer)), + Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + // All calls were emitted; an unclosed wrapper at end of stream is + // tolerated and any buffered between-call noise is dropped. + Mode::ToolBetween | Mode::Done => {} + Mode::ToolHeader | Mode::ToolArgs { .. } => { + return Err(parsing_failed!("incomplete {} tool call", F::NAME)); + } + } + + let _ = self.reset(); + Ok(output) + } + + fn reset(&mut self) -> String { + ConfigDrivenParser::reset(self) + } +} + +/// Parse one engine event from buffered streaming input. +/// +/// In `ToolBetween`, non-marker text is separator noise (typically +/// whitespace) and is consumed without output. +fn parse_next_event( + input: &mut Input<'_>, + mode: &mut Mode, + markers: &ModeMarkerSet, +) -> ModalResult { + match mode { + Mode::Text => boundary_or_content(input, &markers.text, Event::Text), + Mode::Reasoning => boundary_or_content(input, &markers.reasoning, Event::Reasoning), + Mode::ToolBetween => boundary_or_content(input, &markers.between, Event::Ignored), + Mode::ToolHeader => F::tool_header(input).map(|name| Event::CallHeader { name }), + Mode::ToolArgs { scan, .. } => args_event::(input, scan), + Mode::Done => rest.value(Event::Ignored).parse_next(input), + } +} + +/// Parse complete tool-call arguments once the end marker is located. +fn args_event( + input: &mut Input<'_>, + scan: &mut F::ArgsScan, +) -> ModalResult { + let body = scan.scan(input, F::CALL_END)?; + let args = F::tool_args(body)?; + Ok(Event::CallComplete { args }) +} + +/// Match one of the mode's marker candidates at the cursor (in order), or +/// emit a safe content run held back before the earliest possible marker. +fn boundary_or_content( + input: &mut Input<'_>, + markers: &ModeMarkers, + content: Event, +) -> ModalResult { + // Engine invariant: every reachable mode watches at least one marker + // (`ToolBetween` without `WRAPPER` is unreachable). + debug_assert!(!markers.texts.is_empty()); + for (marker, boundary) in markers.texts.iter().zip(&markers.boundaries) { + let matched: ModalResult<&str> = literal(*marker).parse_next(input); + match matched { + Ok(_) => return Ok(Event::Boundary(*boundary)), + Err(ErrMode::Incomplete(needed)) => return Err(ErrMode::Incomplete(needed)), + Err(_) => {} + } + } + + safe_text_len_mul(input, &markers.texts).map(|_| content) +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 0410c568461..fe7c10dd57b 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -1,9 +1,14 @@ //! Unified parser interface for reasoning and tool-call deltas. mod combined; +mod config_driven; mod gemma4; pub use combined::CombinedParser; +pub use config_driven::{ + ArgsEndScan, ConfigDrivenParser, Gemma4ConfigDrivenParser, Gemma4Format, ParserFormat, + PlainArgsScan, +}; pub use gemma4::Gemma4UnifiedParser; use thiserror::Error; use thiserror_ext::Macro;