forked from Karylab-cklius/vllm
525 lines
18 KiB
Rust
525 lines
18 KiB
Rust
// SPDX-License-Identifier: Apache-2.0
|
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
use winnow::ascii::multispace0 as ws0;
|
|
use winnow::combinator::seq;
|
|
use winnow::error::{ModalResult, StrContext};
|
|
use winnow::prelude::*;
|
|
use winnow::token::literal;
|
|
|
|
use super::{
|
|
JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput,
|
|
argument_delta_event, tool_call_header_event,
|
|
};
|
|
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
|
|
use crate::tool::{
|
|
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
|
|
};
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
enum LlamaJsonMode {
|
|
Start,
|
|
Header,
|
|
Arguments { json_scan: JsonObjectScanState },
|
|
AfterCall,
|
|
Passthrough,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
enum LlamaJsonEvent {
|
|
ToolCallHeader { function_name: String },
|
|
Arguments { len: usize },
|
|
ToolCallClose,
|
|
Separator,
|
|
}
|
|
|
|
/// Tool parser for strict Llama JSON-template tool calls.
|
|
///
|
|
/// Example tool call content:
|
|
///
|
|
/// ```text
|
|
/// {"name":"get_weather","parameters":{"location":"Tokyo"}}; {"name":"add","parameters":{"x":1,"y":2}}
|
|
/// ```
|
|
///
|
|
/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
|
|
/// argument deltas without schema conversion or JSON normalization.
|
|
///
|
|
/// Natural text at the beginning of the stream permanently disables tool
|
|
/// parsing for that assistant output.
|
|
pub struct Llama3JsonToolParser {
|
|
buffer: String,
|
|
mode: LlamaJsonMode,
|
|
active_tool_index: Option<usize>,
|
|
emitted_tool_count: usize,
|
|
}
|
|
|
|
impl Llama3JsonToolParser {
|
|
/// Create a Llama JSON tool parser.
|
|
fn new(_tools: &[Tool]) -> Self {
|
|
Self {
|
|
buffer: String::new(),
|
|
mode: LlamaJsonMode::Start,
|
|
active_tool_index: None,
|
|
emitted_tool_count: 0,
|
|
}
|
|
}
|
|
|
|
/// Commit the stream to JSON parsing or permanent passthrough.
|
|
fn commit_start(&mut self) -> bool {
|
|
if !matches!(self.mode, LlamaJsonMode::Start) {
|
|
return true;
|
|
}
|
|
|
|
if self.buffer.is_empty() {
|
|
return false;
|
|
}
|
|
|
|
if self.buffer.starts_with('{') {
|
|
self.mode = LlamaJsonMode::Header;
|
|
} else {
|
|
self.mode = LlamaJsonMode::Passthrough;
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Apply one parsed Llama JSON event to parser state and output.
|
|
fn apply_event(&mut self, event: LlamaJsonEvent, output: &mut ToolParserOutput) -> Result<()> {
|
|
match event {
|
|
LlamaJsonEvent::ToolCallHeader { function_name } => {
|
|
let tool_index = self.emitted_tool_count;
|
|
self.emitted_tool_count += 1;
|
|
self.active_tool_index = Some(tool_index);
|
|
self.mode = LlamaJsonMode::Arguments {
|
|
json_scan: JsonObjectScanState::default(),
|
|
};
|
|
output.push_call(ToolCallDelta {
|
|
tool_index,
|
|
name: Some(function_name),
|
|
arguments: String::new(),
|
|
});
|
|
}
|
|
LlamaJsonEvent::Arguments { len: consumed_len } => {
|
|
let Some(tool_index) = self.active_tool_index else {
|
|
return Err(parsing_failed!(
|
|
"Llama JSON arguments without an active tool call"
|
|
));
|
|
};
|
|
output.push_call(ToolCallDelta {
|
|
tool_index,
|
|
name: None,
|
|
arguments: self.buffer[..consumed_len].to_string(),
|
|
});
|
|
}
|
|
LlamaJsonEvent::ToolCallClose => {
|
|
self.active_tool_index = None;
|
|
self.mode = LlamaJsonMode::AfterCall;
|
|
}
|
|
LlamaJsonEvent::Separator => {
|
|
self.active_tool_index = None;
|
|
self.mode = LlamaJsonMode::Header;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn reset(&mut self) -> String {
|
|
self.mode = LlamaJsonMode::Start;
|
|
self.active_tool_index = None;
|
|
self.emitted_tool_count = 0;
|
|
std::mem::take(&mut self.buffer)
|
|
}
|
|
}
|
|
|
|
impl ToolParser for Llama3JsonToolParser {
|
|
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
|
|
where
|
|
Self: Sized + 'static,
|
|
{
|
|
Ok(Box::new(Self::new(tools)))
|
|
}
|
|
|
|
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
|
Some(xgrammar_structural_tag::Model::Llama.builder())
|
|
}
|
|
|
|
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
|
self.buffer.push_str(chunk);
|
|
|
|
if !self.commit_start() {
|
|
return Ok(());
|
|
}
|
|
|
|
if matches!(self.mode, LlamaJsonMode::Passthrough) {
|
|
output.push_text(&self.buffer);
|
|
self.buffer.clear();
|
|
return Ok(());
|
|
}
|
|
|
|
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
|
|
parse_next_llama_json_event(input, &mut self.mode)
|
|
})? {
|
|
self.apply_event(event, output)?;
|
|
self.buffer.drain(..consumed_len);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn finish(&mut self) -> Result<ToolParserOutput> {
|
|
let mut output = ToolParserOutput::default();
|
|
match &self.mode {
|
|
LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
|
|
output.push_text(&self.buffer);
|
|
}
|
|
LlamaJsonMode::AfterCall if self.buffer.trim().is_empty() => {}
|
|
LlamaJsonMode::Header | LlamaJsonMode::Arguments { .. } => {
|
|
return Err(parsing_failed!("incomplete Llama JSON tool call"));
|
|
}
|
|
LlamaJsonMode::AfterCall => {
|
|
return Err(parsing_failed!("invalid Llama JSON"));
|
|
}
|
|
}
|
|
let _ = self.reset();
|
|
Ok(output)
|
|
}
|
|
|
|
fn reset(&mut self) -> String {
|
|
Llama3JsonToolParser::reset(self)
|
|
}
|
|
}
|
|
|
|
/// Parse a Llama JSON event for the current parser mode.
|
|
fn parse_next_llama_json_event(
|
|
input: &mut JsonToolInput<'_>,
|
|
mode: &mut LlamaJsonMode,
|
|
) -> ModalResult<LlamaJsonEvent> {
|
|
match mode {
|
|
LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
|
|
unreachable!("Llama JSON parser driver must commit before parsing events")
|
|
}
|
|
LlamaJsonMode::Header => llama_tool_call_header_event(input),
|
|
LlamaJsonMode::Arguments { json_scan } => parse_llama_arguments_event(input, json_scan),
|
|
LlamaJsonMode::AfterCall => after_call_event(input),
|
|
}
|
|
}
|
|
|
|
/// Parse a Llama JSON tool-call header.
|
|
fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
|
|
const CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
|
parser_name: "Llama JSON",
|
|
start_marker: "",
|
|
end_marker: "",
|
|
marker_whitespace: JsonToolCallWhitespace::Optional,
|
|
delimiter: Some(";"),
|
|
name_key: "name",
|
|
arguments_key: &["parameters"],
|
|
};
|
|
|
|
match tool_call_header_event(input, CONFIG)? {
|
|
JsonToolCallEvent::ToolCallHeader { function_name } => {
|
|
Ok(LlamaJsonEvent::ToolCallHeader { function_name })
|
|
}
|
|
_ => unreachable!("tool_call_header_event only emits ToolCallHeader"),
|
|
}
|
|
}
|
|
|
|
/// Parse one event inside a Llama JSON arguments payload.
|
|
fn parse_llama_arguments_event(
|
|
input: &mut JsonToolInput<'_>,
|
|
json_scan: &mut JsonObjectScanState,
|
|
) -> ModalResult<LlamaJsonEvent> {
|
|
if json_scan.complete() {
|
|
tool_call_close_event(input)
|
|
} else {
|
|
match argument_delta_event(input, json_scan)? {
|
|
JsonToolCallEvent::Arguments { len } => Ok(LlamaJsonEvent::Arguments { len }),
|
|
_ => unreachable!("argument_delta_event only emits Arguments"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Parse the outer closing brace for one Llama JSON tool call.
|
|
fn tool_call_close_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
|
|
seq!(_: ws0, _: literal("}"))
|
|
.value(LlamaJsonEvent::ToolCallClose)
|
|
.parse_next(input)
|
|
}
|
|
|
|
/// Parse a semicolon separator after one Llama JSON tool call.
|
|
fn after_call_event(input: &mut JsonToolInput<'_>) -> ModalResult<LlamaJsonEvent> {
|
|
seq!(
|
|
_: ws0,
|
|
_: literal(";"),
|
|
_: ws0,
|
|
)
|
|
.value(LlamaJsonEvent::Separator)
|
|
.context(StrContext::Label("Llama JSON"))
|
|
.parse_next(input)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use expect_test::expect;
|
|
use thiserror_ext::AsReport;
|
|
|
|
use super::Llama3JsonToolParser;
|
|
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
|
use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
|
|
|
|
fn build_tool_call(function_name: &str, parameters: &str) -> String {
|
|
format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#)
|
|
}
|
|
|
|
#[test]
|
|
fn llama_tolerates_whitespace_before_outer_brace() {
|
|
// Whitespace between the parameters object's `}` and the outer `}` must
|
|
// still parse (json.loads / raw_decode parity).
|
|
let mut whole = Llama3JsonToolParser::new(&test_tools());
|
|
let whole_output = whole.parse_complete(r#"{"name":"f","parameters":{"x":1} }"#).unwrap();
|
|
assert_eq!(whole_output.calls().len(), 1);
|
|
assert_eq!(whole_output.calls()[0].name.as_deref(), Some("f"));
|
|
assert_eq!(whole_output.calls()[0].arguments, r#"{"x":1}"#);
|
|
|
|
// Same input, whitespace before the outer `}` split across a chunk boundary.
|
|
let mut chunked = Llama3JsonToolParser::new(&test_tools());
|
|
let mut output = ToolParserOutput::default();
|
|
for chunk in [r#"{"name":"f","parameters":{"x":1}"#, " ", "}"] {
|
|
output.append(chunked.parse_chunk(chunk).unwrap());
|
|
}
|
|
output.append(chunked.finish().unwrap());
|
|
let output = output.coalesce();
|
|
assert_eq!(output.calls().len(), 1);
|
|
assert_eq!(output.calls()[0].arguments, r#"{"x":1}"#);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_parse_complete_without_tool_call_keeps_text() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let output = parser.parse_complete("Hello, world!").unwrap();
|
|
|
|
assert_eq!(output.normal_text(), "Hello, world!");
|
|
assert!(output.calls().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_passthrough_never_reenters_tool_parsing() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let mut output = parser.parse_chunk("plain text first ").unwrap();
|
|
output.append(
|
|
parser
|
|
.parse_chunk(&build_tool_call("get_weather", r#"{"location":"Tokyo"}"#))
|
|
.unwrap(),
|
|
);
|
|
output.append(parser.finish().unwrap());
|
|
|
|
assert_eq!(
|
|
output.normal_text(),
|
|
r#"plain text first {"name":"get_weather","parameters":{"location":"Tokyo"}}"#
|
|
);
|
|
assert!(output.calls().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_does_not_support_python_tag_prefix() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let input = format!(
|
|
"<|python_tag|>{}",
|
|
build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
|
|
);
|
|
let output = parser.parse_complete(&input).unwrap();
|
|
|
|
assert_eq!(output.normal_text(), input);
|
|
assert!(output.calls().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_rejects_leading_whitespace_before_tool_call() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let input = format!(
|
|
"\n {}",
|
|
build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
|
|
);
|
|
let output = parser.parse_complete(&input).unwrap();
|
|
|
|
assert_eq!(output.normal_text(), input);
|
|
assert!(output.calls().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_extracts_raw_parameters_object() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let arguments = r#"{ "location": "Tokyo", "days": 3 }"#;
|
|
let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
|
|
|
|
assert_eq!(output.calls().len(), 1);
|
|
assert_eq!(output.calls()[0].tool_index, 0);
|
|
assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather"));
|
|
assert_eq!(output.calls()[0].arguments, arguments);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_rejects_arguments_key() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let error = parser
|
|
.parse_complete(r#"{"name":"get_weather","arguments":{"location":"Tokyo"}}"#)
|
|
.unwrap_err();
|
|
|
|
expect![[r#"
|
|
tool parser parsing failed: near "{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Tokyo\"}}": invalid Llama JSON
|
|
expected `parameters`"#]]
|
|
.assert_eq(&error.to_report_string());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_extracts_multiple_semicolon_separated_calls() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let input = format!(
|
|
"{} \n; {}",
|
|
build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
|
|
build_tool_call("add", r#"{"x":1,"y":2}"#),
|
|
);
|
|
let output = parser.parse_complete(&input).unwrap();
|
|
|
|
expect![[r#"
|
|
ToolParserOutput {
|
|
events: [
|
|
ToolCall(
|
|
ToolCallDelta {
|
|
tool_index: 0,
|
|
name: Some(
|
|
"get_weather",
|
|
),
|
|
arguments: "{\"location\":\"Shanghai\"}",
|
|
},
|
|
),
|
|
ToolCall(
|
|
ToolCallDelta {
|
|
tool_index: 1,
|
|
name: Some(
|
|
"add",
|
|
),
|
|
arguments: "{\"x\":1,\"y\":2}",
|
|
},
|
|
),
|
|
],
|
|
}
|
|
"#]]
|
|
.assert_debug_eq(&output);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_streaming_emits_argument_deltas() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let chunks = [
|
|
"{\"name\":\"get_weather\",\"parameters\":",
|
|
"{\"location\":",
|
|
"\"Beijing\"",
|
|
"}}",
|
|
];
|
|
|
|
let mut output = ToolParserOutput::default();
|
|
let mut observed_arguments = Vec::new();
|
|
for chunk in chunks {
|
|
let next = parser.parse_chunk(chunk).unwrap();
|
|
observed_arguments.extend(
|
|
next.calls()
|
|
.iter()
|
|
.filter(|call| call.name.is_none())
|
|
.map(|call| call.arguments.clone()),
|
|
);
|
|
output.append(next);
|
|
}
|
|
output.append(parser.finish().unwrap());
|
|
|
|
assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
|
|
assert_eq!(
|
|
output.coalesce().calls()[0].arguments,
|
|
r#"{"location":"Beijing"}"#
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_streaming_handles_split_objects_and_separator() {
|
|
let input = format!(
|
|
"{};{}",
|
|
build_tool_call("get_weather", r#"{"location":"Dallas","state":"TX"}"#),
|
|
build_tool_call("add", r#"{"x":4,"y":5}"#),
|
|
);
|
|
let chunks = split_by_chars(&input, 6);
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
|
|
let output = collect_stream(&mut parser, &chunks);
|
|
|
|
assert_eq!(output.normal_text(), "");
|
|
assert_eq!(output.calls().len(), 2);
|
|
assert_eq!(
|
|
output.calls()[0].arguments,
|
|
r#"{"location":"Dallas","state":"TX"}"#
|
|
);
|
|
assert_eq!(output.calls()[1].name.as_deref(), Some("add"));
|
|
assert_eq!(output.calls()[1].arguments, r#"{"x":4,"y":5}"#);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_handles_nested_multiline_and_escaped_string_parameters() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let arguments = r#"{
|
|
"payload": {"items": [1, {"value": "literal { brace } and \"quote\""}]},
|
|
"flag": true
|
|
}"#;
|
|
let output = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap();
|
|
|
|
assert_eq!(output.calls()[0].arguments, arguments);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_keeps_trailing_whitespace_after_tool_call() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let output = parser
|
|
.parse_complete(&format!(
|
|
"{}\n\t ",
|
|
build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
|
|
))
|
|
.unwrap();
|
|
|
|
assert_eq!(output.normal_text(), "");
|
|
assert_eq!(output.calls().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_finish_fails_incomplete_tool_call() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
parser.parse_chunk(r#"{"name":"get_weather","parameters":{"location""#).unwrap();
|
|
|
|
let error = parser.finish().unwrap_err();
|
|
|
|
expect!["tool parser parsing failed: incomplete Llama JSON tool call"]
|
|
.assert_eq(&error.to_report_string());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_malformed_field_order_fails_fast() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let error = parser.parse_chunk(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
|
|
|
|
expect![[r#"
|
|
tool parser parsing failed: near "{\"parameters\":{},\"name\":\"get_weather\"}": invalid Llama JSON
|
|
expected `name`"#]]
|
|
.assert_eq(&error.to_report_string());
|
|
}
|
|
|
|
#[test]
|
|
fn llama_json_trailing_non_separator_content_errors() {
|
|
let mut parser = Llama3JsonToolParser::new(&test_tools());
|
|
let error = parser
|
|
.parse_chunk(&format!(
|
|
"{} trailing",
|
|
build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
|
|
))
|
|
.unwrap_err();
|
|
|
|
expect![[r#"tool parser parsing failed: near " trailing": invalid Llama JSON"#]]
|
|
.assert_eq(&error.to_report_string());
|
|
}
|
|
}
|