use std::collections::HashMap;
use regex::RegexBuilder;
use serde_json::Value;
use uuid::Uuid;
use super::super::ToolDefinition;
use super::config::JsonParserConfig;
use super::response::{CalledFunction, ToolCallResponse, ToolCallType};
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct CalledFunctionParameters {
pub name: String,
pub parameters: HashMap<String, Value>,
}
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct CalledFunctionArguments {
pub name: String,
pub arguments: HashMap<String, Value>,
}
fn extract_tool_call_content(input: &str, start_token: &str, end_token: &str) -> Option<String> {
let escaped_start = regex::escape(start_token);
let escaped_end = regex::escape(end_token);
let pattern = format!(r"{}(.*?){}", escaped_start, escaped_end);
match RegexBuilder::new(&pattern)
.dot_matches_new_line(true)
.build()
{
Ok(regex) => {
let matches: Vec<_> = regex
.captures_iter(input)
.filter_map(|captures| captures.get(1))
.map(|m| m.as_str().trim().to_string())
.collect();
if !matches.is_empty() {
if matches.len() == 1 {
return Some(matches.last().unwrap().clone());
} else {
return Some(format!("[{}]", matches.join(",")));
}
}
None
}
Err(_) => None,
}
}
fn extract_tool_call_content_eof_recovery(input: &str, start_token: &str) -> Option<String> {
let start_pos = input.find(start_token)?;
let tail = input[start_pos + start_token.len()..].trim();
if tail.starts_with('{') || tail.starts_with('[') {
Some(tail.to_string())
} else {
None
}
}
fn handle_single_token_tool_calls(input: &str, start_token: &str) -> Option<String> {
if !input.contains(start_token) {
return None;
}
let mut items: Vec<String> = Vec::new();
for seg in input.split(start_token) {
let s = seg.trim();
if s.is_empty() {
continue;
}
if s.starts_with('{') {
if let Some(pos) = s.rfind('}') {
let candidate = &s[..=pos].trim();
if serde_json::from_str::<serde_json::Value>(candidate).is_ok() {
items.push(candidate.to_string());
}
}
} else if s.starts_with('[') {
if let Some(pos) = s.rfind(']') {
let candidate = &s[..=pos].trim();
if serde_json::from_str::<serde_json::Value>(candidate).is_ok() {
if let Ok(serde_json::Value::Array(arr)) =
serde_json::from_str::<serde_json::Value>(candidate)
{
for item in arr {
if let Ok(item_str) = serde_json::to_string(&item) {
items.push(item_str);
}
}
}
}
}
}
}
if items.is_empty() {
return Some(String::new());
}
Some(format!("[{}]", items.join(",")))
}
pub(crate) fn try_repair_truncated_json(s: &str) -> Option<String> {
let mut stack: Vec<char> = Vec::new();
let mut in_string = false;
let mut escape = false;
for c in s.chars() {
if escape {
escape = false;
continue;
}
if in_string {
match c {
'\\' => escape = true,
'"' => in_string = false,
_ => {}
}
continue;
}
match c {
'"' => in_string = true,
'{' => stack.push('}'),
'[' => stack.push(']'),
'}' | ']' => {
stack.pop();
}
_ => {}
}
}
if !escape && !in_string && stack.is_empty() {
return None;
}
let mut repaired = s.to_string();
if escape {
repaired.push('\\');
}
if in_string {
repaired.push('"');
}
while let Some(closer) = stack.pop() {
repaired.push(closer);
}
Some(repaired)
}
fn try_parse_normal_text(input: &str, start_token: &str) -> String {
if let Some(idx) = input.find(start_token) {
return input[..idx].trim().to_string();
}
String::new()
}
pub fn try_tool_call_parse_basic_json(
message: &str,
config: &JsonParserConfig,
_tools: Option<&[ToolDefinition]>,
) -> anyhow::Result<(Vec<ToolCallResponse>, Option<String>)> {
tracing::debug!("Using JSON parser config: {:?}", config);
let trimmed = message.trim();
if trimmed.is_empty() {
return Ok((vec![], Some(String::new())));
}
let tool_call_start_tokens = &config.tool_call_start_tokens;
let tool_call_end_tokens = &config.tool_call_end_tokens;
if tool_call_start_tokens.is_empty() && !config.bare_json_mode {
return Ok((vec![], Some(trimmed.to_string())));
}
let mut json = trimmed.to_string();
let mut normal_text = trimmed.to_string();
let mut found_start_token_with_no_valid_json = false;
let has_start_token = !config.bare_json_mode
&& tool_call_start_tokens
.iter()
.any(|token| !token.is_empty() && normal_text.contains(token));
if !has_start_token {
if let Some(idx) = normal_text.find(['{', '[']) {
let extracted_normal = normal_text[..idx].trim().to_string();
let extracted_json = normal_text[idx..].trim().to_string();
if !extracted_json.is_empty() {
normal_text = extracted_normal;
json = extracted_json;
}
}
} else {
'outer: for start_token in tool_call_start_tokens.iter() {
for end_token in tool_call_end_tokens.iter() {
let new_normal_text = try_parse_normal_text(&normal_text, start_token);
match (start_token.is_empty(), end_token.is_empty()) {
(false, true) => {
let result = handle_single_token_tool_calls(&json, start_token);
if let Some(content) = result {
if content.is_empty() {
found_start_token_with_no_valid_json = true;
}
json = content;
normal_text = new_normal_text;
break 'outer; }
}
(false, false) => {
let mut result = extract_tool_call_content(&json, start_token, end_token);
if result.is_none()
&& config.allow_eof_recovery
&& json.contains(start_token.as_str())
{
result = extract_tool_call_content_eof_recovery(&json, start_token);
}
if let Some(content) = result {
if content.is_empty() {
found_start_token_with_no_valid_json = true;
}
json = content;
normal_text = new_normal_text;
break 'outer; }
}
_ => {
continue;
}
}
}
}
}
let json = json.as_str();
let parse = |name: String, args: HashMap<String, Value>| -> anyhow::Result<ToolCallResponse> {
Ok(ToolCallResponse {
id: format!("call-{}", Uuid::new_v4()),
tp: ToolCallType::Function,
function: CalledFunction {
name,
arguments: serde_json::to_string(&args)?,
},
})
};
if let Ok(single) = serde_json::from_str::<CalledFunctionParameters>(json) {
return Ok((
vec![parse(single.name, single.parameters)?],
Some(normal_text),
));
} else if let Ok(single) = serde_json::from_str::<CalledFunctionArguments>(json) {
return Ok((
vec![parse(single.name, single.arguments)?],
Some(normal_text),
));
} else if let Ok(array) = serde_json::from_str::<Vec<serde_json::Value>>(json) {
let mut results = Vec::new();
for item in array {
if let Ok(func_args) = serde_json::from_value::<CalledFunctionArguments>(item.clone()) {
results.push(parse(func_args.name, func_args.arguments)?);
} else if let Ok(func_params) = serde_json::from_value::<CalledFunctionParameters>(item)
{
results.push(parse(func_params.name, func_params.parameters)?);
}
}
return Ok((results, Some(normal_text)));
}
if config.allow_eof_recovery
&& let Some(repaired) = try_repair_truncated_json(json)
{
let repaired = repaired.as_str();
if let Ok(single) = serde_json::from_str::<CalledFunctionParameters>(repaired) {
return Ok((
vec![parse(single.name, single.parameters)?],
Some(normal_text),
));
} else if let Ok(single) = serde_json::from_str::<CalledFunctionArguments>(repaired) {
return Ok((
vec![parse(single.name, single.arguments)?],
Some(normal_text),
));
} else if let Ok(array) = serde_json::from_str::<Vec<serde_json::Value>>(repaired) {
let mut results = Vec::new();
for item in array {
if let Ok(func_args) =
serde_json::from_value::<CalledFunctionArguments>(item.clone())
{
results.push(parse(func_args.name, func_args.arguments)?);
} else if let Ok(func_params) =
serde_json::from_value::<CalledFunctionParameters>(item)
{
results.push(parse(func_params.name, func_params.parameters)?);
}
}
if !results.is_empty() {
return Ok((results, Some(normal_text)));
}
}
}
if found_start_token_with_no_valid_json {
Ok((vec![], Some(String::new())))
} else {
Ok((vec![], Some(trimmed.to_string())))
}
}
pub fn detect_tool_call_start_basic_json(chunk: &str, config: &JsonParserConfig) -> bool {
let trimmed = chunk.trim();
if trimmed.is_empty() {
return false;
}
let contains_complete_token = config
.tool_call_start_tokens
.iter()
.any(|token| !token.is_empty() && trimmed.contains(token));
if contains_complete_token {
return true;
}
let has_partial_token = config.tool_call_start_tokens.iter().any(|token| {
if token.is_empty() {
return false;
}
for i in 1..=token.chars().count() {
if let Some(prefix) = token.chars().take(i).collect::<String>().get(..) {
let prefix_str = &prefix[..prefix.len()];
if trimmed == prefix_str {
return true;
}
if prefix_str.len() >= 3 && trimmed.contains(prefix_str) {
return true;
}
if prefix_str.len() < 3 && trimmed.ends_with(prefix_str) {
return true;
}
}
}
false
});
has_partial_token || trimmed.contains('{') || trimmed.contains('[')
}
#[cfg(test)]
mod repair_tests {
use super::*;
#[test]
fn test_repair_eof_after_backslash() {
let repaired = try_repair_truncated_json(r#"{"k":"a\"#).expect("must repair");
assert!(
serde_json::from_str::<serde_json::Value>(&repaired).is_ok(),
"repaired must parse: {:?}",
repaired
);
}
}
#[cfg(test)]
mod detect_parser_tests {
use super::*;
#[test] fn detect_tool_call_start_basic_json_chunk_with_tool_call_start_token_hermes() {
let text =
r#"<tool_call>{"name": "search", "parameters": { "query": "rust" } }</tool_call>"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<tool_call>".to_string()],
tool_call_end_tokens: vec!["</tool_call>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_without_tool_call_start_token() {
let text = r#"{"name": "search", "parameters": { "query": "rust" } }"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<tool_call>".to_string()],
tool_call_end_tokens: vec!["</tool_call>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_without_tool_call_start_token_with_normal_text() {
let text = r#"Here it is {"name": "#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<tool_call>".to_string()],
tool_call_end_tokens: vec!["</tool_call>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_with_square_brackets() {
let text = r#"Here it is [{"name": "search","#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<tool_call>".to_string()],
tool_call_end_tokens: vec!["</tool_call>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_false_positive() {
let text = r#"Here it is { Whats up"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<tool_call>".to_string()],
tool_call_end_tokens: vec!["</tool_call>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_with_tool_call_start_token_nemotron_deci() {
let text =
r#"<TOOLCALL>[{"name": "search", "parameters": { "query": "rust" } }]</TOOLCALL>"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<TOOLCALL>".to_string()],
tool_call_end_tokens: vec!["</TOOLCALL>".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_with_lllama3_json_token() {
let text = r#"<|python_tag|>{ "name": }"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["<|python_tag|>".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_mistral_token() {
let text = r#"Hello Yo ! [TOOL_CALLS]{"name": "search", "#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["[TOOL_CALLS]".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_token() {
let text = r#"functools{"name": "search", "#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_partial_token_fun() {
let text = r#"fun"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(
result,
"Should detect 'fun' as potential start of 'functools'"
);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_partial_token_func() {
let text = r#"func"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(
result,
"Should detect 'func' as potential start of 'functools'"
);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_partial_token_f() {
let text = r#"f"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(
result,
"Should detect 'f' as potential start of 'functools'"
);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_partial_with_prefix() {
let text = r#"Hello fun"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(
result,
"Should detect text ending with 'fun' as potential tool call start"
);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_avoid_false_positive() {
let text = r#"funny joke"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(result);
}
#[test] fn detect_tool_call_start_basic_json_chunk_phi4_no_match() {
let text = r#"hello world"#;
let config = JsonParserConfig {
tool_call_start_tokens: vec!["functools".to_string()],
tool_call_end_tokens: vec!["".to_string()],
..Default::default()
};
let result = detect_tool_call_start_basic_json(text, &config);
assert!(
!result,
"Should not detect unrelated text as tool call start"
);
}
}