diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml
index 67c74bb560..09c3d6b5dc 100644
--- a/rust/src/parser/Cargo.toml
+++ b/rust/src/parser/Cargo.toml
@@ -74,5 +74,10 @@ name = "gemma4"
harness = false
required-features = ["test-util"]
+[[bench]]
+name = "granite4"
+harness = false
+required-features = ["test-util"]
+
[lints]
workspace = true
diff --git a/rust/src/parser/benches/granite4.rs b/rust/src/parser/benches/granite4.rs
new file mode 100644
index 0000000000..17ef867160
--- /dev/null
+++ b/rust/src/parser/benches/granite4.rs
@@ -0,0 +1,75 @@
+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::{Granite4ToolParser, Tool, ToolParser};
+
+mod utils;
+use utils::feed_parser;
+
+const CHUNK_CHARS: usize = 7;
+const LONG_ARGUMENT_BYTES: usize = 64 * 1024;
+
+fn string_args_fixture() -> String {
+ let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(LONG_ARGUMENT_BYTES));
+ let encoded_arguments = serde_json::to_string(&arguments).unwrap();
+ format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#)
+}
+
+fn object_args_fixture() -> String {
+ format!(
+ r#"{{"name":"f","arguments":{{"data":"{}"}}}}"#,
+ "x".repeat(LONG_ARGUMENT_BYTES)
+ )
+}
+
+fn parser(tools: &[Tool]) -> Box {
+ Granite4ToolParser::create(tools).expect("Granite4 parser should initialize")
+}
+
+fn run_stream_group(c: &mut Criterion, name: &str, tools: &[Tool], text: &str) {
+ let chunks = split_by_chars(text, CHUNK_CHARS);
+
+ let mut group = c.benchmark_group(name);
+ group.sample_size(50);
+ group.warm_up_time(Duration::from_millis(300));
+ group.measurement_time(Duration::from_secs(2));
+ group.throughput(Throughput::Bytes(text.len() as u64));
+
+ group.bench_function("reuse_parser", |b| {
+ let mut parser = parser(tools);
+ b.iter(|| {
+ let result = feed_parser(&mut *parser, black_box(&chunks));
+ debug_assert_eq!(result.0, "");
+ debug_assert_eq!(result.1, 1);
+ black_box(result);
+ })
+ });
+
+ group.bench_function("create_parser", |b| {
+ b.iter_batched(
+ || parser(tools),
+ |mut parser| {
+ let result = feed_parser(&mut *parser, black_box(&chunks));
+ debug_assert_eq!(result.0, "");
+ debug_assert_eq!(result.1, 1);
+ black_box(result);
+ },
+ BatchSize::SmallInput,
+ )
+ });
+
+ group.finish();
+}
+
+fn bench_granite4(c: &mut Criterion) {
+ let tools = test_tools();
+ let string_args = string_args_fixture();
+ let object_args = object_args_fixture();
+
+ run_stream_group(c, "granite4/long_string_arguments", &tools, &string_args);
+ run_stream_group(c, "granite4/long_object_arguments", &tools, &object_args);
+}
+
+criterion_group!(benches, bench_granite4);
+criterion_main!(benches);
diff --git a/rust/src/parser/src/tool/json/granite4.rs b/rust/src/parser/src/tool/json/granite4.rs
index fe1cf19022..4989578011 100644
--- a/rust/src/parser/src/tool/json/granite4.rs
+++ b/rust/src/parser/src/tool/json/granite4.rs
@@ -9,7 +9,8 @@ use super::{
tool_call_header_event,
};
use crate::tool::utils::{
- JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object,
+ JsonObjectScanState, JsonStringScanState, decode_json_str, parse_buffered_event, safe_text_len,
+ take_json_object, take_json_string,
};
use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
@@ -22,14 +23,20 @@ enum Granite4Mode {
Header,
/// Parsing the arguments value:
/// `None` until the first byte decides object vs string;
- /// `Some` while streaming an object value.
+ /// `Some` while streaming the selected value shape.
Args {
- json_scan: Option,
+ args_scan: Option,
},
/// Arguments done; consume the object's closing `}` and ``.
Close,
}
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum Granite4ArgsScan {
+ Object(JsonObjectScanState),
+ String(JsonStringScanState),
+}
+
#[derive(Debug, Clone, PartialEq, Eq)]
enum Granite4Event {
Text {
@@ -92,7 +99,7 @@ impl Granite4ToolParser {
let tool_index = self.emitted_tool_count;
self.emitted_tool_count += 1;
self.active_tool_index = Some(tool_index);
- self.mode = Granite4Mode::Args { json_scan: None };
+ self.mode = Granite4Mode::Args { args_scan: None };
output.push_call(ToolCallDelta {
tool_index,
name: Some(function_name),
@@ -187,7 +194,7 @@ fn parse_next_granite4_event(
match mode {
Granite4Mode::Text => text_event(input),
Granite4Mode::Header => header_event(input),
- Granite4Mode::Args { json_scan } => args_event(input, json_scan),
+ Granite4Mode::Args { args_scan } => args_event(input, args_scan),
Granite4Mode::Close => close_event(input),
}
}
@@ -237,14 +244,19 @@ fn header_event(input: &mut JsonToolInput<'_>) -> ModalResult {
/// once seen whole and unescaped.
fn args_event(
input: &mut JsonToolInput<'_>,
- json_scan: &mut Option,
+ args_scan: &mut Option,
) -> ModalResult {
- if let Some(scan) = json_scan {
- let len = take_json_object(input, scan)?;
- return Ok(Granite4Event::ObjectArgsDelta {
- len,
- complete: scan.complete(),
- });
+ if let Some(scan) = args_scan {
+ return match scan {
+ Granite4ArgsScan::Object(scan) => {
+ let len = take_json_object(input, scan)?;
+ Ok(Granite4Event::ObjectArgsDelta {
+ len,
+ complete: scan.complete(),
+ })
+ }
+ Granite4ArgsScan::String(scan) => string_args_event(input, scan),
+ };
}
match peek(any).parse_next(input)? {
@@ -252,12 +264,16 @@ fn args_event(
let mut scan = JsonObjectScanState::default();
let len = take_json_object(input, &mut scan)?;
let complete = scan.complete();
- *json_scan = Some(scan);
+ *args_scan = Some(Granite4ArgsScan::Object(scan));
Ok(Granite4Event::ObjectArgsDelta { len, complete })
}
- '"' => Ok(Granite4Event::StringArgs {
- decoded: json_str(input)?,
- }),
+ '"' => {
+ *args_scan = Some(Granite4ArgsScan::String(JsonStringScanState::default()));
+ let Some(Granite4ArgsScan::String(scan)) = args_scan else {
+ unreachable!("Granite4 string scan state was just initialized")
+ };
+ string_args_event(input, scan)
+ }
_ => {
let mut error = ContextError::new();
error.push(StrContext::Label("Granite4 arguments"));
@@ -266,6 +282,17 @@ fn args_event(
}
}
+fn string_args_event(
+ input: &mut JsonToolInput<'_>,
+ scan: &mut JsonStringScanState,
+) -> ModalResult {
+ let text = **input;
+ let len = take_json_string(input, scan)?;
+ Ok(Granite4Event::StringArgs {
+ decoded: decode_json_str(&text[..len])?,
+ })
+}
+
/// Parse the tool-call object's closing `}` and the `` end marker.
fn close_event(input: &mut JsonToolInput<'_>) -> ModalResult {
seq!(_: ws0, _: literal("}"), _: ws0, _: literal(TOOL_CALL_END))
@@ -420,6 +447,22 @@ mod tests {
assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#);
}
+ #[test]
+ fn granite4_long_string_args_stream_without_reparse() {
+ let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(64 * 1024));
+ let encoded_arguments = serde_json::to_string(&arguments).unwrap();
+ let input =
+ format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#);
+ let chunks = split_by_chars(&input, 7);
+ let mut parser = Granite4ToolParser::new(&test_tools());
+
+ let output = collect_stream(&mut parser, &chunks);
+
+ assert_eq!(output.calls().len(), 1);
+ assert_eq!(output.calls()[0].name.as_deref(), Some("f"));
+ assert_eq!(output.calls()[0].arguments, arguments);
+ }
+
#[test]
fn granite4_streaming_handles_marker_and_json_whitespace() {
// Granite spaces the markers (` {…} `) and the JSON
diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs
index 7025521539..bd8f6f48e9 100644
--- a/rust/src/parser/src/utils.rs
+++ b/rust/src/parser/src/utils.rs
@@ -290,8 +290,22 @@ pub fn take_json_object(
Ok(text.len())
}
-/// Parse a JSON string literal.
-pub fn json_str(input: &mut Partial<&str>) -> ModalResult {
+/// Streaming lexical state for a JSON string literal.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct JsonStringScanState {
+ scanned_len: usize,
+ escape: bool,
+}
+
+/// Parse a raw JSON string literal, resuming from the last scanned byte.
+///
+/// The returned length covers the quoted JSON string. This only scans for the
+/// string boundary; callers that need the decoded value should pass the raw
+/// slice to [`decode_json_str`].
+pub fn take_json_string(
+ input: &mut Partial<&str>,
+ state: &mut JsonStringScanState,
+) -> ModalResult {
let text = **input;
if text.is_empty() {
return incomplete();
@@ -305,37 +319,58 @@ pub fn json_str(input: &mut Partial<&str>) -> ModalResult {
));
}
- let mut escape = false;
- let mut index = 1;
+ let mut index = if state.scanned_len == 0 {
+ 1
+ } else if state.scanned_len <= bytes.len() {
+ state.scanned_len
+ } else {
+ return incomplete();
+ };
+
while index < bytes.len() {
let byte = bytes[index];
index += 1;
- if escape {
- escape = false;
+ if state.escape {
+ state.escape = false;
continue;
}
match byte {
- b'\\' => escape = true,
+ b'\\' => state.escape = true,
b'"' => {
- let raw = &text[..index];
- let value = serde_json::from_str::(raw).map_err(|_| {
- json_scan_error(
- "JSON string",
- StrContextValue::Description("valid JSON string"),
- )
- })?;
input.next_slice(index);
- return Ok(value);
+ return Ok(index);
}
_ => {}
}
}
+ state.scanned_len = text.len();
incomplete()
}
+/// Parse a JSON string literal.
+pub fn json_str(input: &mut Partial<&str>) -> ModalResult {
+ let text = **input;
+ let checkpoint = input.checkpoint();
+ let mut state = JsonStringScanState::default();
+ let len = take_json_string(input, &mut state)?;
+ decode_json_str(&text[..len]).inspect_err(|_| {
+ input.reset(&checkpoint);
+ })
+}
+
+/// Decode a complete JSON string literal.
+pub fn decode_json_str(raw: &str) -> ModalResult {
+ serde_json::from_str::(raw).map_err(|_| {
+ json_scan_error(
+ "JSON string",
+ StrContextValue::Description("valid JSON string"),
+ )
+ })
+}
+
fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode {
let mut error = ContextError::new();
error.push(StrContext::Label(label));
@@ -388,8 +423,8 @@ mod tests {
use winnow::stream::{Offset, Partial, Stream};
use super::{
- JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len,
- safe_text_len_mul, take_json_object, take_until_marker,
+ JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len,
+ safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker,
};
#[test]
@@ -714,6 +749,71 @@ mod tests {
.assert_eq(&error.to_string());
}
+ #[test]
+ fn take_json_string_consumes_complete_string() {
+ let mut state = JsonStringScanState::default();
+ let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);
+ let checkpoint = input.checkpoint();
+
+ let len = take_json_string(&mut input, &mut state).unwrap();
+
+ assert_eq!(len, r#""say_\"hi\u0021""#.len());
+ assert_eq!(input.offset_from(&checkpoint), len);
+ assert_eq!(*input, " rest");
+ }
+
+ #[test]
+ fn take_json_string_resumes_after_incomplete_input() {
+ let mut state = JsonStringScanState::default();
+ let mut input = Partial::new(r#""{\"data\":\"partial"#);
+ let checkpoint = input.checkpoint();
+
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
+
+ assert!(matches!(error, ErrMode::Incomplete(_)));
+ assert_eq!(input.offset_from(&checkpoint), 0);
+ assert_eq!(state.scanned_len, r#""{\"data\":\"partial"#.len());
+
+ let mut input = Partial::new(r#""{\"data\":\"partial string\"}" tail"#);
+ let len = take_json_string(&mut input, &mut state).unwrap();
+
+ assert_eq!(len, r#""{\"data\":\"partial string\"}""#.len());
+ assert_eq!(*input, " tail");
+ }
+
+ #[test]
+ fn take_json_string_tracks_escape_across_chunks() {
+ let mut state = JsonStringScanState::default();
+ let mut input = Partial::new(r#""abc\"#);
+
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
+
+ assert!(matches!(error, ErrMode::Incomplete(_)));
+ assert!(state.escape);
+
+ let mut input = Partial::new(r#""abc\"def" tail"#);
+ let len = take_json_string(&mut input, &mut state).unwrap();
+
+ assert_eq!(len, r#""abc\"def""#.len());
+ assert_eq!(*input, " tail");
+ }
+
+ #[test]
+ fn take_json_string_rejects_non_string_start() {
+ let mut state = JsonStringScanState::default();
+ let mut input = Partial::new("42");
+
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
+
+ let ErrMode::Cut(error) = error else {
+ panic!("expected cut error");
+ };
+ expect![[r#"
+ invalid JSON string
+ expected `"`"#]]
+ .assert_eq(&error.to_string());
+ }
+
#[test]
fn json_str_decodes_escaped_content() {
let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);