karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

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

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-06-25 16:59:57 -07:00
committed by GitHub
parent c53994e134
commit f9e684499f
23 changed files with 1000 additions and 394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+2 -1
View File
@@ -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
rust/Cargo.toml
+1 -1
View File
@@ -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",
rust/src/chat/src/output/default/mod.rs
+149 -15
View File
@@ -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<Self> {
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<dyn UnifiedParser>
};
let reasoning_parser =
Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?;
let parser: Box<dyn UnifiedParser> =
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<Box<dyn ToolParser>> {
@@ -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<Option<Box<dyn UnifiedParser>>> {
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<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
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<u32> {
match token {
"<|channel>" => Some(1),
"<channel|>" => Some(2),
_ => None,
}
}
}
fn tokenizer() -> Arc<FakeTokenizer> {
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"
);
}
}
rust/src/chat/src/output/default/structural_tag.rs
+12 -12
View File
@@ -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<StructuralTagModel>,
) -> 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<StructuralTagTool
mod tests {
use serde_json::{Value, json};
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
use vllm_parser::tool::{Qwen3CoderToolParser, Tool};
use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
use super::*;
@@ -134,7 +134,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -147,7 +147,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -163,7 +163,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -179,7 +179,7 @@ mod tests {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -197,7 +197,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -218,7 +218,7 @@ mod tests {
);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request).to_string();
@@ -231,7 +231,7 @@ mod tests {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -247,7 +247,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.as_ref())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
let params = structured_outputs(&request);
rust/src/chat/src/output/default/unified.rs
+1 -1
View File
@@ -317,7 +317,7 @@ mod tests {
where
Self: Sized + 'static,
{
unreachable!("ScriptedParser is constructed directly in tests")
Ok(Box::new(Self::new([])))
}
fn parse_into(
rust/src/chat/src/parser/mod.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod reasoning;
pub mod tool;
pub mod unified;
use std::collections::HashMap;
use std::convert::Infallible;
rust/src/chat/src/parser/reasoning/mod.rs
+21 -10
View File
@@ -1,13 +1,13 @@
//! Reasoning parser registration and selection boundary for `vllm-chat`.
use std::sync::LazyLock;
use std::sync::{Arc, LazyLock};
pub use vllm_parser::reasoning::{
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser, Step3p5ReasoningParser,
DeepSeekV4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser,
MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, NemotronV3ReasoningParser,
Qwen3ReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser,
Step3ReasoningParser, Step3p5ReasoningParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -33,8 +33,9 @@ pub mod names {
}
/// Constructor signature for one registered reasoning parser implementation.
type ReasoningParserCreator =
fn(DynTokenizer) -> vllm_parser::reasoning::Result<Box<dyn ReasoningParser>>;
type ReasoningParserCreator = Arc<
dyn Fn(DynTokenizer) -> vllm_parser::reasoning::Result<Box<dyn ReasoningParser>> + Send + Sync,
>;
/// Registry and model matcher for reasoning parsers.
pub type ReasoningParserFactory = ParserFactory<ReasoningParserCreator>;
@@ -58,7 +59,7 @@ impl ReasoningParserFactory {
.register_parser::<DeepSeekR1ReasoningParser>(names::DEEPSEEK_R1)
.register_parser::<DeepSeekV3ReasoningParser>(names::DEEPSEEK_V3)
.register_parser::<DeepSeekV4ReasoningParser>(names::DEEPSEEK_V4)
.register_parser::<Gemma4ReasoningParser>(names::GEMMA4)
.register_unified_dummy(names::GEMMA4)
.register_parser::<Glm45ReasoningParser>(names::GLM45)
.register_parser::<KimiReasoningParser>(names::KIMI)
.register_parser::<KimiK2ReasoningParser>(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(
&registered_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(),
rust/src/chat/src/parser/reasoning/tests.rs
+2
View File
@@ -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]
rust/src/chat/src/parser/tool/mod.rs
+20 -9
View File
@@ -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<Box<dyn ToolParser>>;
type ToolParserCreator =
Arc<dyn Fn(&[ChatTool]) -> vllm_parser::tool::Result<Box<dyn ToolParser>> + Send + Sync>;
/// Registry and model matcher for tool parsers.
pub type ToolParserFactory = ParserFactory<ToolParserCreator>;
@@ -65,7 +66,7 @@ impl ToolParserFactory {
.register_parser::<DeepSeekV4ToolParser>(names::DEEPSEEK_V4)
.register_parser::<Glm45MoeToolParser>(names::GLM45)
.register_parser::<Glm47MoeToolParser>(names::GLM47)
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
.register_unified_dummy(names::GEMMA4)
.register_parser::<Granite4ToolParser>(names::GRANITE4)
.register_parser::<HermesToolParser>(names::HERMES)
.register_parser::<HyV3ToolParser>(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(
&registered_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(),
rust/src/chat/src/parser/unified.rs
+124
View File
@@ -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<Box<dyn UnifiedParser>>;
/// Registry and model matcher for unified parsers.
pub type UnifiedParserFactory = ParserFactory<UnifiedParserCreator>;
impl UnifiedParserFactory {
/// Get the global unified parser factory with built-in registrations and
/// model mappings.
pub fn global() -> &'static Self {
static INSTANCE: LazyLock<UnifiedParserFactory> = 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::<Gemma4UnifiedParser>(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<T>(&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<Box<dyn UnifiedParser>> {
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<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
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<u32> {
match token {
"<|channel>" => Some(1),
"<channel|>" => 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();
}
}
rust/src/chat/tests/roundtrip.rs
+50 -3
View File
@@ -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<String> {
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::<Vec<_>>();
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<vllm_chat::LoadedModelBackends> {
load_model_backends(
rust/src/parser/benches/gemma4.rs
+5 -3
View File
@@ -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<dyn ToolParser> {
Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
UnifiedToolParserAdapter::<Gemma4UnifiedParser>::create(tools)
.expect("Gemma4 unified parser should initialize")
}
fn run_stream_group(
rust/src/parser/benches/utils/adapter.rs
+106
View File
@@ -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<Vec<u32>> {
Ok(text.chars().map(|_| u32::MAX).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok("\u{FFFD}".repeat(token_ids.len()))
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
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<T> {
inner: Box<dyn UnifiedParser>,
_marker: std::marker::PhantomData<T>,
}
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<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
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<StructuralTagModel> {
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<ToolParserOutput> {
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()
}
}
rust/src/parser/benches/utils/mod.rs
+7
View File
@@ -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<OpenAiTool> {
tools
.iter()
rust/src/parser/src/lib.rs
+1
View File
@@ -3,3 +3,4 @@
pub mod reasoning;
pub mod tool;
pub mod unified;
pub(crate) mod utils;
rust/src/parser/src/reasoning/delimited.rs
+5 -5
View File
@@ -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<bool> {
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;
}
}
rust/src/parser/src/reasoning/gemma4.rs
-273
View File
@@ -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> ... <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:
/// <https://github.com/vllm-project/vllm/blob/18b1c77211d8f6fe800bcfb89524d2b598708032/vllm/reasoning/gemma4_reasoning_parser.py#L23>
pub struct Gemma4ReasoningParser {
inner: DelimitedReasoningParser,
reasoning_text: String,
prefix_stripped: bool,
}
impl Gemma4ReasoningParser {
/// Create a Gemma4 parser.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(tokenizer, "<|channel>", "<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<String> {
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<Box<dyn ReasoningParser>>
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<ReasoningDelta> {
let result = self.inner.push(delta);
Ok(self.post_process(result))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
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<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
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<u32> {
match token {
"<|channel>" => Some(1000),
"<channel|>" => Some(1001),
_ => None,
}
}
}
fn run_streaming(output: &[&str]) -> (Option<String>, Option<String>) {
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 section<channel|>This is the rest"],
Some("This is a reasoning section"),
Some("This is the rest"),
),
(
"complete_reasoning",
vec!["<|channel>This is a reasoning section<channel|>"],
Some("This is a reasoning section"),
None,
),
(
"multiple_lines",
vec!["<|channel>This\nThat<channel|>This 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<channel|>\nThis is the rest"],
Some("This is a reasoning section"),
Some("Before\n\nThis is the rest"),
),
(
"thought_prefix",
vec!["<|channel>thought\nActual reasoning here<channel|>Final answer"],
Some("Actual reasoning here"),
Some("Final answer"),
),
(
"thought_prefix_only",
vec!["<|channel>thought\n<channel|>"],
None,
None,
),
(
"thought_prefix_multiline",
vec!["<|channel>thought\nLine1\nLine2<channel|>Answer"],
Some("Line1\nLine2"),
Some("Answer"),
),
(
"thought_prefix_diverge",
vec!["<|channel>thousand reasons<channel|>Done"],
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", "<channel|>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());
}
}
rust/src/parser/src/reasoning/mod.rs
+5 -3
View File
@@ -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)]
rust/src/parser/src/tool/error.rs
+4
View File
@@ -10,4 +10,8 @@ pub type Result<T> = std::result::Result<T, ToolParserError>;
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 },
}
rust/src/parser/src/tool/mod.rs
+1 -3
View File
@@ -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::{
rust/src/parser/src/tool/gemma4.rs → rust/src/parser/src/unified/gemma4.rs
+371 -38
View File
@@ -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 = "<channel|>";
const TOOL_CALL_START: &str = "<|tool_call>";
const TOOL_CALL_END: &str = "<tool_call|>";
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<String, Value> },
@@ -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:
/// <https://github.com/vllm-project/vllm/blob/bf45e6d0a558da2b8d7b60efb07b4aa394f3b60b/vllm/tool_parsers/gemma4_tool_parser.py>
/// <https://github.com/vllm-project/vllm/blob/main/vllm/parser/gemma4.py>
///
/// Handles the Gemma4 function call format:
/// Handles Gemma4 reasoning and function-call formats:
///
/// `<|channel>thought\nreasoning<channel|>`
///
/// `<|tool_call>call:func_name{key:<|"|>value<|"|>}<tool_call|>`
///
/// 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<Self> {
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<Box<dyn ToolParser>>
impl UnifiedParser for Gemma4UnifiedParser {
fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tools)))
Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box<dyn UnifiedParser>)
}
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<ToolParserOutput> {
let mut output = ToolParserOutput::default();
fn finish(&mut self) -> Result<UnifiedParserOutput> {
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<Gemma4Event> {
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<Gemma4Event> {
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<Gemma4Event> {
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<Gemma4Event> {
literal(REASONING_START).value(Gemma4Event::ReasoningStart).parse_next(input)
}
/// Parse a Gemma4 reasoning end marker.
fn reasoning_end_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
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<String> {
/// Parse a safe text run before the next Gemma4 marker.
fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
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<Gemma4Event> {
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<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
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<u32> {
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<Vec<u32>> {
Ok(text.chars().map(u32::from).collect())
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
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<u32> {
None
}
}
trait UnifiedParserTestExt {
fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput>;
fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput>;
}
impl UnifiedParserTestExt for Gemma4UnifiedParser {
fn parse_chunk(&mut self, chunk: &str) -> super::Result<UnifiedParserOutput> {
let mut output = UnifiedParserOutput::default();
self.parse_into(chunk, &mut output)?;
Ok(output)
}
fn parse_complete(&mut self, text: &str) -> super::Result<UnifiedParserOutput> {
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<Vec<Value>> {
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<|\"|>}<tool_call|>")
.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\nreason<channel|>answer"]);
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("son<channel|>answer").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("reason<channel|>answer").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\nreason<channel|>answer").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<|\"|>}",
"<tool_call|>",
]);
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::<Value>(&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}<tool_call|>";
let _error = parser.parse_chunk(input).unwrap_err();
rust/src/parser/src/unified/mod.rs
+9 -1
View File
@@ -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)]
rust/src/parser/src/tool/utils.rs → rust/src/parser/src/utils.rs
+103 -16
View File
@@ -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<usize> {
pub fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalResult<usize> {
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<usize> {
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<usize> {
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<Partial<&'i str>, &'i str, ErrMode<ContextError>> + '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<usize> {
@@ -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<String> {
pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
let text = **input;
if text.is_empty() {
return incomplete();
@@ -311,7 +350,7 @@ fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode<Co
/// of bytes consumed from the buffer.
/// - `Ok(None)` if the buffer does not contain a full event yet, and more data is needed.
/// - `Err` if a parsing error occurred.
pub(super) fn parse_buffered_event<E>(
pub fn parse_buffered_event<E>(
buffer: &str,
parse: impl FnOnce(&mut Partial<&str>) -> ModalResult<E>,
) -> Result<Option<(E, usize)>> {
@@ -322,7 +361,9 @@ pub(super) fn parse_buffered_event<E>(
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<E>(
}
/// Returns an error indicating that we need more data to continue parsing.
pub(super) fn incomplete<T>() -> ModalResult<T> {
pub fn incomplete<T>() -> ModalResult<T> {
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<channel|><|tool_call>");
let checkpoint = input.checkpoint();
let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<channel|>"]).unwrap();
assert_eq!(len, "hello".len());
assert_eq!(input.offset_from(&checkpoint), "hello".len());
assert_eq!(*input, "<channel|><|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<not_marker><|tool_call>");
let checkpoint = input.checkpoint();
let len = safe_text_len_mul(&mut input, &["<|tool_call>", "<|channel>thought\n"]).unwrap();
assert_eq!(len, "hello<not_marker>".len());
assert_eq!(input.offset_from(&checkpoint), "hello<not_marker>".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();