[misc] replace existing tool call code with new crate package (#17720)
This commit is contained in:
Simo Lin
@@ -1,32 +0,0 @@
|
||||
use thiserror::Error;
|
||||
|
||||
/// Result type for tool parser operations
|
||||
pub type ParserResult<T> = Result<T, ParserError>;
|
||||
|
||||
/// Errors that can occur during tool parsing
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ParserError {
|
||||
#[error("Parsing failed: {0}")]
|
||||
ParsingFailed(String),
|
||||
|
||||
#[error("Model not supported: {0}")]
|
||||
ModelNotSupported(String),
|
||||
|
||||
#[error("Parse depth exceeded: max {0}")]
|
||||
DepthExceeded(usize),
|
||||
|
||||
#[error("Invalid JSON: {0}")]
|
||||
JsonError(#[from] serde_json::Error),
|
||||
|
||||
#[error("Regex error: {0}")]
|
||||
RegexError(#[from] regex::Error),
|
||||
|
||||
#[error("Incomplete tool call")]
|
||||
Incomplete,
|
||||
|
||||
#[error("Invalid tool name: {0}")]
|
||||
InvalidToolName(String),
|
||||
|
||||
#[error("Token not found: {0}")]
|
||||
TokenNotFound(String),
|
||||
}
|
||||
@@ -1,395 +0,0 @@
|
||||
// Factory and pool for creating model-specific tool parsers with pooling support.
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, RwLock},
|
||||
};
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::tool_parser::{
|
||||
parsers::{
|
||||
DeepSeekParser, Glm4MoeParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser,
|
||||
MistralParser, PassthroughParser, PythonicParser, QwenCoderParser, QwenParser, Step3Parser,
|
||||
},
|
||||
traits::ToolParser,
|
||||
};
|
||||
|
||||
/// Type alias for pooled parser instances.
|
||||
pub type PooledParser = Arc<Mutex<Box<dyn ToolParser>>>;
|
||||
|
||||
/// Type alias for parser creator functions.
|
||||
type ParserCreator = Arc<dyn Fn() -> Box<dyn ToolParser> + Send + Sync>;
|
||||
|
||||
/// Registry for model-specific tool parsers with pooling support.
|
||||
#[derive(Clone)]
|
||||
pub struct ParserRegistry {
|
||||
/// Creator functions for parsers (used when pool is empty)
|
||||
creators: Arc<RwLock<HashMap<String, ParserCreator>>>,
|
||||
/// Pooled parser instances for reuse
|
||||
pool: Arc<RwLock<HashMap<String, PooledParser>>>,
|
||||
/// Model pattern to parser name mappings
|
||||
model_mapping: Arc<RwLock<HashMap<String, String>>>,
|
||||
/// Default parser name
|
||||
default_parser: Arc<RwLock<String>>,
|
||||
}
|
||||
|
||||
impl ParserRegistry {
|
||||
/// Create a new empty registry.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
creators: Arc::new(RwLock::new(HashMap::new())),
|
||||
pool: Arc::new(RwLock::new(HashMap::new())),
|
||||
model_mapping: Arc::new(RwLock::new(HashMap::new())),
|
||||
default_parser: Arc::new(RwLock::new("passthrough".to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a parser creator for a given parser type.
|
||||
pub fn register_parser<F>(&self, name: &str, creator: F)
|
||||
where
|
||||
F: Fn() -> Box<dyn ToolParser> + Send + Sync + 'static,
|
||||
{
|
||||
let mut creators = self.creators.write().unwrap();
|
||||
creators.insert(name.to_string(), Arc::new(creator));
|
||||
}
|
||||
|
||||
/// Map a model name/pattern to a parser
|
||||
pub fn map_model(&self, model: impl Into<String>, parser: impl Into<String>) {
|
||||
let mut mapping = self.model_mapping.write().unwrap();
|
||||
mapping.insert(model.into(), parser.into());
|
||||
}
|
||||
|
||||
/// Get a pooled parser by exact name.
|
||||
/// Returns a shared parser instance from the pool, creating one if needed.
|
||||
pub fn get_pooled_parser(&self, name: &str) -> Option<PooledParser> {
|
||||
// First check if we have a pooled instance
|
||||
{
|
||||
let pool = self.pool.read().unwrap();
|
||||
if let Some(parser) = pool.get(name) {
|
||||
return Some(Arc::clone(parser));
|
||||
}
|
||||
}
|
||||
|
||||
// If not in pool, create one and add to pool
|
||||
let creators = self.creators.read().unwrap();
|
||||
if let Some(creator) = creators.get(name) {
|
||||
let parser = Arc::new(Mutex::new(creator()));
|
||||
|
||||
// Add to pool for future use
|
||||
let mut pool = self.pool.write().unwrap();
|
||||
pool.insert(name.to_string(), Arc::clone(&parser));
|
||||
|
||||
Some(parser)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a parser with the given name is registered.
|
||||
pub fn has_parser(&self, name: &str) -> bool {
|
||||
let creators = self.creators.read().unwrap();
|
||||
creators.contains_key(name)
|
||||
}
|
||||
|
||||
/// Create a fresh (non-pooled) parser instance by exact name.
|
||||
/// Returns a new parser instance for each call - useful for streaming where state isolation is needed.
|
||||
pub fn create_parser(&self, name: &str) -> Option<Box<dyn ToolParser>> {
|
||||
let creators = self.creators.read().unwrap();
|
||||
creators.get(name).map(|creator| creator())
|
||||
}
|
||||
|
||||
/// Check if a parser can be created for a specific model without actually creating it.
|
||||
/// Returns true if a parser is available (registered) for this model.
|
||||
pub fn has_parser_for_model(&self, model: &str) -> bool {
|
||||
// Try exact match first
|
||||
{
|
||||
let mapping = self.model_mapping.read().unwrap();
|
||||
if let Some(parser_name) = mapping.get(model) {
|
||||
let creators = self.creators.read().unwrap();
|
||||
if creators.contains_key(parser_name) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try prefix matching
|
||||
let model_mapping = self.model_mapping.read().unwrap();
|
||||
let best_match = model_mapping
|
||||
.iter()
|
||||
.filter(|(pattern, _)| {
|
||||
pattern.ends_with('*') && model.starts_with(&pattern[..pattern.len() - 1])
|
||||
})
|
||||
.max_by_key(|(pattern, _)| pattern.len());
|
||||
|
||||
if let Some((_, parser_name)) = best_match {
|
||||
let creators = self.creators.read().unwrap();
|
||||
if creators.contains_key(parser_name) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// Return false if no specific parser found for this model
|
||||
// (get_pooled will still fall back to default parser)
|
||||
false
|
||||
}
|
||||
|
||||
/// Create a fresh (non-pooled) parser instance for a specific model.
|
||||
/// Returns a new parser instance for each call - useful for streaming where state isolation is needed.
|
||||
pub fn create_for_model(&self, model: &str) -> Option<Box<dyn ToolParser>> {
|
||||
// Try exact match first
|
||||
{
|
||||
let mapping = self.model_mapping.read().unwrap();
|
||||
if let Some(parser_name) = mapping.get(model) {
|
||||
if let Some(parser) = self.create_parser(parser_name) {
|
||||
return Some(parser);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try prefix matching with more specific patterns first
|
||||
let model_mapping = self.model_mapping.read().unwrap();
|
||||
let best_match = model_mapping
|
||||
.iter()
|
||||
.filter(|(pattern, _)| {
|
||||
pattern.ends_with('*') && model.starts_with(&pattern[..pattern.len() - 1])
|
||||
})
|
||||
.max_by_key(|(pattern, _)| pattern.len());
|
||||
|
||||
// Return the best matching parser
|
||||
if let Some((_, parser_name)) = best_match {
|
||||
if let Some(parser) = self.create_parser(parser_name) {
|
||||
return Some(parser);
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to default parser
|
||||
let default = self.default_parser.read().unwrap().clone();
|
||||
self.create_parser(&default)
|
||||
}
|
||||
|
||||
/// Get parser for a specific model
|
||||
pub fn get_pooled_for_model(&self, model: &str) -> Option<PooledParser> {
|
||||
// Try exact match first
|
||||
{
|
||||
let mapping = self.model_mapping.read().unwrap();
|
||||
if let Some(parser_name) = mapping.get(model) {
|
||||
if let Some(parser) = self.get_pooled_parser(parser_name) {
|
||||
return Some(parser);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try prefix matching with more specific patterns first
|
||||
let model_mapping = self.model_mapping.read().unwrap();
|
||||
let best_match = model_mapping
|
||||
.iter()
|
||||
.filter(|(pattern, _)| {
|
||||
pattern.ends_with('*') && model.starts_with(&pattern[..pattern.len() - 1])
|
||||
})
|
||||
.max_by_key(|(pattern, _)| pattern.len());
|
||||
|
||||
// Return the best matching parser
|
||||
if let Some((_, parser_name)) = best_match {
|
||||
if let Some(parser) = self.get_pooled_parser(parser_name) {
|
||||
return Some(parser);
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to default parser
|
||||
let default = self.default_parser.read().unwrap().clone();
|
||||
self.get_pooled_parser(&default)
|
||||
}
|
||||
|
||||
/// Clear the parser pool, forcing new instances to be created.
|
||||
pub fn clear_pool(&self) {
|
||||
let mut pool = self.pool.write().unwrap();
|
||||
pool.clear();
|
||||
}
|
||||
|
||||
/// Set the default parser
|
||||
pub fn set_default_parser(&self, name: impl Into<String>) {
|
||||
let mut default = self.default_parser.write().unwrap();
|
||||
*default = name.into();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ParserRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Factory for creating tool parsers based on model type.
|
||||
#[derive(Clone)]
|
||||
pub struct ParserFactory {
|
||||
registry: ParserRegistry,
|
||||
}
|
||||
|
||||
impl ParserFactory {
|
||||
/// Create a new factory with default parsers registered.
|
||||
pub fn new() -> Self {
|
||||
let registry = ParserRegistry::new();
|
||||
|
||||
// Register default parsers
|
||||
registry.register_parser("passthrough", || Box::new(PassthroughParser::new()));
|
||||
registry.register_parser("json", || Box::new(JsonParser::new()));
|
||||
registry.register_parser("mistral", || Box::new(MistralParser::new()));
|
||||
registry.register_parser("qwen", || Box::new(QwenParser::new()));
|
||||
registry.register_parser("qwen_coder", || Box::new(QwenCoderParser::new()));
|
||||
registry.register_parser("pythonic", || Box::new(PythonicParser::new()));
|
||||
registry.register_parser("llama", || Box::new(LlamaParser::new()));
|
||||
registry.register_parser("deepseek", || Box::new(DeepSeekParser::new()));
|
||||
registry.register_parser("glm45_moe", || Box::new(Glm4MoeParser::glm45()));
|
||||
registry.register_parser("glm47_moe", || Box::new(Glm4MoeParser::glm47()));
|
||||
registry.register_parser("step3", || Box::new(Step3Parser::new()));
|
||||
registry.register_parser("kimik2", || Box::new(KimiK2Parser::new()));
|
||||
registry.register_parser("minimax_m2", || Box::new(MinimaxM2Parser::new()));
|
||||
|
||||
// Register default model mappings
|
||||
Self::register_default_mappings(®istry);
|
||||
|
||||
Self { registry }
|
||||
}
|
||||
|
||||
fn register_default_mappings(registry: &ParserRegistry) {
|
||||
// OpenAI models
|
||||
registry.map_model("gpt-4*", "json");
|
||||
registry.map_model("gpt-3.5*", "json");
|
||||
registry.map_model("gpt-4o*", "json");
|
||||
|
||||
// Anthropic models
|
||||
registry.map_model("claude-*", "json");
|
||||
|
||||
// Mistral models
|
||||
registry.map_model("mistral-*", "mistral");
|
||||
registry.map_model("mixtral-*", "mistral");
|
||||
|
||||
// Qwen models (more specific patterns first - longer patterns take precedence)
|
||||
// Qwen Coder models use XML format: <tool_call><function=name><parameter=key>value</parameter></function></tool_call>
|
||||
registry.map_model("Qwen/Qwen3-Coder*", "qwen_coder");
|
||||
registry.map_model("Qwen3-Coder*", "qwen_coder");
|
||||
registry.map_model("qwen3-coder*", "qwen_coder");
|
||||
registry.map_model("Qwen/Qwen2.5-Coder*", "qwen_coder");
|
||||
registry.map_model("Qwen2.5-Coder*", "qwen_coder");
|
||||
registry.map_model("qwen2.5-coder*", "qwen_coder");
|
||||
// Generic Qwen models use JSON format
|
||||
registry.map_model("qwen*", "qwen");
|
||||
registry.map_model("Qwen*", "qwen");
|
||||
|
||||
// Llama models
|
||||
registry.map_model("llama-4*", "pythonic");
|
||||
registry.map_model("meta-llama-4*", "pythonic");
|
||||
registry.map_model("llama-3.2*", "llama");
|
||||
registry.map_model("meta-llama-3.2*", "llama");
|
||||
registry.map_model("llama-*", "json");
|
||||
registry.map_model("meta-llama-*", "json");
|
||||
|
||||
// DeepSeek models
|
||||
registry.map_model("deepseek-v3*", "deepseek");
|
||||
registry.map_model("deepseek-ai/DeepSeek-V3*", "deepseek");
|
||||
registry.map_model("deepseek-*", "pythonic");
|
||||
|
||||
// GLM models
|
||||
registry.map_model("glm-4.5*", "glm45_moe");
|
||||
registry.map_model("glm-4.6*", "glm45_moe");
|
||||
registry.map_model("glm-4.7*", "glm47_moe");
|
||||
registry.map_model("glm-*", "json");
|
||||
|
||||
// Step3 models
|
||||
registry.map_model("step3*", "step3");
|
||||
registry.map_model("Step-3*", "step3");
|
||||
|
||||
// Kimi models
|
||||
registry.map_model("kimi-k2*", "kimik2");
|
||||
registry.map_model("Kimi-K2*", "kimik2");
|
||||
registry.map_model("moonshot*/Kimi-K2*", "kimik2");
|
||||
|
||||
// MiniMax models
|
||||
registry.map_model("minimax*", "minimax_m2");
|
||||
registry.map_model("MiniMax*", "minimax_m2");
|
||||
|
||||
// Other models
|
||||
registry.map_model("gemini-*", "json");
|
||||
registry.map_model("palm-*", "json");
|
||||
registry.map_model("gemma-*", "json");
|
||||
}
|
||||
|
||||
/// Get a pooled parser for the given model ID.
|
||||
/// Returns a shared instance that can be used concurrently.
|
||||
/// Falls back to passthrough parser if model is not recognized.
|
||||
pub fn get_pooled(&self, model_id: &str) -> PooledParser {
|
||||
self.registry
|
||||
.get_pooled_for_model(model_id)
|
||||
.unwrap_or_else(|| {
|
||||
// Fallback to passthrough parser (no-op, returns text unchanged)
|
||||
self.registry
|
||||
.get_pooled_parser("passthrough")
|
||||
.expect("Passthrough parser should always be registered")
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the internal registry for custom registration.
|
||||
pub fn registry(&self) -> &ParserRegistry {
|
||||
&self.registry
|
||||
}
|
||||
|
||||
/// Clear the parser pool.
|
||||
pub fn clear_pool(&self) {
|
||||
self.registry.clear_pool();
|
||||
}
|
||||
|
||||
/// Get a non-pooled parser for the given model ID (creates a fresh instance each time).
|
||||
/// This is useful for benchmarks and testing where you want independent parser instances.
|
||||
pub fn get_parser(&self, model_id: &str) -> Option<Arc<dyn ToolParser>> {
|
||||
// Determine which parser type to use
|
||||
let parser_type = {
|
||||
let mapping = self.registry.model_mapping.read().unwrap();
|
||||
|
||||
// Try exact match first
|
||||
if let Some(parser_name) = mapping.get(model_id) {
|
||||
parser_name.clone()
|
||||
} else {
|
||||
// Try prefix matching
|
||||
let best_match = mapping
|
||||
.iter()
|
||||
.filter(|(pattern, _)| {
|
||||
pattern.ends_with('*')
|
||||
&& model_id.starts_with(&pattern[..pattern.len() - 1])
|
||||
})
|
||||
.max_by_key(|(pattern, _)| pattern.len());
|
||||
|
||||
if let Some((_, parser_name)) = best_match {
|
||||
parser_name.clone()
|
||||
} else {
|
||||
// Fall back to default
|
||||
self.registry.default_parser.read().unwrap().clone()
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let creators = self.registry.creators.read().unwrap();
|
||||
creators.get(&parser_type).map(|creator| {
|
||||
// Call the creator to get a Box<dyn ToolParser>, then convert to Arc
|
||||
let boxed_parser = creator();
|
||||
Arc::from(boxed_parser)
|
||||
})
|
||||
}
|
||||
|
||||
/// List all registered parsers (for compatibility with old API).
|
||||
pub fn list_parsers(&self) -> Vec<String> {
|
||||
self.registry
|
||||
.creators
|
||||
.read()
|
||||
.unwrap()
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ParserFactory {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -1,24 +0,0 @@
|
||||
/// Tool parser module for handling function/tool calls in model outputs
|
||||
///
|
||||
/// This module provides infrastructure for parsing tool calls from various model formats.
|
||||
// Core modules
|
||||
pub mod errors;
|
||||
pub mod factory;
|
||||
pub mod partial_json;
|
||||
pub mod traits;
|
||||
pub mod types;
|
||||
|
||||
// Parser implementations
|
||||
pub mod parsers;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
// Re-export types used outside this module
|
||||
pub use factory::{ParserFactory, PooledParser};
|
||||
pub use parsers::{
|
||||
DeepSeekParser, Glm4MoeParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser,
|
||||
MistralParser, PythonicParser, QwenParser, Step3Parser,
|
||||
};
|
||||
pub use traits::ToolParser;
|
||||
pub use types::{FunctionCall, PartialToolCall, StreamingParseResult, ToolCall};
|
||||
@@ -1,328 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// DeepSeek V3 format parser for tool calls
|
||||
///
|
||||
/// Handles the DeepSeek V3 specific format that uses Unicode tokens:
|
||||
/// `<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>{name}\n```json\n{args}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Unicode token delimiters
|
||||
/// - JSON arguments in code blocks
|
||||
/// - Support for multiple sequential tool calls
|
||||
///
|
||||
/// Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3-0324?chat_template=default
|
||||
pub struct DeepSeekParser {
|
||||
/// Regex for extracting complete tool calls
|
||||
tool_call_extractor: Regex,
|
||||
/// Regex for extracting function details
|
||||
func_detail_extractor: Regex,
|
||||
/// Regex for matching partial tool calls during streaming
|
||||
partial_tool_call_regex: Regex,
|
||||
/// Regex pattern for removing completed tool calls from buffer
|
||||
tool_call_end_pattern: Regex,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
}
|
||||
|
||||
impl DeepSeekParser {
|
||||
/// Create a new DeepSeek parser
|
||||
pub fn new() -> Self {
|
||||
// Use (?s) flag for DOTALL mode to handle newlines
|
||||
let tool_call_pattern = r"(?s)<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>";
|
||||
let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
|
||||
|
||||
let func_detail_pattern = r"(?s)<｜tool▁call▁begin｜>(.*?)<｜tool▁sep｜>(.*?)\n```json\n(.*?)\n```<｜tool▁call▁end｜>";
|
||||
let func_detail_extractor = Regex::new(func_detail_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Partial pattern for streaming - uses .* (greedy) not .*? to match all partial content
|
||||
let partial_pattern = r"(?s)<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)";
|
||||
let partial_tool_call_regex = Regex::new(partial_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Pattern for removing completed tool calls
|
||||
let end_pattern = r"(?s)<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>";
|
||||
let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
tool_call_extractor,
|
||||
func_detail_extractor,
|
||||
partial_tool_call_regex,
|
||||
tool_call_end_pattern,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a single tool call block - throws error if parsing fails
|
||||
fn parse_tool_call(&self, block: &str) -> ParserResult<ToolCall> {
|
||||
let captures = self.func_detail_extractor.captures(block).ok_or_else(|| {
|
||||
ParserError::ParsingFailed("Failed to match tool call pattern".to_string())
|
||||
})?;
|
||||
|
||||
// Get function type (should be "function")
|
||||
let func_type = captures.get(1).map_or("", |m| m.as_str());
|
||||
if func_type != "function" {
|
||||
return Err(ParserError::ParsingFailed(format!(
|
||||
"Invalid function type: {}",
|
||||
func_type
|
||||
)));
|
||||
}
|
||||
|
||||
// Get function name
|
||||
let func_name = captures.get(2).map_or("", |m| m.as_str()).trim();
|
||||
if func_name.is_empty() {
|
||||
return Err(ParserError::ParsingFailed(
|
||||
"Empty function name".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Get JSON arguments
|
||||
let json_args = captures.get(3).map_or("{}", |m| m.as_str()).trim();
|
||||
|
||||
// Parse JSON arguments
|
||||
let value = serde_json::from_str::<Value>(json_args)
|
||||
.map_err(|e| ParserError::ParsingFailed(format!("Invalid JSON: {}", e)))?;
|
||||
|
||||
// Create arguments object
|
||||
let args = if value.is_object() {
|
||||
value
|
||||
} else {
|
||||
// If not an object, wrap it
|
||||
serde_json::json!({ "value": value })
|
||||
};
|
||||
|
||||
let arguments =
|
||||
serde_json::to_string(&args).map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: func_name.to_string(),
|
||||
arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for DeepSeekParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for DeepSeekParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where tool calls begin
|
||||
let idx = text.find("<｜tool▁calls▁begin｜>").unwrap();
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Try to extract tool calls, log warnings for failures
|
||||
let mut tools = Vec::new();
|
||||
for mat in self.tool_call_extractor.find_iter(text) {
|
||||
match self.parse_tool_call(mat.as_str()) {
|
||||
Ok(tool) => tools.push(tool),
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if we have a tool call (either the start token or individual tool call)
|
||||
let has_tool_call =
|
||||
self.has_tool_markers(current_text) || current_text.contains("<｜tool▁call▁begin｜>");
|
||||
|
||||
if !has_tool_call {
|
||||
// No tool markers detected - return all buffered content as normal text
|
||||
// Strip out end tokens if present
|
||||
let mut normal_text = std::mem::take(&mut self.buffer);
|
||||
for e_token in ["<｜tool▁calls▁end｜>", "```", "<｜tool▁call▁end｜>"] {
|
||||
normal_text = normal_text.replace(e_token, "");
|
||||
}
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
let mut calls: Vec<ToolCallItem> = Vec::new();
|
||||
|
||||
// Try to match the partial tool call pattern
|
||||
if let Some(captures) = self.partial_tool_call_regex.captures(current_text) {
|
||||
let func_name = captures.get(2).map_or("", |m| m.as_str()).trim();
|
||||
let func_args_raw = captures.get(3).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Validate tool name
|
||||
if !tool_indices.contains_key(func_name) {
|
||||
// Invalid tool name - skip this tool, preserve indexing for next tool
|
||||
tracing::debug!("Invalid tool name '{}' - skipping", func_name);
|
||||
helpers::reset_current_tool_state(
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&self.prev_tool_call_arr,
|
||||
);
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
|
||||
// Initialize state if this is the first tool call
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
self.prev_tool_call_arr = Vec::new();
|
||||
self.streamed_args_for_tool = vec![String::new()];
|
||||
}
|
||||
|
||||
// Ensure we have enough entries in our tracking arrays
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Send tool name if not sent yet
|
||||
if !self.current_tool_name_sent {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: Some(func_name.to_string()),
|
||||
parameters: String::new(),
|
||||
});
|
||||
self.current_tool_name_sent = true;
|
||||
|
||||
// Store the tool call info for serving layer completions endpoint
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if self.prev_tool_call_arr.len() <= tool_id {
|
||||
self.prev_tool_call_arr
|
||||
.resize_with(tool_id + 1, || Value::Null);
|
||||
}
|
||||
self.prev_tool_call_arr[tool_id] = serde_json::json!({
|
||||
"name": func_name,
|
||||
"arguments": {},
|
||||
});
|
||||
} else {
|
||||
// Compute incremental diff
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
let last_sent = self
|
||||
.streamed_args_for_tool
|
||||
.get(tool_id)
|
||||
.map(|s| s.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let argument_diff = func_args_raw
|
||||
.strip_prefix(last_sent)
|
||||
.unwrap_or(func_args_raw);
|
||||
|
||||
if !argument_diff.is_empty() {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: None,
|
||||
parameters: argument_diff.to_string(),
|
||||
});
|
||||
if tool_id < self.streamed_args_for_tool.len() {
|
||||
self.streamed_args_for_tool[tool_id].push_str(argument_diff);
|
||||
}
|
||||
}
|
||||
|
||||
// Check if JSON is complete
|
||||
if helpers::is_complete_json(func_args_raw) {
|
||||
// Update the stored arguments
|
||||
if let Ok(parsed_args) = serde_json::from_str::<Value>(func_args_raw) {
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.prev_tool_call_arr.len() {
|
||||
if let Some(obj) = self.prev_tool_call_arr[tool_id].as_object_mut() {
|
||||
obj.insert("arguments".to_string(), parsed_args);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Find the end of the current tool call and remove only that part from buffer
|
||||
if let Some(mat) = self.tool_call_end_pattern.find(current_text) {
|
||||
// Remove the completed tool call from buffer, keep any remaining content
|
||||
self.buffer = current_text[mat.end()..].to_string();
|
||||
} else {
|
||||
self.buffer.clear();
|
||||
}
|
||||
|
||||
let result = StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
};
|
||||
|
||||
self.current_tool_id += 1;
|
||||
self.current_tool_name_sent = false;
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
})
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains("<｜tool▁calls▁begin｜>")
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.prev_tool_call_arr.clear();
|
||||
self.current_tool_id = -1;
|
||||
self.current_tool_name_sent = false;
|
||||
self.streamed_args_for_tool.clear();
|
||||
}
|
||||
}
|
||||
@@ -1,340 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// GLM-4 MoE format parser for tool calls
|
||||
///
|
||||
/// Handles both GLM-4 MoE and GLM-4.7 MoE formats:
|
||||
/// - GLM-4: `<tool_call>{name}\n<arg_key>{key}</arg_key>\n<arg_value>{value}</arg_value>\n</tool_call>`
|
||||
/// - GLM-4.7: `<tool_call>{name}<arg_key>{key}</arg_key><arg_value>{value}</arg_value></tool_call>`
|
||||
///
|
||||
/// Features:
|
||||
/// - XML-style tags for tool calls
|
||||
/// - Key-value pairs for arguments
|
||||
/// - Support for multiple sequential tool calls
|
||||
pub struct Glm4MoeParser {
|
||||
/// Regex for extracting complete tool calls
|
||||
tool_call_extractor: Regex,
|
||||
/// Regex for extracting function details
|
||||
func_detail_extractor: Regex,
|
||||
/// Regex for extracting argument key-value pairs
|
||||
arg_extractor: Regex,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Token configuration
|
||||
bot_token: &'static str,
|
||||
eot_token: &'static str,
|
||||
}
|
||||
|
||||
impl Glm4MoeParser {
|
||||
/// Create a new generic GLM MoE parser with a custom func_detail_extractor pattern
|
||||
///
|
||||
/// # Arguments
|
||||
/// - `func_detail_pattern`: Regex pattern for extracting function name and arguments
|
||||
/// - For GLM-4: `r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>"`
|
||||
/// - For GLM-4.7: `r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>"`
|
||||
pub(crate) fn new(func_detail_pattern: &str) -> Self {
|
||||
// Use (?s) flag for DOTALL mode to handle newlines
|
||||
let tool_call_pattern = r"(?s)<tool_call>.*?</tool_call>";
|
||||
let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
|
||||
|
||||
let func_detail_extractor = Regex::new(func_detail_pattern).expect("Valid regex pattern");
|
||||
|
||||
let arg_pattern = r"(?s)<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>";
|
||||
let arg_extractor = Regex::new(arg_pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
tool_call_extractor,
|
||||
func_detail_extractor,
|
||||
arg_extractor,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
bot_token: "<tool_call>",
|
||||
eot_token: "</tool_call>",
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new GLM-4.5/4.6 MoE parser (with newline-based format)
|
||||
pub fn glm45() -> Self {
|
||||
Self::new(r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>")
|
||||
}
|
||||
|
||||
/// Create a new GLM-4.7 MoE parser (with whitespace-based format)
|
||||
pub fn glm47() -> Self {
|
||||
Self::new(r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>")
|
||||
}
|
||||
|
||||
/// Parse arguments from key-value pairs
|
||||
fn parse_arguments(&self, args_text: &str) -> ParserResult<serde_json::Map<String, Value>> {
|
||||
let mut arguments = serde_json::Map::new();
|
||||
|
||||
for capture in self.arg_extractor.captures_iter(args_text) {
|
||||
let key = capture.get(1).map_or("", |m| m.as_str()).trim();
|
||||
let value_str = capture.get(2).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Try to parse the value as JSON first, fallback to string
|
||||
let value = if let Ok(json_val) = serde_json::from_str::<Value>(value_str) {
|
||||
json_val
|
||||
} else {
|
||||
// Try parsing as Python literal (similar to Python's ast.literal_eval)
|
||||
if value_str == "true" || value_str == "True" {
|
||||
Value::Bool(true)
|
||||
} else if value_str == "false" || value_str == "False" {
|
||||
Value::Bool(false)
|
||||
} else if value_str == "null" || value_str == "None" {
|
||||
Value::Null
|
||||
} else if let Ok(num) = value_str.parse::<i64>() {
|
||||
Value::Number(num.into())
|
||||
} else if let Ok(num) = value_str.parse::<f64>() {
|
||||
if let Some(n) = serde_json::Number::from_f64(num) {
|
||||
Value::Number(n)
|
||||
} else {
|
||||
Value::String(value_str.to_string())
|
||||
}
|
||||
} else {
|
||||
Value::String(value_str.to_string())
|
||||
}
|
||||
};
|
||||
|
||||
arguments.insert(key.to_string(), value);
|
||||
}
|
||||
|
||||
Ok(arguments)
|
||||
}
|
||||
|
||||
/// Parse a single tool call block
|
||||
fn parse_tool_call(&self, block: &str) -> ParserResult<Option<ToolCall>> {
|
||||
if let Some(captures) = self.func_detail_extractor.captures(block) {
|
||||
// Get function name
|
||||
let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Get arguments text
|
||||
let args_text = captures.get(2).map_or("", |m| m.as_str());
|
||||
|
||||
// Parse arguments
|
||||
let arguments = self.parse_arguments(args_text)?;
|
||||
|
||||
let arguments_str = serde_json::to_string(&arguments)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: func_name.to_string(),
|
||||
arguments: arguments_str,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse all tool calls from text (shared logic for complete and incremental parsing)
|
||||
fn parse_tool_calls_from_text(&self, text: &str) -> ParserResult<Vec<ToolCall>> {
|
||||
let mut tools = Vec::new();
|
||||
|
||||
for mat in self.tool_call_extractor.find_iter(text) {
|
||||
match self.parse_tool_call(mat.as_str()) {
|
||||
Ok(Some(tool)) => tools.push(tool),
|
||||
Ok(None) => continue,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(tools)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Glm4MoeParser {
|
||||
fn default() -> Self {
|
||||
Self::glm45()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for Glm4MoeParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Check if text contains GLM-4 MoE format
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where tool calls begin
|
||||
let idx = text.find("<tool_call>").unwrap();
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Parse all tool calls using shared helper
|
||||
let tools = self.parse_tool_calls_from_text(text)?;
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Python logic: Wait for complete tool call, then parse it all at once
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if we have bot_token
|
||||
let start = current_text.find(self.bot_token);
|
||||
if start.is_none() {
|
||||
self.buffer.clear();
|
||||
// If we're in the middle of streaming (current_tool_id > 0), don't return text
|
||||
let normal_text = if self.current_tool_id > 0 {
|
||||
String::new()
|
||||
} else {
|
||||
current_text.clone()
|
||||
};
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Check if we have eot_token (end of tool call)
|
||||
let end = current_text.find(self.eot_token);
|
||||
if let Some(end_pos) = end {
|
||||
// We have a complete tool call!
|
||||
|
||||
// Initialize state if this is the first tool call
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
self.prev_tool_call_arr = Vec::new();
|
||||
self.streamed_args_for_tool = vec![String::new()];
|
||||
}
|
||||
|
||||
// Ensure we have enough entries in our tracking arrays
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Parse the complete block using shared helper
|
||||
let block_end = end_pos + self.eot_token.len();
|
||||
let parsed_tools = self.parse_tool_calls_from_text(¤t_text[..block_end])?;
|
||||
|
||||
// Extract normal text before tool calls
|
||||
let idx = current_text.find(self.bot_token);
|
||||
let normal_text = if let Some(pos) = idx {
|
||||
current_text[..pos].trim().to_string()
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
let mut calls = Vec::new();
|
||||
|
||||
if !parsed_tools.is_empty() {
|
||||
// Take the first tool and convert to ToolCallItem
|
||||
let tool_call = &parsed_tools[0];
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
|
||||
// Validate tool name
|
||||
if !tool_indices.contains_key(&tool_call.function.name) {
|
||||
// Invalid tool name - skip this tool, preserve indexing for next tool
|
||||
tracing::debug!("Invalid tool name '{}' - skipping", tool_call.function.name);
|
||||
helpers::reset_current_tool_state(
|
||||
&mut self.buffer,
|
||||
&mut false, // glm45_moe/glm47_moe doesn't track name_sent per tool
|
||||
&mut self.streamed_args_for_tool,
|
||||
&self.prev_tool_call_arr,
|
||||
);
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: Some(tool_call.function.name.clone()),
|
||||
parameters: tool_call.function.arguments.clone(),
|
||||
});
|
||||
|
||||
// Store in tracking arrays
|
||||
if self.prev_tool_call_arr.len() <= tool_id {
|
||||
self.prev_tool_call_arr
|
||||
.resize_with(tool_id + 1, || Value::Null);
|
||||
}
|
||||
|
||||
// Parse parameters as JSON and store
|
||||
if let Ok(args) = serde_json::from_str::<Value>(&tool_call.function.arguments) {
|
||||
self.prev_tool_call_arr[tool_id] = serde_json::json!({
|
||||
"name": tool_call.function.name,
|
||||
"arguments": args,
|
||||
});
|
||||
}
|
||||
|
||||
if self.streamed_args_for_tool.len() <= tool_id {
|
||||
self.streamed_args_for_tool
|
||||
.resize_with(tool_id + 1, String::new);
|
||||
}
|
||||
self.streamed_args_for_tool[tool_id] = tool_call.function.arguments.clone();
|
||||
|
||||
self.current_tool_id += 1;
|
||||
}
|
||||
|
||||
// Remove processed portion from buffer
|
||||
self.buffer = current_text[block_end..].to_string();
|
||||
return Ok(StreamingParseResult { normal_text, calls });
|
||||
}
|
||||
|
||||
// No complete tool call yet - return normal text before start token
|
||||
let start_pos = start.unwrap();
|
||||
let normal_text = current_text[..start_pos].to_string();
|
||||
self.buffer = current_text[start_pos..].to_string();
|
||||
|
||||
Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains(self.bot_token)
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.prev_tool_call_arr.clear();
|
||||
self.current_tool_id = -1;
|
||||
self.streamed_args_for_tool.clear();
|
||||
}
|
||||
}
|
||||
@@ -1,488 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
types::{StreamingParseResult, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// Get a mapping of tool names to their indices
|
||||
pub fn get_tool_indices(tools: &[Tool]) -> HashMap<String, usize> {
|
||||
tools
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, tool)| (tool.function.name.clone(), i))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Find the common prefix of two strings
|
||||
/// Used for incremental argument streaming when partial JSON returns different intermediate states
|
||||
pub fn find_common_prefix(s1: &str, s2: &str) -> String {
|
||||
s1.chars()
|
||||
.zip(s2.chars())
|
||||
.take_while(|(c1, c2)| c1 == c2)
|
||||
.map(|(c1, _)| c1)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get unstreamed tool call arguments
|
||||
/// Returns tool call items for arguments that have been parsed but not yet streamed
|
||||
/// This ensures tool calls are properly completed even if the model generates final arguments in the last chunk
|
||||
pub fn get_unstreamed_args(
|
||||
prev_tool_call_arr: &[Value],
|
||||
streamed_args_for_tool: &[String],
|
||||
) -> Option<Vec<ToolCallItem>> {
|
||||
// Check if we have tool calls being tracked
|
||||
if prev_tool_call_arr.is_empty() || streamed_args_for_tool.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Get the last tool call that was being processed
|
||||
let tool_index = prev_tool_call_arr.len() - 1;
|
||||
if tool_index >= streamed_args_for_tool.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Get expected vs actual arguments
|
||||
let expected_args = prev_tool_call_arr[tool_index].get("arguments")?;
|
||||
let expected_str = serde_json::to_string(expected_args).ok()?;
|
||||
let actual_str = &streamed_args_for_tool[tool_index];
|
||||
|
||||
// Check if there are remaining arguments to send
|
||||
let remaining = if expected_str.starts_with(actual_str) {
|
||||
&expected_str[actual_str.len()..]
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
|
||||
if remaining.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Return the remaining arguments as a ToolCallItem
|
||||
Some(vec![ToolCallItem {
|
||||
tool_index,
|
||||
name: None, // No name for argument deltas
|
||||
parameters: remaining.to_string(),
|
||||
}])
|
||||
}
|
||||
|
||||
/// Check if a buffer ends with a partial occurrence of a token
|
||||
/// Returns Some(length) if there's a partial match, None otherwise
|
||||
pub fn ends_with_partial_token(buffer: &str, token: &str) -> Option<usize> {
|
||||
if buffer.is_empty() || token.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
(1..token.len()).find(|&i| buffer.ends_with(&token[..i]))
|
||||
}
|
||||
|
||||
/// Reset state for the current tool being parsed (used when skipping invalid tools).
|
||||
/// This preserves the parser's overall state (current_tool_id, prev_tool_call_arr)
|
||||
/// but clears the state specific to the current incomplete tool.
|
||||
pub fn reset_current_tool_state(
|
||||
buffer: &mut String,
|
||||
current_tool_name_sent: &mut bool,
|
||||
streamed_args_for_tool: &mut Vec<String>,
|
||||
prev_tool_call_arr: &[Value],
|
||||
) {
|
||||
buffer.clear();
|
||||
*current_tool_name_sent = false;
|
||||
|
||||
// Only pop if we added an entry for the current (invalid) tool
|
||||
// streamed_args_for_tool should match prev_tool_call_arr length for completed tools
|
||||
if streamed_args_for_tool.len() > prev_tool_call_arr.len() {
|
||||
streamed_args_for_tool.pop();
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the entire parser state (used at the start of a new request).
|
||||
/// Clears all accumulated tool calls and resets all state to initial values.
|
||||
pub fn reset_parser_state(
|
||||
buffer: &mut String,
|
||||
prev_tool_call_arr: &mut Vec<Value>,
|
||||
current_tool_id: &mut i32,
|
||||
current_tool_name_sent: &mut bool,
|
||||
streamed_args_for_tool: &mut Vec<String>,
|
||||
) {
|
||||
buffer.clear();
|
||||
prev_tool_call_arr.clear();
|
||||
*current_tool_id = -1;
|
||||
*current_tool_name_sent = false;
|
||||
streamed_args_for_tool.clear();
|
||||
}
|
||||
|
||||
/// Ensure arrays have capacity for the given tool ID
|
||||
pub fn ensure_capacity(
|
||||
current_tool_id: i32,
|
||||
prev_tool_call_arr: &mut Vec<Value>,
|
||||
streamed_args_for_tool: &mut Vec<String>,
|
||||
) {
|
||||
if current_tool_id < 0 {
|
||||
return;
|
||||
}
|
||||
let needed = (current_tool_id + 1) as usize;
|
||||
|
||||
if prev_tool_call_arr.len() < needed {
|
||||
prev_tool_call_arr.resize_with(needed, || Value::Null);
|
||||
}
|
||||
if streamed_args_for_tool.len() < needed {
|
||||
streamed_args_for_tool.resize_with(needed, String::new);
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a string contains complete, valid JSON
|
||||
pub fn is_complete_json(input: &str) -> bool {
|
||||
serde_json::from_str::<Value>(input).is_ok()
|
||||
}
|
||||
|
||||
/// Normalize the arguments/parameters field in a tool call object.
|
||||
/// If the object has "parameters" but not "arguments", copy parameters to arguments.
|
||||
///
|
||||
/// # Background
|
||||
/// Different LLM formats use different field names:
|
||||
/// - Llama and JSON parsers use "parameters" (correct per JSON Schema spec)
|
||||
/// - Mistral and Qwen use "arguments"
|
||||
///
|
||||
/// This function normalizes to "arguments" for consistent downstream processing.
|
||||
pub fn normalize_arguments_field(mut obj: Value) -> Value {
|
||||
if obj.get("arguments").is_none() {
|
||||
if let Some(params) = obj.get("parameters").cloned() {
|
||||
if let Value::Object(ref mut map) = obj {
|
||||
map.insert("arguments".to_string(), params);
|
||||
}
|
||||
}
|
||||
}
|
||||
obj
|
||||
}
|
||||
|
||||
/// Handle the entire JSON tool call streaming process for JSON-based parsers.
|
||||
///
|
||||
/// This unified function handles all aspects of streaming tool calls:
|
||||
/// - Parsing partial JSON from the buffer
|
||||
/// - Validating tool names against available tools
|
||||
/// - Streaming tool names (Case 1)
|
||||
/// - Streaming tool arguments (Case 2)
|
||||
/// - Managing parser state and buffer updates
|
||||
///
|
||||
/// Used by JSON, Llama, Mistral, and Qwen parsers.
|
||||
///
|
||||
/// # Parameters
|
||||
/// - `current_text`: The current buffered text being parsed
|
||||
/// - `start_idx`: Start index of JSON content in current_text
|
||||
/// - `partial_json`: Mutable reference to partial JSON parser
|
||||
/// - `tool_indices`: Map of valid tool names to their indices
|
||||
/// - `buffer`: Mutable parser buffer
|
||||
/// - `current_tool_id`: Mutable current tool index (-1 means no active tool)
|
||||
/// - `current_tool_name_sent`: Mutable flag for whether current tool's name was sent
|
||||
/// - `streamed_args_for_tool`: Mutable accumulator of streamed arguments per tool
|
||||
/// - `prev_tool_call_arr`: Mutable array of previous tool call states
|
||||
///
|
||||
/// # Returns
|
||||
/// - `Ok(StreamingParseResult)` with any tool call items to stream
|
||||
/// - `Err(ParserError)` if JSON parsing or serialization fails
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn handle_json_tool_streaming(
|
||||
current_text: &str,
|
||||
start_idx: usize,
|
||||
partial_json: &mut crate::tool_parser::partial_json::PartialJson,
|
||||
tool_indices: &HashMap<String, usize>,
|
||||
buffer: &mut String,
|
||||
current_tool_id: &mut i32,
|
||||
current_tool_name_sent: &mut bool,
|
||||
streamed_args_for_tool: &mut Vec<String>,
|
||||
prev_tool_call_arr: &mut Vec<Value>,
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Check if we have content to parse
|
||||
if start_idx >= current_text.len() {
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
|
||||
// Extract JSON string from current position
|
||||
let json_str = ¤t_text[start_idx..];
|
||||
|
||||
// When current_tool_name_sent is false, don't allow partial strings to avoid
|
||||
// parsing incomplete tool names as empty strings
|
||||
let allow_partial_strings = *current_tool_name_sent;
|
||||
|
||||
// Parse partial JSON
|
||||
let (obj, end_idx) = match partial_json.parse_value(json_str, allow_partial_strings) {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
};
|
||||
|
||||
// Check if JSON is complete - validate only the parsed portion
|
||||
// Ensure end_idx is on a valid UTF-8 character boundary
|
||||
let safe_end_idx = if json_str.is_char_boundary(end_idx) {
|
||||
end_idx
|
||||
} else {
|
||||
// Find the nearest valid character boundary before end_idx
|
||||
(0..end_idx)
|
||||
.rev()
|
||||
.find(|&i| json_str.is_char_boundary(i))
|
||||
.unwrap_or(0)
|
||||
};
|
||||
let is_complete = serde_json::from_str::<Value>(&json_str[..safe_end_idx]).is_ok();
|
||||
|
||||
// Validate tool name if present
|
||||
if let Some(name) = obj.get("name").and_then(|v| v.as_str()) {
|
||||
if !tool_indices.contains_key(name) {
|
||||
// Invalid tool name - skip this tool, preserve indexing for next tool
|
||||
tracing::debug!("Invalid tool name '{}' - skipping", name);
|
||||
reset_current_tool_state(
|
||||
buffer,
|
||||
current_tool_name_sent,
|
||||
streamed_args_for_tool,
|
||||
prev_tool_call_arr,
|
||||
);
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize parameters/arguments field
|
||||
let current_tool_call = normalize_arguments_field(obj);
|
||||
|
||||
let mut result = StreamingParseResult::default();
|
||||
|
||||
// Case 1: Handle tool name streaming
|
||||
if !*current_tool_name_sent {
|
||||
if let Some(function_name) = current_tool_call.get("name").and_then(|v| v.as_str()) {
|
||||
if tool_indices.contains_key(function_name) {
|
||||
// Initialize if first tool
|
||||
if *current_tool_id == -1 {
|
||||
*current_tool_id = 0;
|
||||
streamed_args_for_tool.push(String::new());
|
||||
} else if *current_tool_id as usize >= streamed_args_for_tool.len() {
|
||||
// Ensure capacity for subsequent tools
|
||||
ensure_capacity(*current_tool_id, prev_tool_call_arr, streamed_args_for_tool);
|
||||
}
|
||||
|
||||
// Send tool name with empty parameters
|
||||
*current_tool_name_sent = true;
|
||||
result.calls.push(ToolCallItem {
|
||||
tool_index: *current_tool_id as usize,
|
||||
name: Some(function_name.to_string()),
|
||||
parameters: String::new(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
// Case 2: Handle streaming arguments
|
||||
else if let Some(cur_arguments) = current_tool_call.get("arguments") {
|
||||
let tool_id = *current_tool_id as usize;
|
||||
let sent = streamed_args_for_tool
|
||||
.get(tool_id)
|
||||
.map(|s| s.len())
|
||||
.unwrap_or(0);
|
||||
let cur_args_json = serde_json::to_string(cur_arguments)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
// Get prev_arguments (matches Python's structure)
|
||||
let prev_arguments = if tool_id < prev_tool_call_arr.len() {
|
||||
prev_tool_call_arr[tool_id].get("arguments")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Calculate diff: everything after we've already sent
|
||||
let mut argument_diff = None;
|
||||
|
||||
if is_complete {
|
||||
// Python: argument_diff = cur_args_json[sent:]
|
||||
// Rust needs bounds check (Python returns "" automatically)
|
||||
argument_diff = if sent < cur_args_json.len() {
|
||||
Some(cur_args_json[sent..].to_string())
|
||||
} else {
|
||||
Some(String::new())
|
||||
};
|
||||
} else if let Some(prev_args) = prev_arguments {
|
||||
let prev_args_json = serde_json::to_string(prev_args)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
if cur_args_json != prev_args_json {
|
||||
let prefix = find_common_prefix(&prev_args_json, &cur_args_json);
|
||||
argument_diff = if sent < prefix.len() {
|
||||
Some(prefix[sent..].to_string())
|
||||
} else {
|
||||
Some(String::new())
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Send diff if present
|
||||
if let Some(diff) = argument_diff {
|
||||
if !diff.is_empty() {
|
||||
if tool_id < streamed_args_for_tool.len() {
|
||||
streamed_args_for_tool[tool_id].push_str(&diff);
|
||||
}
|
||||
result.calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: None,
|
||||
parameters: diff,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Update prev_tool_call_arr with current state
|
||||
if *current_tool_id >= 0 {
|
||||
ensure_capacity(*current_tool_id, prev_tool_call_arr, streamed_args_for_tool);
|
||||
|
||||
if tool_id < prev_tool_call_arr.len() {
|
||||
prev_tool_call_arr[tool_id] = current_tool_call;
|
||||
}
|
||||
}
|
||||
|
||||
// If complete, advance to next tool
|
||||
if is_complete {
|
||||
*buffer = current_text[start_idx + end_idx..].to_string();
|
||||
*current_tool_name_sent = false;
|
||||
*current_tool_id += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_ends_with_partial_token() {
|
||||
assert!(ends_with_partial_token("hello <|py", "<|python_tag|>").is_some());
|
||||
assert!(ends_with_partial_token("hello <|python_tag", "<|python_tag|>").is_some());
|
||||
assert!(ends_with_partial_token("hello <|python_tag|>", "<|python_tag|>").is_none());
|
||||
assert!(ends_with_partial_token("", "<|python_tag|>").is_none());
|
||||
assert!(ends_with_partial_token("hello world", "<|python_tag|>").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_current_tool_state() {
|
||||
let mut buffer = String::from("partial json");
|
||||
let mut current_tool_name_sent = true;
|
||||
let mut streamed_args = vec!["tool0_args".to_string(), "tool1_partial".to_string()];
|
||||
let prev_tools = vec![serde_json::json!({"name": "tool0"})];
|
||||
|
||||
reset_current_tool_state(
|
||||
&mut buffer,
|
||||
&mut current_tool_name_sent,
|
||||
&mut streamed_args,
|
||||
&prev_tools,
|
||||
);
|
||||
|
||||
assert_eq!(buffer, "");
|
||||
assert!(!current_tool_name_sent);
|
||||
assert_eq!(streamed_args.len(), 1); // Popped the partial tool1 args
|
||||
assert_eq!(streamed_args[0], "tool0_args");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_current_tool_state_no_pop_when_synced() {
|
||||
let mut buffer = String::from("partial json");
|
||||
let mut current_tool_name_sent = true;
|
||||
let mut streamed_args = vec!["tool0_args".to_string()];
|
||||
let prev_tools = vec![serde_json::json!({"name": "tool0"})];
|
||||
|
||||
reset_current_tool_state(
|
||||
&mut buffer,
|
||||
&mut current_tool_name_sent,
|
||||
&mut streamed_args,
|
||||
&prev_tools,
|
||||
);
|
||||
|
||||
assert_eq!(buffer, "");
|
||||
assert!(!current_tool_name_sent);
|
||||
assert_eq!(streamed_args.len(), 1); // No pop, lengths matched
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_parser_state() {
|
||||
let mut buffer = String::from("some buffer");
|
||||
let mut prev_tools = vec![serde_json::json!({"name": "tool0"})];
|
||||
let mut current_tool_id = 5;
|
||||
let mut current_tool_name_sent = true;
|
||||
let mut streamed_args = vec!["args".to_string()];
|
||||
|
||||
reset_parser_state(
|
||||
&mut buffer,
|
||||
&mut prev_tools,
|
||||
&mut current_tool_id,
|
||||
&mut current_tool_name_sent,
|
||||
&mut streamed_args,
|
||||
);
|
||||
|
||||
assert_eq!(buffer, "");
|
||||
assert_eq!(prev_tools.len(), 0);
|
||||
assert_eq!(current_tool_id, -1);
|
||||
assert!(!current_tool_name_sent);
|
||||
assert_eq!(streamed_args.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_capacity() {
|
||||
let mut prev_tools = vec![];
|
||||
let mut streamed_args = vec![];
|
||||
|
||||
ensure_capacity(2, &mut prev_tools, &mut streamed_args);
|
||||
|
||||
assert_eq!(prev_tools.len(), 3);
|
||||
assert_eq!(streamed_args.len(), 3);
|
||||
assert_eq!(prev_tools[0], Value::Null);
|
||||
assert_eq!(streamed_args[0], "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_capacity_negative_id() {
|
||||
let mut prev_tools = vec![];
|
||||
let mut streamed_args = vec![];
|
||||
|
||||
ensure_capacity(-1, &mut prev_tools, &mut streamed_args);
|
||||
|
||||
// Should not resize for negative ID
|
||||
assert_eq!(prev_tools.len(), 0);
|
||||
assert_eq!(streamed_args.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_complete_json() {
|
||||
assert!(is_complete_json(r#"{"name": "test"}"#));
|
||||
assert!(is_complete_json("[1, 2, 3]"));
|
||||
assert!(is_complete_json("42"));
|
||||
assert!(is_complete_json("true"));
|
||||
assert!(!is_complete_json(r#"{"name": "#));
|
||||
assert!(!is_complete_json("[1, 2,"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_arguments_field() {
|
||||
// Case 1: Has parameters, no arguments
|
||||
let obj = serde_json::json!({
|
||||
"name": "test",
|
||||
"parameters": {"key": "value"}
|
||||
});
|
||||
let normalized = normalize_arguments_field(obj);
|
||||
assert_eq!(
|
||||
normalized.get("arguments").unwrap(),
|
||||
&serde_json::json!({"key": "value"})
|
||||
);
|
||||
|
||||
// Case 2: Already has arguments
|
||||
let obj = serde_json::json!({
|
||||
"name": "test",
|
||||
"arguments": {"key": "value"}
|
||||
});
|
||||
let normalized = normalize_arguments_field(obj.clone());
|
||||
assert_eq!(normalized, obj);
|
||||
|
||||
// Case 3: No parameters or arguments
|
||||
let obj = serde_json::json!({"name": "test"});
|
||||
let normalized = normalize_arguments_field(obj.clone());
|
||||
assert_eq!(normalized, obj);
|
||||
}
|
||||
}
|
||||
@@ -1,305 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
partial_json::PartialJson,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// JSON format parser for tool calls
|
||||
///
|
||||
/// Handles pure JSON formats for function calling:
|
||||
/// - Single tool call: {"name": "fn", "arguments": {...}}
|
||||
/// - Multiple tool calls: [{"name": "fn1", "arguments": {...}}, ...]
|
||||
/// - With parameters instead of arguments: {"name": "fn", "parameters": {...}}
|
||||
pub struct JsonParser {
|
||||
/// Parser for handling incomplete JSON during streaming
|
||||
partial_json: PartialJson,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Separator between multiple tool calls
|
||||
tool_call_separator: &'static str,
|
||||
|
||||
/// Track whether we're parsing array format `[...]` vs single object `{...}`
|
||||
is_array_format: bool,
|
||||
|
||||
/// Track whether we've already stripped the closing ] bracket (for array format)
|
||||
array_closed: bool,
|
||||
}
|
||||
|
||||
impl JsonParser {
|
||||
/// Create a new JSON parser
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
partial_json: PartialJson::default(),
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
tool_call_separator: ",",
|
||||
is_array_format: false,
|
||||
array_closed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to extract a first valid JSON object or array from text that may contain other content
|
||||
/// Returns (json_string, normal_text) where normal_text is text before and after the JSON
|
||||
fn extract_json_from_text(&self, text: &str) -> Option<(String, String)> {
|
||||
let mut in_string = false;
|
||||
let mut escape = false;
|
||||
let mut stack: Vec<char> = Vec::with_capacity(8);
|
||||
let mut start: Option<usize> = None;
|
||||
|
||||
for (i, ch) in text.char_indices() {
|
||||
if escape {
|
||||
escape = false;
|
||||
continue;
|
||||
}
|
||||
|
||||
match ch {
|
||||
'\\' if in_string => escape = true,
|
||||
'"' => in_string = !in_string,
|
||||
_ if in_string => {}
|
||||
'{' | '[' => {
|
||||
if start.is_none() {
|
||||
start = Some(i);
|
||||
}
|
||||
stack.push(ch);
|
||||
}
|
||||
'}' | ']' => {
|
||||
let Some(open) = stack.pop() else {
|
||||
// Stray closer - reset and continue looking for next valid JSON
|
||||
start = None;
|
||||
continue;
|
||||
};
|
||||
|
||||
let valid = (open == '{' && ch == '}') || (open == '[' && ch == ']');
|
||||
if !valid {
|
||||
// Mismatch - reset and continue looking
|
||||
start = None;
|
||||
stack.clear();
|
||||
continue;
|
||||
}
|
||||
|
||||
if stack.is_empty() {
|
||||
let s = start.unwrap();
|
||||
let e = i + ch.len_utf8();
|
||||
let potential_json = &text[s..e];
|
||||
|
||||
// Validate that this is actually valid JSON before returning
|
||||
if serde_json::from_str::<Value>(potential_json).is_ok() {
|
||||
let json = potential_json.to_string();
|
||||
let normal = format!("{}{}", &text[..s], &text[e..]);
|
||||
return Some((json, normal));
|
||||
} else {
|
||||
// Not valid JSON, reset and continue looking
|
||||
start = None;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Parse a single JSON object into a ToolCall
|
||||
fn parse_single_object(&self, obj: &Value) -> ParserResult<Option<ToolCall>> {
|
||||
// Check if this looks like a tool call
|
||||
let name = obj
|
||||
.get("name")
|
||||
.or_else(|| obj.get("function"))
|
||||
.and_then(|v| v.as_str());
|
||||
|
||||
if let Some(name) = name {
|
||||
// Get arguments - support both "arguments" and "parameters" keys
|
||||
let empty_obj = Value::Object(serde_json::Map::new());
|
||||
let args = obj
|
||||
.get("arguments")
|
||||
.or_else(|| obj.get("parameters"))
|
||||
.unwrap_or(&empty_obj);
|
||||
|
||||
// Convert arguments to JSON string
|
||||
let arguments = serde_json::to_string(args)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse JSON value(s) into tool calls
|
||||
fn parse_json_value(&self, value: &Value) -> ParserResult<Vec<ToolCall>> {
|
||||
let mut tools = Vec::new();
|
||||
|
||||
match value {
|
||||
Value::Array(arr) => {
|
||||
// Parse each element in the array
|
||||
for item in arr {
|
||||
if let Some(tool) = self.parse_single_object(item)? {
|
||||
tools.push(tool);
|
||||
}
|
||||
}
|
||||
}
|
||||
Value::Object(_) => {
|
||||
// Single tool call
|
||||
if let Some(tool) = self.parse_single_object(value)? {
|
||||
tools.push(tool);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// Not a valid tool call format
|
||||
return Ok(vec![]);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(tools)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for JsonParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for JsonParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Always use extract_json_from_text to handle both pure JSON and mixed content
|
||||
if let Some((extracted_json, normal_text)) = self.extract_json_from_text(text) {
|
||||
let parsed = serde_json::from_str::<Value>(&extracted_json)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))
|
||||
.and_then(|v| self.parse_json_value(&v));
|
||||
|
||||
match parsed {
|
||||
Ok(tools) => return Ok((normal_text, tools)),
|
||||
Err(e) => tracing::debug!("parse_complete failed: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// No valid JSON found, return original text as normal text
|
||||
Ok((text.to_string(), vec![]))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Append new text to buffer
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Determine format on first parse (array vs single object)
|
||||
if self.current_tool_id == -1 && self.has_tool_markers(current_text) {
|
||||
self.is_array_format = current_text.trim().starts_with('[');
|
||||
}
|
||||
|
||||
// Check if current_text has tool_call
|
||||
// Once array is closed, don't treat [ or { as tool markers
|
||||
let has_tool_start = (!self.array_closed && self.has_tool_markers(current_text))
|
||||
|| (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
|
||||
|
||||
if !has_tool_start {
|
||||
let mut normal_text = self.buffer.clone();
|
||||
self.buffer.clear();
|
||||
|
||||
// Strip ] only once (the closing bracket of JSON array format)
|
||||
// Only for array format and only if we haven't already closed it
|
||||
if self.is_array_format
|
||||
&& !self.array_closed
|
||||
&& self.current_tool_id > 0
|
||||
&& normal_text.starts_with("]")
|
||||
{
|
||||
normal_text = normal_text.strip_prefix("]").unwrap().to_string();
|
||||
self.array_closed = true;
|
||||
}
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Build tool indices
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
// Determine start index for JSON parsing
|
||||
// JSON can start with [ (array) or { (single object)
|
||||
let start_idx = if let Some(bracket_pos) = current_text.find('[') {
|
||||
let brace_pos = current_text.find('{');
|
||||
match brace_pos {
|
||||
Some(bp) => bp,
|
||||
_ => bracket_pos,
|
||||
}
|
||||
} else if let Some(brace_pos) = current_text.find('{') {
|
||||
brace_pos
|
||||
} else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
|
||||
self.tool_call_separator.len()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
helpers::handle_json_tool_streaming(
|
||||
current_text,
|
||||
start_idx,
|
||||
&mut self.partial_json,
|
||||
&tool_indices,
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&mut self.prev_tool_call_arr,
|
||||
)
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
let trimmed = text.trim();
|
||||
trimmed.starts_with('[') || trimmed.starts_with('{')
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
helpers::reset_parser_state(
|
||||
&mut self.buffer,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
self.is_array_format = false;
|
||||
self.array_closed = false;
|
||||
}
|
||||
}
|
||||
@@ -1,348 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::ParserResult,
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// Kimi K2 format parser for tool calls
|
||||
///
|
||||
/// Handles the Kimi K2 specific format:
|
||||
/// `<|tool_calls_section_begin|><|tool_call_begin|>functions.{name}:{index}<|tool_call_argument_begin|>{json_args}<|tool_call_end|><|tool_calls_section_end|>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Token-based delimiters
|
||||
/// - Function calls with explicit indexing
|
||||
/// - JSON arguments
|
||||
///
|
||||
/// Reference: https://huggingface.co/moonshotai/Kimi-K2-Instruct/blob/main/docs/tool_call_guidance.md
|
||||
pub struct KimiK2Parser {
|
||||
/// Regex for extracting complete tool calls
|
||||
tool_call_extractor: Regex,
|
||||
/// Regex for extracting partial tool calls (streaming)
|
||||
stream_tool_call_extractor: Regex,
|
||||
/// Regex pattern for removing completed tool calls from buffer
|
||||
tool_call_end_pattern: Regex,
|
||||
/// Robust parser for ids like "functions.search:0" or fallback "search:0"
|
||||
tool_call_id_regex: Regex,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Tracks the last arguments sent for incremental diffing
|
||||
last_arguments: String,
|
||||
}
|
||||
|
||||
impl KimiK2Parser {
|
||||
/// Create a new Kimi K2 parser
|
||||
pub fn new() -> Self {
|
||||
// Pattern for complete tool calls
|
||||
let tool_call_pattern = r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*?\})\s*<\|tool_call_end\|>";
|
||||
let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Pattern for streaming (partial) tool calls
|
||||
let stream_pattern = r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*)";
|
||||
let stream_tool_call_extractor = Regex::new(stream_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Pattern for removing completed tool calls
|
||||
let end_pattern = r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>";
|
||||
let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Robust parser for ids like "functions.search:0" or fallback "search:0"
|
||||
let id_pattern = r"^(?:functions\.)?(?P<name>[\w\.]+):(?P<index>\d+)$";
|
||||
let tool_call_id_regex = Regex::new(id_pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
tool_call_extractor,
|
||||
stream_tool_call_extractor,
|
||||
tool_call_end_pattern,
|
||||
tool_call_id_regex,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
last_arguments: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse function ID to extract name and index
|
||||
fn parse_function_id(&self, id: &str) -> Option<(String, usize)> {
|
||||
if let Some(captures) = self.tool_call_id_regex.captures(id) {
|
||||
let name = captures.name("name")?.as_str().to_string();
|
||||
let index = captures.name("index")?.as_str().parse::<usize>().ok()?;
|
||||
Some((name, index))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for KimiK2Parser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for KimiK2Parser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where tool calls begin
|
||||
let idx = text.find("<|tool_calls_section_begin|>").unwrap();
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Try to extract tool calls
|
||||
let mut tools = Vec::new();
|
||||
for captures in self.tool_call_extractor.captures_iter(text) {
|
||||
if let (Some(id_match), Some(args_match)) = (
|
||||
captures.name("tool_call_id"),
|
||||
captures.name("function_arguments"),
|
||||
) {
|
||||
let function_id = id_match.as_str();
|
||||
let function_args = args_match.as_str();
|
||||
|
||||
// Parse function ID
|
||||
if let Some((func_name, _index)) = self.parse_function_id(function_id) {
|
||||
// Try to parse JSON arguments
|
||||
match serde_json::from_str::<Value>(function_args) {
|
||||
Ok(_) => {
|
||||
tools.push(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: func_name,
|
||||
arguments: function_args.to_string(),
|
||||
},
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
"Failed to parse JSON arguments for {}: {}",
|
||||
func_name,
|
||||
e
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tracing::debug!("Failed to parse function ID: {}", function_id);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if we have a tool call (either the start token or individual tool call)
|
||||
let has_tool_call =
|
||||
self.has_tool_markers(current_text) || current_text.contains("<|tool_call_begin|>");
|
||||
|
||||
if !has_tool_call {
|
||||
// No tool markers detected - return all buffered content as normal text
|
||||
let mut normal_text = std::mem::take(&mut self.buffer);
|
||||
// Remove end tokens if present
|
||||
for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>"] {
|
||||
normal_text = normal_text.replace(e_token, "");
|
||||
}
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
let mut calls: Vec<ToolCallItem> = Vec::new();
|
||||
|
||||
// Try to match streaming pattern
|
||||
if let Some(captures) = self.stream_tool_call_extractor.captures(current_text) {
|
||||
if let (Some(id_match), Some(args_match)) = (
|
||||
captures.name("tool_call_id"),
|
||||
captures.name("function_arguments"),
|
||||
) {
|
||||
let function_id = id_match.as_str();
|
||||
let function_args = args_match.as_str();
|
||||
|
||||
// Parse function ID
|
||||
if let Some((func_name, _index)) = self.parse_function_id(function_id) {
|
||||
// Validate tool name
|
||||
if !tool_indices.contains_key(&func_name) {
|
||||
// Invalid tool name - skip this tool, preserve indexing for next tool
|
||||
tracing::debug!("Invalid tool name '{}' - skipping", func_name);
|
||||
helpers::reset_current_tool_state(
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&self.prev_tool_call_arr,
|
||||
);
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
|
||||
// Initialize state if this is the first tool call
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
self.prev_tool_call_arr = Vec::new();
|
||||
self.streamed_args_for_tool = vec![String::new()];
|
||||
}
|
||||
|
||||
// Ensure we have enough entries in our tracking arrays
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Send tool name if not sent yet
|
||||
if !self.current_tool_name_sent {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: Some(func_name.clone()),
|
||||
parameters: String::new(),
|
||||
});
|
||||
self.current_tool_name_sent = true;
|
||||
|
||||
// Store the tool call info for serving layer completions endpoint
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if self.prev_tool_call_arr.len() <= tool_id {
|
||||
self.prev_tool_call_arr
|
||||
.resize_with(tool_id + 1, || Value::Null);
|
||||
}
|
||||
self.prev_tool_call_arr[tool_id] = serde_json::json!({
|
||||
"name": func_name,
|
||||
"arguments": {},
|
||||
});
|
||||
} else {
|
||||
// Compute incremental diff
|
||||
let argument_diff = if function_args.starts_with(&self.last_arguments) {
|
||||
&function_args[self.last_arguments.len()..]
|
||||
} else {
|
||||
function_args
|
||||
};
|
||||
|
||||
// Split by end token before sending (like Python does)
|
||||
let parsed_args_diff =
|
||||
if let Some(pos) = argument_diff.find("<|tool_call_end|>") {
|
||||
&argument_diff[..pos]
|
||||
} else {
|
||||
argument_diff
|
||||
};
|
||||
|
||||
if !parsed_args_diff.is_empty() {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: parsed_args_diff.to_string(),
|
||||
});
|
||||
// Note: Python adds full diff to _last_arguments, not just parsed part
|
||||
self.last_arguments.push_str(argument_diff);
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.streamed_args_for_tool.len() {
|
||||
self.streamed_args_for_tool[tool_id].push_str(parsed_args_diff);
|
||||
}
|
||||
}
|
||||
|
||||
// Check completeness - split by end token first
|
||||
let parsed_args = if let Some(pos) = function_args.find("<|tool_call_end|>")
|
||||
{
|
||||
&function_args[..pos]
|
||||
} else {
|
||||
function_args
|
||||
};
|
||||
|
||||
if helpers::is_complete_json(parsed_args) {
|
||||
// Update the stored arguments
|
||||
if let Ok(parsed_args_value) =
|
||||
serde_json::from_str::<Value>(parsed_args)
|
||||
{
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.prev_tool_call_arr.len() {
|
||||
if let Some(obj) =
|
||||
self.prev_tool_call_arr[tool_id].as_object_mut()
|
||||
{
|
||||
obj.insert("arguments".to_string(), parsed_args_value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Find the end of the current tool call and remove only that part from buffer
|
||||
if let Some(mat) = self.tool_call_end_pattern.find(current_text) {
|
||||
// Remove the completed tool call from buffer, keep any remaining content
|
||||
self.buffer = current_text[mat.end()..].to_string();
|
||||
} else {
|
||||
self.buffer.clear();
|
||||
}
|
||||
|
||||
let result = StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
};
|
||||
|
||||
self.current_tool_id += 1;
|
||||
self.last_arguments.clear();
|
||||
self.current_tool_name_sent = false;
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
})
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains("<|tool_calls_section_begin|>")
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.prev_tool_call_arr.clear();
|
||||
self.current_tool_id = -1;
|
||||
self.current_tool_name_sent = false;
|
||||
self.streamed_args_for_tool.clear();
|
||||
self.last_arguments.clear();
|
||||
}
|
||||
}
|
||||
@@ -1,245 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
partial_json::PartialJson,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall},
|
||||
},
|
||||
};
|
||||
|
||||
/// Llama 3.2 format parser for tool calls
|
||||
///
|
||||
/// Handles the Llama 3.2 specific format:
|
||||
/// `<|python_tag|>{"name": "func", "parameters": {...}}`
|
||||
///
|
||||
/// Also supports plain JSON without the python_tag prefix
|
||||
pub struct LlamaParser {
|
||||
/// Parser for handling incomplete JSON during streaming
|
||||
partial_json: PartialJson,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Token configuration
|
||||
bot_token: &'static str,
|
||||
tool_call_separator: &'static str,
|
||||
}
|
||||
|
||||
impl LlamaParser {
|
||||
/// Create a new Llama parser
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
partial_json: PartialJson::default(),
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
bot_token: "<|python_tag|>",
|
||||
tool_call_separator: ";",
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract content after python_tag token
|
||||
fn extract_content_after_python_tag(&self, text: &str) -> Option<(String, String)> {
|
||||
const PYTHON_TAG: &str = "<|python_tag|>";
|
||||
|
||||
if let Some(tag_pos) = text.find(PYTHON_TAG) {
|
||||
let normal_text = text[..tag_pos].to_string();
|
||||
let json_content = text[tag_pos + PYTHON_TAG.len()..].to_string();
|
||||
Some((normal_text, json_content))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a single JSON object into a ToolCall (Llama format: name + parameters)
|
||||
fn parse_single_object(&self, obj: &Value) -> ParserResult<Option<ToolCall>> {
|
||||
// Llama format only: {"name": "function_name", "parameters": {...}}
|
||||
let name = obj.get("name").and_then(|v| v.as_str());
|
||||
|
||||
if let Some(name) = name {
|
||||
// Llama uses "parameters" key
|
||||
let empty_obj = Value::Object(serde_json::Map::new());
|
||||
let parameters = obj.get("parameters").unwrap_or(&empty_obj);
|
||||
|
||||
// Convert parameters to JSON string
|
||||
let arguments = serde_json::to_string(parameters)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse semicolon-separated JSON objects
|
||||
fn parse_semicolon_separated(&self, content: &str) -> ParserResult<Vec<ToolCall>> {
|
||||
let mut all_tools = Vec::new();
|
||||
|
||||
// Split by semicolon and parse each JSON object
|
||||
for part in content.split(';') {
|
||||
let trimmed = part.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Try to parse this part as a single JSON object
|
||||
match serde_json::from_str::<Value>(trimmed) {
|
||||
Ok(value) => {
|
||||
if let Some(tool) = self.parse_single_object(&value)? {
|
||||
all_tools.push(tool);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Skip invalid JSON parts in semicolon-separated list
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_tools)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LlamaParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for LlamaParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Extract normal text and JSON content
|
||||
let (normal_text, json_content) =
|
||||
if let Some((normal, json)) = self.extract_content_after_python_tag(text) {
|
||||
(normal, json)
|
||||
} else if text.trim_start().starts_with('{') {
|
||||
(String::new(), text.to_string())
|
||||
} else {
|
||||
// No JSON structure found
|
||||
return Ok((text.to_string(), vec![]));
|
||||
};
|
||||
|
||||
// Parse the JSON content (may contain semicolon-separated objects)
|
||||
let tools = if json_content.contains(';') {
|
||||
self.parse_semicolon_separated(&json_content)?
|
||||
} else {
|
||||
// Try single JSON object
|
||||
let parsed = serde_json::from_str::<Value>(json_content.trim())
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))
|
||||
.and_then(|v| {
|
||||
self.parse_single_object(&v)
|
||||
.map(|opt| opt.map_or_else(Vec::new, |tool| vec![tool]))
|
||||
});
|
||||
|
||||
parsed.unwrap_or_else(|e| {
|
||||
tracing::debug!("Failed to parse tool call: {:?}", e);
|
||||
vec![]
|
||||
})
|
||||
};
|
||||
|
||||
// If we couldn't parse any tools, return the original text
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Append new text to buffer
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if current_text has tool_call
|
||||
let has_tool_start = self.has_tool_markers(current_text)
|
||||
|| (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
|
||||
|
||||
if !has_tool_start {
|
||||
// Only clear buffer if we're sure no tool call is starting
|
||||
if helpers::ends_with_partial_token(&self.buffer, self.bot_token).is_none() {
|
||||
let normal_text = self.buffer.clone();
|
||||
self.buffer.clear();
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
} else {
|
||||
// Might be partial bot_token, keep buffering
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// Build tool indices
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
// Determine start index for JSON parsing
|
||||
let start_idx = if let Some(pos) = current_text.find(self.bot_token) {
|
||||
pos + self.bot_token.len()
|
||||
} else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
|
||||
self.tool_call_separator.len()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
helpers::handle_json_tool_streaming(
|
||||
current_text,
|
||||
start_idx,
|
||||
&mut self.partial_json,
|
||||
&tool_indices,
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&mut self.prev_tool_call_arr,
|
||||
)
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
// Llama format if contains python_tag or starts with JSON object
|
||||
text.contains("<|python_tag|>") || text.trim_start().starts_with('{')
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::tool_parser::types::ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
helpers::reset_parser_state(
|
||||
&mut self.buffer,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,549 +0,0 @@
|
||||
use std::{collections::HashMap, fmt::Write as FmtWrite};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// MiniMax M2 format parser for tool calls
|
||||
///
|
||||
/// Handles the MiniMax M2 specific format:
|
||||
/// `<minimax:tool_call><invoke name="func"><parameter name="key">value</parameter></invoke></minimax:tool_call>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Namespaced XML tags (`minimax:tool_call`)
|
||||
/// - Function wrapped in `<invoke name="...">` tags
|
||||
/// - Parameters as `<parameter name="key">value</parameter>`
|
||||
/// - Incremental JSON streaming for parameters
|
||||
///
|
||||
/// Reference: https://huggingface.co/MiniMaxAI/MiniMax-M2?chat_template=default
|
||||
pub struct MinimaxM2Parser {
|
||||
// Regex patterns
|
||||
tool_call_extractor: Regex,
|
||||
invoke_extractor: Regex,
|
||||
param_extractor: Regex,
|
||||
|
||||
// Streaming state
|
||||
buffer: String,
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
current_tool_id: i32,
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
current_function_name: String,
|
||||
current_parameters: HashMap<String, Value>,
|
||||
in_tool_call: bool,
|
||||
function_name_sent: bool,
|
||||
waiting_for_tool_call_end: bool,
|
||||
|
||||
// Token configuration
|
||||
tool_call_start_token: &'static str,
|
||||
tool_call_end_token: &'static str,
|
||||
invoke_end_token: &'static str,
|
||||
}
|
||||
|
||||
impl MinimaxM2Parser {
|
||||
/// Parse a value from string with consistent logic
|
||||
#[inline]
|
||||
fn parse_value(text: &str) -> Value {
|
||||
// Try parsing as common literals first
|
||||
match text {
|
||||
"true" | "True" => return Value::Bool(true),
|
||||
"false" | "False" => return Value::Bool(false),
|
||||
"null" | "None" => return Value::Null,
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Try parsing as number
|
||||
if let Ok(num) = text.parse::<i64>() {
|
||||
return Value::Number(num.into());
|
||||
}
|
||||
|
||||
if let Ok(num) = text.parse::<f64>() {
|
||||
if let Some(n) = serde_json::Number::from_f64(num) {
|
||||
return Value::Number(n);
|
||||
}
|
||||
}
|
||||
|
||||
// Default to string
|
||||
Value::String(text.to_string())
|
||||
}
|
||||
|
||||
/// Create a new MiniMax M2 parser
|
||||
pub fn new() -> Self {
|
||||
// Use (?s) flag for DOTALL mode to handle newlines
|
||||
let tool_call_pattern = r"(?s)<minimax:tool_call>.*?</minimax:tool_call>";
|
||||
let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
|
||||
|
||||
let invoke_pattern = r#"(?s)<invoke\s+name="([^"]+)">(.*?)</invoke>"#;
|
||||
let invoke_extractor = Regex::new(invoke_pattern).expect("Valid regex pattern");
|
||||
|
||||
let param_pattern = r#"(?s)<parameter\s+name="([^"]+)">(.*?)</parameter>"#;
|
||||
let param_extractor = Regex::new(param_pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
tool_call_extractor,
|
||||
invoke_extractor,
|
||||
param_extractor,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
current_function_name: String::new(),
|
||||
current_parameters: HashMap::new(),
|
||||
in_tool_call: false,
|
||||
function_name_sent: false,
|
||||
waiting_for_tool_call_end: false,
|
||||
tool_call_start_token: "<minimax:tool_call>",
|
||||
tool_call_end_token: "</minimax:tool_call>",
|
||||
invoke_end_token: "</invoke>",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse parameters from parameter tags
|
||||
fn parse_parameters(&self, params_text: &str) -> ParserResult<serde_json::Map<String, Value>> {
|
||||
let mut parameters = serde_json::Map::new();
|
||||
|
||||
for capture in self.param_extractor.captures_iter(params_text) {
|
||||
let key = capture.get(1).map_or("", |m| m.as_str()).trim();
|
||||
let value_str = capture.get(2).map_or("", |m| m.as_str());
|
||||
|
||||
// Decode XML entities and parse value
|
||||
let decoded_value = self.decode_xml_entities(value_str);
|
||||
|
||||
// Note: We keep JSON-like strings as strings (not parsed JSON)
|
||||
// This matches the behavior of other parsers like GLM4 MOE
|
||||
let value = Self::parse_value(&decoded_value);
|
||||
|
||||
parameters.insert(key.to_string(), value);
|
||||
}
|
||||
|
||||
Ok(parameters)
|
||||
}
|
||||
|
||||
/// Decode common XML entities
|
||||
fn decode_xml_entities(&self, text: &str) -> String {
|
||||
text.replace("<", "<")
|
||||
.replace(">", ">")
|
||||
.replace("&", "&")
|
||||
.replace(""", "\"")
|
||||
.replace("'", "'")
|
||||
}
|
||||
|
||||
/// Parse a single tool call block
|
||||
fn parse_tool_call(&self, block: &str) -> ParserResult<Option<ToolCall>> {
|
||||
if let Some(captures) = self.invoke_extractor.captures(block) {
|
||||
// Get function name from invoke tag attribute
|
||||
let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Get parameters text
|
||||
let params_text = captures.get(2).map_or("", |m| m.as_str());
|
||||
|
||||
// Parse parameters
|
||||
let parameters = self.parse_parameters(params_text)?;
|
||||
|
||||
let arguments_str = serde_json::to_string(¶meters)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: func_name.to_string(),
|
||||
arguments: arguments_str,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse all tool calls from text and return first valid position
|
||||
fn parse_tool_calls_from_text(
|
||||
&self,
|
||||
text: &str,
|
||||
) -> ParserResult<(Vec<ToolCall>, Option<usize>)> {
|
||||
let mut tools = Vec::new();
|
||||
let mut first_valid_pos = None;
|
||||
|
||||
for mat in self.tool_call_extractor.find_iter(text) {
|
||||
match self.parse_tool_call(mat.as_str()) {
|
||||
Ok(Some(tool)) => {
|
||||
if first_valid_pos.is_none() {
|
||||
first_valid_pos = Some(mat.start());
|
||||
}
|
||||
tools.push(tool);
|
||||
}
|
||||
Ok(None) => continue,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((tools, first_valid_pos))
|
||||
}
|
||||
|
||||
/// Parse and stream parameters incrementally
|
||||
fn parse_and_stream_parameters(&mut self, text: &str, _tools: &[Tool]) -> Vec<ToolCallItem> {
|
||||
let mut calls = Vec::new();
|
||||
|
||||
// Find all complete parameter patterns in the buffer
|
||||
let param_matches: Vec<_> = self
|
||||
.param_extractor
|
||||
.captures_iter(text)
|
||||
.map(|cap| {
|
||||
let name = cap.get(1).map_or("", |m| m.as_str()).trim().to_string();
|
||||
let value_str = cap.get(2).map_or("", |m| m.as_str());
|
||||
let decoded = self.decode_xml_entities(value_str);
|
||||
|
||||
// Try parsing as JSON first (for nested objects/arrays)
|
||||
let value = if decoded.starts_with('{') || decoded.starts_with('[') {
|
||||
if let Ok(json_val) = serde_json::from_str::<Value>(&decoded) {
|
||||
json_val
|
||||
} else {
|
||||
Self::parse_value(&decoded)
|
||||
}
|
||||
} else {
|
||||
Self::parse_value(&decoded)
|
||||
};
|
||||
|
||||
(name, value)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Build new parameters map
|
||||
let mut new_params = HashMap::new();
|
||||
for (name, value) in param_matches {
|
||||
new_params.insert(name, value);
|
||||
}
|
||||
|
||||
// If we have new parameters that weren't in current_parameters, stream them
|
||||
if !new_params.is_empty() && new_params != self.current_parameters {
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
|
||||
// Ensure we have enough capacity
|
||||
while self.streamed_args_for_tool.len() <= tool_id {
|
||||
self.streamed_args_for_tool.push(String::new());
|
||||
}
|
||||
|
||||
// Build incremental JSON with single allocation
|
||||
if self.current_parameters.is_empty() {
|
||||
// First parameters - start JSON object but don't close it
|
||||
let mut json_fragment = String::with_capacity(256);
|
||||
json_fragment.push('{');
|
||||
|
||||
let mut first = true;
|
||||
for (key, value) in &new_params {
|
||||
if !first {
|
||||
json_fragment.push_str(", ");
|
||||
}
|
||||
write!(
|
||||
&mut json_fragment,
|
||||
"{}: {}",
|
||||
serde_json::to_string(key).unwrap(),
|
||||
serde_json::to_string(value).unwrap()
|
||||
)
|
||||
.unwrap();
|
||||
first = false;
|
||||
}
|
||||
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: None,
|
||||
parameters: json_fragment.clone(),
|
||||
});
|
||||
|
||||
self.streamed_args_for_tool[tool_id] = json_fragment;
|
||||
} else {
|
||||
// Additional parameters - add them incrementally
|
||||
let new_keys: Vec<_> = new_params
|
||||
.keys()
|
||||
.filter(|k| !self.current_parameters.contains_key(*k))
|
||||
.collect();
|
||||
|
||||
if !new_keys.is_empty() {
|
||||
let mut json_fragment = String::with_capacity(128);
|
||||
|
||||
for key in new_keys {
|
||||
let value = &new_params[key];
|
||||
write!(
|
||||
&mut json_fragment,
|
||||
", {}: {}",
|
||||
serde_json::to_string(key).unwrap(),
|
||||
serde_json::to_string(value).unwrap()
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: None,
|
||||
parameters: json_fragment.clone(),
|
||||
});
|
||||
|
||||
self.streamed_args_for_tool[tool_id].push_str(&json_fragment);
|
||||
}
|
||||
}
|
||||
|
||||
// Update current parameters
|
||||
self.current_parameters = new_params;
|
||||
|
||||
// Update prev_tool_call_arr
|
||||
while self.prev_tool_call_arr.len() <= tool_id {
|
||||
self.prev_tool_call_arr.push(Value::Null);
|
||||
}
|
||||
self.prev_tool_call_arr[tool_id] = serde_json::json!({
|
||||
"name": self.current_function_name,
|
||||
"arguments": self.current_parameters,
|
||||
});
|
||||
}
|
||||
|
||||
calls
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MinimaxM2Parser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for MinimaxM2Parser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Check if text contains MiniMax M2 format
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Parse all tool calls and get first valid position
|
||||
let (tools, first_valid_tool_pos) = self.parse_tool_calls_from_text(text)?;
|
||||
|
||||
// If no tools were successfully parsed, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Determine what text to return as normal_text
|
||||
let normal_text = if let Some(pos) = first_valid_tool_pos {
|
||||
// Return text up to the first valid tool call
|
||||
text[..pos].to_string()
|
||||
} else {
|
||||
// No valid tool calls found, return entire text
|
||||
text.to_string()
|
||||
};
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
let mut normal_text = String::new();
|
||||
let mut calls = Vec::new();
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
loop {
|
||||
// If we're waiting for the tool call end tag, check for it first
|
||||
if self.waiting_for_tool_call_end {
|
||||
if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
|
||||
// Complete tool call found
|
||||
self.buffer =
|
||||
self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
|
||||
self.in_tool_call = false;
|
||||
self.waiting_for_tool_call_end = false;
|
||||
self.function_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
self.current_tool_id += 1;
|
||||
continue;
|
||||
} else {
|
||||
// End tag not complete yet, wait for more text
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// If we're not in a tool call and don't see a start token, return normal text
|
||||
if !self.in_tool_call && !self.buffer.contains(self.tool_call_start_token) {
|
||||
// Check if buffer might contain a partial start token at the end
|
||||
if let Some(partial_len) =
|
||||
helpers::ends_with_partial_token(&self.buffer, self.tool_call_start_token)
|
||||
{
|
||||
// Return everything except the potential partial token
|
||||
let end = self.buffer.len() - partial_len;
|
||||
normal_text = self.buffer[..end].to_string();
|
||||
self.buffer = self.buffer[end..].to_string();
|
||||
} else {
|
||||
// No partial token, return all as normal text
|
||||
normal_text = self.buffer.clone();
|
||||
self.buffer.clear();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
// Look for tool call start
|
||||
if !self.in_tool_call {
|
||||
if let Some(start) = self.buffer.find(self.tool_call_start_token) {
|
||||
normal_text = self.buffer[..start].to_string();
|
||||
self.buffer =
|
||||
self.buffer[start + self.tool_call_start_token.len()..].to_string();
|
||||
|
||||
self.in_tool_call = true;
|
||||
self.function_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
|
||||
continue;
|
||||
} else {
|
||||
// No start token found
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// We're in a tool call, try to parse function name if not sent yet
|
||||
if !self.function_name_sent {
|
||||
// Use regex to extract function name from <invoke name="..."> pattern
|
||||
// Check if we have enough text to match the invoke pattern
|
||||
if let Some(captures) = self.invoke_extractor.captures(&self.buffer) {
|
||||
let function_name = captures
|
||||
.get(1)
|
||||
.map_or("", |m| m.as_str())
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
// Validate function name
|
||||
if tool_indices.contains_key(&function_name) {
|
||||
self.current_function_name = function_name.clone();
|
||||
self.function_name_sent = true;
|
||||
|
||||
// Initialize tool call tracking
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
}
|
||||
|
||||
// Ensure tracking arrays are large enough
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Send tool name with empty parameters
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: Some(function_name),
|
||||
parameters: String::new(),
|
||||
});
|
||||
|
||||
// Find the position after the opening invoke tag (after the >)
|
||||
// We only want to remove up to the opening tag, not the full match
|
||||
if let Some(pos) = self.buffer.find('>') {
|
||||
self.buffer = self.buffer[pos + 1..].to_string();
|
||||
}
|
||||
continue;
|
||||
} else {
|
||||
// Invalid function name, reset state
|
||||
tracing::debug!("Invalid function name: {}", function_name);
|
||||
self.in_tool_call = false;
|
||||
normal_text.push_str(&self.buffer);
|
||||
self.buffer.clear();
|
||||
break;
|
||||
}
|
||||
}
|
||||
// No complete invoke pattern found yet, wait for more text
|
||||
break;
|
||||
}
|
||||
|
||||
// Parse parameters incrementally
|
||||
if self.function_name_sent {
|
||||
// Process parameters and get any calls to emit
|
||||
// Note: We need to be careful here - parse_and_stream_parameters needs
|
||||
// to work with the buffer but we can't pass &self.buffer directly
|
||||
// due to borrow checker. Instead, we'll refactor slightly.
|
||||
// For now, keep the clone but mark it as a TODO for future optimization
|
||||
let buffer_copy = self.buffer.clone(); // TODO: Optimize this
|
||||
let parameter_calls = self.parse_and_stream_parameters(&buffer_copy, tools);
|
||||
calls.extend(parameter_calls);
|
||||
|
||||
// Check if tool call is complete (</invoke> found)
|
||||
if let Some(invoke_end) = self.buffer.find(self.invoke_end_token) {
|
||||
// Add closing brace to complete the JSON object
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.streamed_args_for_tool.len() {
|
||||
let current_streamed = &self.streamed_args_for_tool[tool_id];
|
||||
if !current_streamed.is_empty() && !current_streamed.ends_with('}') {
|
||||
// Count opening and closing braces to check if JSON is complete
|
||||
let open_braces = current_streamed.matches('{').count();
|
||||
let close_braces = current_streamed.matches('}').count();
|
||||
if open_braces > close_braces {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: tool_id,
|
||||
name: None,
|
||||
parameters: "}".to_string(),
|
||||
});
|
||||
self.streamed_args_for_tool[tool_id].push('}');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Move buffer past the </invoke>
|
||||
self.buffer =
|
||||
self.buffer[invoke_end + self.invoke_end_token.len()..].to_string();
|
||||
|
||||
// Check if we have the closing </minimax:tool_call>
|
||||
if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
|
||||
// Complete tool call found
|
||||
self.buffer =
|
||||
self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
|
||||
self.in_tool_call = false;
|
||||
self.function_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
self.current_tool_id += 1;
|
||||
continue;
|
||||
} else {
|
||||
// End tag not complete yet, mark that we're waiting for it
|
||||
self.waiting_for_tool_call_end = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Tool call not complete yet, wait for more text
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult { normal_text, calls })
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains(self.tool_call_start_token)
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.buffer.clear();
|
||||
self.prev_tool_call_arr.clear();
|
||||
self.current_tool_id = -1;
|
||||
self.streamed_args_for_tool.clear();
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
self.in_tool_call = false;
|
||||
self.function_name_sent = false;
|
||||
self.waiting_for_tool_call_end = false;
|
||||
}
|
||||
}
|
||||
@@ -1,287 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
partial_json::PartialJson,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall},
|
||||
},
|
||||
};
|
||||
|
||||
/// Mistral format parser for tool calls
|
||||
///
|
||||
/// Handles the Mistral-specific format:
|
||||
/// `[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`
|
||||
///
|
||||
/// Reference: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3?chat_template=default
|
||||
pub struct MistralParser {
|
||||
/// Parser for handling incomplete JSON during streaming
|
||||
partial_json: PartialJson,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Token configuration
|
||||
bot_token: &'static str,
|
||||
eot_token: &'static str,
|
||||
tool_call_separator: &'static str,
|
||||
|
||||
/// Track whether we've already stripped the closing ] bracket
|
||||
array_closed: bool,
|
||||
}
|
||||
|
||||
impl MistralParser {
|
||||
/// Create a new Mistral parser
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
partial_json: PartialJson::default(),
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
bot_token: "[TOOL_CALLS] [",
|
||||
eot_token: "]",
|
||||
tool_call_separator: ", ",
|
||||
array_closed: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_json_array_with_pos<'a>(&self, text: &'a str) -> Option<(usize, &'a str)> {
|
||||
const BOT_TOKEN: &str = "[TOOL_CALLS] [";
|
||||
|
||||
// Find the start of the token
|
||||
let start_idx = text.find(BOT_TOKEN)?;
|
||||
|
||||
// Start from the opening bracket after [TOOL_CALLS]
|
||||
// The -1 is to include the opening bracket that's part of the token
|
||||
let json_start = start_idx + BOT_TOKEN.len() - 1;
|
||||
|
||||
let mut bracket_count = 0;
|
||||
let mut in_string = false;
|
||||
let mut escape_next = false;
|
||||
|
||||
let bytes = text.as_bytes();
|
||||
|
||||
for i in json_start..text.len() {
|
||||
let char = bytes[i];
|
||||
|
||||
if escape_next {
|
||||
escape_next = false;
|
||||
continue;
|
||||
}
|
||||
|
||||
if char == b'\\' {
|
||||
escape_next = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
if char == b'"' && !escape_next {
|
||||
in_string = !in_string;
|
||||
continue;
|
||||
}
|
||||
|
||||
if !in_string {
|
||||
if char == b'[' {
|
||||
bracket_count += 1;
|
||||
} else if char == b']' {
|
||||
bracket_count -= 1;
|
||||
if bracket_count == 0 {
|
||||
// Found the matching closing bracket
|
||||
return Some((start_idx, &text[json_start..=i]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Incomplete array (no matching closing bracket found)
|
||||
None
|
||||
}
|
||||
|
||||
/// Parse tool calls from a JSON array
|
||||
fn parse_json_array(&self, json_str: &str) -> ParserResult<Vec<ToolCall>> {
|
||||
let value: Value = serde_json::from_str(json_str)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
let mut tools = Vec::new();
|
||||
|
||||
if let Value::Array(arr) = value {
|
||||
for item in arr.iter() {
|
||||
if let Some(tool) = self.parse_single_object(item)? {
|
||||
tools.push(tool);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Single object case (shouldn't happen with Mistral format, but handle it)
|
||||
if let Some(tool) = self.parse_single_object(&value)? {
|
||||
tools.push(tool);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(tools)
|
||||
}
|
||||
|
||||
/// Parse a single JSON object into a ToolCall
|
||||
fn parse_single_object(&self, obj: &Value) -> ParserResult<Option<ToolCall>> {
|
||||
let name = obj.get("name").and_then(|v| v.as_str());
|
||||
|
||||
if let Some(name) = name {
|
||||
// Get arguments - Mistral uses "arguments" key
|
||||
let empty_obj = Value::Object(serde_json::Map::new());
|
||||
let args = obj.get("arguments").unwrap_or(&empty_obj);
|
||||
|
||||
// Convert arguments to JSON string
|
||||
let arguments = serde_json::to_string(args)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MistralParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for MistralParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Check if text contains Mistral format
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Extract JSON array from Mistral format with position
|
||||
if let Some((start_idx, json_array)) = self.extract_json_array_with_pos(text) {
|
||||
// Extract normal text before BOT_TOKEN
|
||||
let normal_text_before = if start_idx > 0 {
|
||||
text[..start_idx].to_string()
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
match self.parse_json_array(json_array) {
|
||||
Ok(tools) => Ok((normal_text_before, tools)),
|
||||
Err(e) => {
|
||||
// If JSON parsing fails, return the original text as normal text
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
Ok((text.to_string(), vec![]))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Markers present but no complete array found
|
||||
Ok((text.to_string(), vec![]))
|
||||
}
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Append new text to buffer
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if current_text has tool_call
|
||||
let has_tool_start = self.has_tool_markers(current_text)
|
||||
|| (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
|
||||
|
||||
if !has_tool_start {
|
||||
// Only clear buffer if we're sure no tool call is starting
|
||||
if helpers::ends_with_partial_token(&self.buffer, self.bot_token).is_none() {
|
||||
let mut normal_text = self.buffer.clone();
|
||||
self.buffer.clear();
|
||||
|
||||
// Strip ] only once (the closing bracket of [TOOL_CALLS] array)
|
||||
// current_tool_id > 0 means we've parsed at least one tool
|
||||
if !self.array_closed
|
||||
&& self.current_tool_id > 0
|
||||
&& normal_text.starts_with(self.eot_token)
|
||||
{
|
||||
normal_text = normal_text
|
||||
.strip_prefix(self.eot_token)
|
||||
.unwrap()
|
||||
.to_string();
|
||||
self.array_closed = true;
|
||||
}
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
} else {
|
||||
// Might be partial bot_token, keep buffering
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// Build tool indices
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
// Determine start index for JSON parsing
|
||||
let start_idx = if let Some(pos) = current_text.find(self.bot_token) {
|
||||
pos + self.bot_token.len()
|
||||
} else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
|
||||
self.tool_call_separator.len()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
helpers::handle_json_tool_streaming(
|
||||
current_text,
|
||||
start_idx,
|
||||
&mut self.partial_json,
|
||||
&tool_indices,
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&mut self.prev_tool_call_arr,
|
||||
)
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains("[TOOL_CALLS]")
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::tool_parser::types::ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
helpers::reset_parser_state(
|
||||
&mut self.buffer,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
self.array_closed = false;
|
||||
}
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
/// Parser implementations for different model formats
|
||||
///
|
||||
/// This module contains concrete parser implementations for various model-specific
|
||||
/// tool/function call formats.
|
||||
// Individual parser modules
|
||||
pub mod deepseek;
|
||||
pub mod glm4_moe;
|
||||
pub mod json;
|
||||
pub mod kimik2;
|
||||
pub mod llama;
|
||||
pub mod minimax_m2;
|
||||
pub mod mistral;
|
||||
pub mod passthrough;
|
||||
pub mod pythonic;
|
||||
pub mod qwen;
|
||||
pub mod qwen_coder;
|
||||
pub mod step3;
|
||||
|
||||
// Shared helpers and utilities
|
||||
pub mod helpers;
|
||||
|
||||
// Re-export parser types for convenience
|
||||
pub use deepseek::DeepSeekParser;
|
||||
pub use glm4_moe::Glm4MoeParser;
|
||||
pub use json::JsonParser;
|
||||
pub use kimik2::KimiK2Parser;
|
||||
pub use llama::LlamaParser;
|
||||
pub use minimax_m2::MinimaxM2Parser;
|
||||
pub use mistral::MistralParser;
|
||||
pub(crate) use passthrough::PassthroughParser;
|
||||
pub use pythonic::PythonicParser;
|
||||
pub use qwen::QwenParser;
|
||||
pub use qwen_coder::QwenCoderParser;
|
||||
pub use step3::Step3Parser;
|
||||
@@ -1,55 +0,0 @@
|
||||
//! Passthrough parser that returns text unchanged
|
||||
//!
|
||||
//! This parser is used as a fallback for unknown models where no specific
|
||||
//! tool call parsing should be performed. It simply returns the input text
|
||||
//! with no tool calls detected.
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::ParserResult,
|
||||
traits::ToolParser,
|
||||
types::{StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// Passthrough parser that returns text unchanged with no tool calls
|
||||
#[derive(Default)]
|
||||
pub(crate) struct PassthroughParser;
|
||||
|
||||
impl PassthroughParser {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for PassthroughParser {
|
||||
async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Return text unchanged with no tool calls
|
||||
Ok((output.to_string(), vec![]))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
_tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Return chunk unchanged with no tool calls
|
||||
Ok(StreamingParseResult {
|
||||
normal_text: chunk.to_string(),
|
||||
calls: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, _text: &str) -> bool {
|
||||
// Passthrough never detects tool calls
|
||||
false
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
@@ -1,414 +0,0 @@
|
||||
use std::sync::OnceLock;
|
||||
|
||||
/// Pythonic format parser for tool calls
|
||||
///
|
||||
/// Handles Python function call syntax within square brackets:
|
||||
/// ```text
|
||||
/// [tool1(arg1=val1, arg2=val2), tool2(arg1=val3)]
|
||||
/// ```
|
||||
///
|
||||
/// This format is used by Llama models and uses Python literals
|
||||
/// rather than JSON for arguments.
|
||||
/// Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct?chat_template=default
|
||||
use async_trait::async_trait;
|
||||
use num_traits::ToPrimitive;
|
||||
use regex::Regex;
|
||||
use rustpython_parser::{
|
||||
ast::{Constant, Expr, Mod, UnaryOp},
|
||||
parse, Mode,
|
||||
};
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
static PYTHONIC_BLOCK_REGEX: OnceLock<Regex> = OnceLock::new();
|
||||
|
||||
/// Lazily compiled regex that locates pythonic tool call blocks.
|
||||
fn pythonic_block_regex() -> &'static Regex {
|
||||
PYTHONIC_BLOCK_REGEX.get_or_init(|| {
|
||||
// Matches one or more function calls inside a list. The `(?s)` flag allows
|
||||
// newlines inside argument lists while keeping the pattern anchored to
|
||||
// identifiers followed by parentheses, preventing plain lists like
|
||||
// `[1, 2, 3]` from matching.
|
||||
Regex::new(r"(?s)\[\s*[A-Za-z_]\w*\s*\(.*?\)\s*(?:,\s*[A-Za-z_]\w*\s*\(.*?\)\s*)*\]")
|
||||
.expect("pythonic tool call regex must compile")
|
||||
})
|
||||
}
|
||||
|
||||
/// Parser for Pythonic tool call format
|
||||
pub struct PythonicParser {
|
||||
/// Buffer for accumulating chunks
|
||||
buffer: String,
|
||||
}
|
||||
|
||||
impl Default for PythonicParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonicParser {
|
||||
/// Create a new Pythonic parser
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
buffer: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the first pythonic tool call block and return it along with the
|
||||
/// surrounding "normal" content.
|
||||
fn extract_tool_calls(&self, text: &str) -> Option<(String, String)> {
|
||||
pythonic_block_regex().find(text).map(|mat| {
|
||||
let block = mat.as_str().to_string();
|
||||
let normal = format!("{}{}", &text[..mat.start()], &text[mat.end()..]);
|
||||
(block, normal)
|
||||
})
|
||||
}
|
||||
|
||||
/// Strip special tokens that Llama models might output
|
||||
fn strip_special_tokens(text: &str) -> String {
|
||||
text.replace("<|python_start|>", "")
|
||||
.replace("<|python_end|>", "")
|
||||
}
|
||||
|
||||
fn parse_tool_call_block(&self, block: &str) -> ParserResult<Vec<ToolCall>> {
|
||||
let expr = parse_python_expression(block)?;
|
||||
match expr {
|
||||
Expr::List(list_expr) => list_expr
|
||||
.elts
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(idx, call_expr)| build_tool_call(call_expr, idx))
|
||||
.collect(),
|
||||
_ => Err(ParserError::ParsingFailed(
|
||||
"Expected a list of function calls in pythonic tool call".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for PythonicParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
let cleaned = Self::strip_special_tokens(text);
|
||||
|
||||
if let Some((tool_calls_text, normal_text)) = self.extract_tool_calls(&cleaned) {
|
||||
match self.parse_tool_call_block(&tool_calls_text) {
|
||||
Ok(calls) => {
|
||||
if calls.is_empty() {
|
||||
// No tools successfully parsed despite having markers
|
||||
Ok((text.to_string(), vec![]))
|
||||
} else {
|
||||
Ok((normal_text, calls))
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Log and return entire text as fallback
|
||||
tracing::debug!("Failed to parse pythonic tool calls: {}", e);
|
||||
Ok((text.to_string(), vec![]))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Ok((text.to_string(), vec![]))
|
||||
}
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
|
||||
let cleaned = Self::strip_special_tokens(&self.buffer);
|
||||
|
||||
// Look for opening bracket
|
||||
if let Some(start) = cleaned.find('[') {
|
||||
let normal_text = if start > 0 {
|
||||
cleaned[..start].to_string()
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
// Look for matching closing bracket
|
||||
if let Some(end) = find_matching_bracket(&cleaned, start) {
|
||||
// Found complete tool call - extract it and parse using parse_complete
|
||||
let call_text = &cleaned[start..=end];
|
||||
|
||||
match self.parse_complete(call_text).await {
|
||||
Ok((_, calls)) => {
|
||||
// Update buffer with remaining text after tool call
|
||||
let remaining_text = &cleaned[end + 1..];
|
||||
self.buffer = remaining_text.to_string();
|
||||
|
||||
// Validate tool names and convert ToolCall to ToolCallItem
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
let items: Vec<ToolCallItem> = calls
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.filter_map(|(idx, tool)| {
|
||||
if !tool_indices.contains_key(&tool.function.name) {
|
||||
tracing::debug!(
|
||||
"Invalid tool name '{}' - skipping",
|
||||
tool.function.name
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(ToolCallItem {
|
||||
tool_index: idx,
|
||||
name: Some(tool.function.name),
|
||||
parameters: tool.function.arguments,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: items,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse pythonic tool call: {}", e);
|
||||
// Clear buffer on error
|
||||
self.buffer.clear();
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// We have an opening bracket but no closing bracket yet
|
||||
// Put back everything from the bracket onwards
|
||||
self.buffer = cleaned[start..].to_string();
|
||||
|
||||
if !normal_text.is_empty() {
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Still accumulating a potential tool call
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// No tool call bracket found
|
||||
self.buffer.clear();
|
||||
Ok(StreamingParseResult {
|
||||
normal_text: cleaned,
|
||||
calls: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
let cleaned = Self::strip_special_tokens(text);
|
||||
if pythonic_block_regex().is_match(&cleaned) {
|
||||
return true;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the matching closing bracket for the opening bracket at start position.
|
||||
/// Properly handles nested brackets.
|
||||
fn find_matching_bracket(buffer: &str, start: usize) -> Option<usize> {
|
||||
let mut bracket_count = 0;
|
||||
let chars: Vec<char> = buffer.chars().collect();
|
||||
|
||||
for (i, &ch) in chars.iter().enumerate().skip(start) {
|
||||
if ch == '[' {
|
||||
bracket_count += 1;
|
||||
} else if ch == ']' {
|
||||
bracket_count -= 1;
|
||||
if bracket_count == 0 {
|
||||
return Some(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
None // No matching bracket found
|
||||
}
|
||||
|
||||
fn parse_python_expression(source: &str) -> ParserResult<Expr> {
|
||||
let module = parse(source, Mode::Expression, "<pythonic_tool_call>")
|
||||
.map_err(|err| ParserError::ParsingFailed(err.to_string()))?;
|
||||
|
||||
match module {
|
||||
Mod::Expression(expr_mod) => Ok(*expr_mod.body),
|
||||
_ => Err(ParserError::ParsingFailed(
|
||||
"Expected a Python expression".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_tool_call(expr: Expr, _index: usize) -> ParserResult<ToolCall> {
|
||||
match expr {
|
||||
Expr::Call(call_expr) => {
|
||||
if !call_expr.args.is_empty() {
|
||||
return Err(ParserError::ParsingFailed(
|
||||
"Positional arguments are not supported in pythonic tool calls".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let function_name = match *call_expr.func {
|
||||
Expr::Name(name_expr) => name_expr.id.to_string(),
|
||||
_ => {
|
||||
return Err(ParserError::ParsingFailed(
|
||||
"Unsupported function reference in pythonic tool call".to_string(),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
let mut arguments_map = Map::with_capacity(call_expr.keywords.len());
|
||||
for keyword in call_expr.keywords {
|
||||
let arg_name = keyword.arg.ok_or_else(|| {
|
||||
ParserError::ParsingFailed(
|
||||
"pythonic tool calls do not support **kwargs".to_string(),
|
||||
)
|
||||
})?;
|
||||
let value_json = expression_to_json(&keyword.value)?;
|
||||
arguments_map.insert(arg_name.to_string(), value_json);
|
||||
}
|
||||
|
||||
let arguments_json = Value::Object(arguments_map);
|
||||
let arguments_string = serde_json::to_string(&arguments_json)?;
|
||||
|
||||
Ok(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: function_name,
|
||||
arguments: arguments_string,
|
||||
},
|
||||
})
|
||||
}
|
||||
_ => Err(ParserError::ParsingFailed(
|
||||
"Expected function calls inside pythonic tool call list".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn expression_to_json(expr: &Expr) -> ParserResult<Value> {
|
||||
match expr {
|
||||
Expr::Constant(expr_constant) => constant_to_json(&expr_constant.value),
|
||||
Expr::List(list_expr) => collect_sequence(&list_expr.elts).map(Value::Array),
|
||||
Expr::Tuple(tuple_expr) => collect_sequence(&tuple_expr.elts).map(Value::Array),
|
||||
Expr::Dict(dict_expr) => {
|
||||
collect_dict(&dict_expr.keys, &dict_expr.values).map(Value::Object)
|
||||
}
|
||||
Expr::UnaryOp(unary_expr) => match unary_expr.op {
|
||||
UnaryOp::USub => match unary_expr.operand.as_ref() {
|
||||
Expr::Constant(const_expr) => negate_constant(&const_expr.value),
|
||||
_ => Err(ParserError::ParsingFailed(
|
||||
"Unsupported unary operand in pythonic tool call".to_string(),
|
||||
)),
|
||||
},
|
||||
UnaryOp::UAdd => expression_to_json(unary_expr.operand.as_ref()),
|
||||
_ => Err(ParserError::ParsingFailed(format!(
|
||||
"Unsupported unary operator in pythonic tool call: {:?}",
|
||||
unary_expr.op
|
||||
))),
|
||||
},
|
||||
Expr::Name(name_expr) => Ok(Value::String(name_expr.id.to_string())),
|
||||
_ => Err(ParserError::ParsingFailed(format!(
|
||||
"Unsupported expression in pythonic tool call: {:?}",
|
||||
expr
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn constant_to_json(constant: &Constant) -> ParserResult<Value> {
|
||||
match constant {
|
||||
Constant::None => Ok(Value::Null),
|
||||
Constant::Bool(b) => Ok(Value::Bool(*b)),
|
||||
Constant::Int(value) => Ok(integer_constant_to_value(value, false)),
|
||||
Constant::Float(f) => Number::from_f64(*f).map(Value::Number).ok_or_else(|| {
|
||||
ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
|
||||
}),
|
||||
Constant::Str(s) => Ok(Value::String(s.clone())),
|
||||
Constant::Bytes(bytes) => Ok(Value::String(String::from_utf8_lossy(bytes).into_owned())),
|
||||
Constant::Tuple(values) => constant_tuple_to_array(values).map(Value::Array),
|
||||
Constant::Ellipsis | Constant::Complex { .. } => Err(ParserError::ParsingFailed(
|
||||
"Unsupported literal in pythonic tool call".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn negate_constant(constant: &Constant) -> ParserResult<Value> {
|
||||
match constant {
|
||||
Constant::Int(value) => Ok(integer_constant_to_value(value, true)),
|
||||
Constant::Float(f) => Number::from_f64(-f).map(Value::Number).ok_or_else(|| {
|
||||
ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
|
||||
}),
|
||||
_ => Err(ParserError::ParsingFailed(
|
||||
"Unsupported unary operand in pythonic tool call".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_key_string(value: Value) -> ParserResult<String> {
|
||||
match value {
|
||||
Value::String(s) => Ok(s),
|
||||
Value::Number(num) => Ok(num.to_string()),
|
||||
Value::Bool(b) => Ok(b.to_string()),
|
||||
Value::Null => Ok("null".to_string()),
|
||||
other => Err(ParserError::ParsingFailed(format!(
|
||||
"Unsupported key type in pythonic tool call: {:?}",
|
||||
other
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_sequence(elements: &[Expr]) -> ParserResult<Vec<Value>> {
|
||||
elements.iter().map(expression_to_json).collect()
|
||||
}
|
||||
|
||||
fn collect_dict(keys: &[Option<Expr>], values: &[Expr]) -> ParserResult<Map<String, Value>> {
|
||||
let mut map = Map::with_capacity(keys.len());
|
||||
for (key_expr, value_expr) in keys.iter().zip(values.iter()) {
|
||||
let key_expr = key_expr.as_ref().ok_or_else(|| {
|
||||
ParserError::ParsingFailed("pythonic tool calls do not support **kwargs".to_string())
|
||||
})?;
|
||||
let key_value = expression_to_json(key_expr)?;
|
||||
let key = value_to_key_string(key_value)?;
|
||||
let value_json = expression_to_json(value_expr)?;
|
||||
map.insert(key, value_json);
|
||||
}
|
||||
Ok(map)
|
||||
}
|
||||
|
||||
fn constant_tuple_to_array(values: &[Constant]) -> ParserResult<Vec<Value>> {
|
||||
values.iter().map(constant_to_json).collect()
|
||||
}
|
||||
|
||||
fn integer_constant_to_value<T>(value: &T, negate: bool) -> Value
|
||||
where
|
||||
T: ToPrimitive + std::fmt::Display,
|
||||
{
|
||||
if let Some(mut i) = value.to_i64() {
|
||||
if negate {
|
||||
i = -i;
|
||||
}
|
||||
return Value::Number(Number::from(i));
|
||||
}
|
||||
|
||||
if negate {
|
||||
if let Some(u) = value.to_u64() {
|
||||
if u <= i64::MAX as u64 {
|
||||
return Value::Number(Number::from(-(u as i64)));
|
||||
}
|
||||
return Value::String(format!("-{}", value));
|
||||
}
|
||||
Value::String(format!("-{}", value))
|
||||
} else if let Some(u) = value.to_u64() {
|
||||
Value::Number(Number::from(u))
|
||||
} else {
|
||||
Value::String(value.to_string())
|
||||
}
|
||||
}
|
||||
@@ -1,258 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
partial_json::PartialJson,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall},
|
||||
},
|
||||
};
|
||||
|
||||
/// Qwen format parser for tool calls
|
||||
///
|
||||
/// Handles the Qwen 2.5/3 specific format:
|
||||
/// `<tool_call>\n{"name": "func", "arguments": {...}}\n</tool_call>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Tool Call Tags: `<tool_call>` and `</tool_call>` wrap each individual call
|
||||
/// - Each individual call is separated by `\n`
|
||||
/// - Function Call Object: JSON object with "name" and "arguments" fields
|
||||
///
|
||||
/// Reference: https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct?chat_template=default
|
||||
pub struct QwenParser {
|
||||
/// Parser for handling incomplete JSON during streaming
|
||||
partial_json: PartialJson,
|
||||
|
||||
/// Regex for extracting tool calls in parse_complete
|
||||
extractor: Regex,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Buffer for normal text that might precede partial end tokens
|
||||
normal_text_buffer: String,
|
||||
|
||||
/// Token configuration
|
||||
/// Start/end tokens for each individual tool call (not the entire sequence)
|
||||
individual_tool_start_token: &'static str,
|
||||
individual_tool_end_token: &'static str,
|
||||
tool_call_separator: &'static str,
|
||||
}
|
||||
|
||||
impl QwenParser {
|
||||
/// Create a new Qwen parser
|
||||
pub fn new() -> Self {
|
||||
// Use (?s) flag for DOTALL mode to handle newlines
|
||||
let pattern = r"(?s)<tool_call>\n(.*?)\n</tool_call>";
|
||||
let extractor = Regex::new(pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
partial_json: PartialJson::default(),
|
||||
extractor,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
normal_text_buffer: String::new(),
|
||||
individual_tool_start_token: "<tool_call>\n",
|
||||
individual_tool_end_token: "\n</tool_call>",
|
||||
tool_call_separator: "\n",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a single JSON object into a ToolCall
|
||||
fn parse_single_object(&self, obj: &Value) -> ParserResult<Option<ToolCall>> {
|
||||
let name = obj.get("name").and_then(|v| v.as_str());
|
||||
|
||||
if let Some(name) = name {
|
||||
// Get arguments - Qwen uses "arguments" key
|
||||
let empty_obj = Value::Object(serde_json::Map::new());
|
||||
let args = obj.get("arguments").unwrap_or(&empty_obj);
|
||||
|
||||
// Convert arguments to JSON string
|
||||
let arguments = serde_json::to_string(args)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for QwenParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for QwenParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Check if text contains Qwen format
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where the first tool call begins
|
||||
let idx = text.find("<tool_call>").unwrap(); // Safe because has_tool_markers checked
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Extract tool calls
|
||||
let mut tools = Vec::new();
|
||||
for captures in self.extractor.captures_iter(text) {
|
||||
if let Some(json_str) = captures.get(1) {
|
||||
let parsed = serde_json::from_str::<Value>(json_str.as_str().trim())
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))
|
||||
.and_then(|v| self.parse_single_object(&v));
|
||||
|
||||
match parsed {
|
||||
Ok(Some(tool)) => tools.push(tool),
|
||||
Ok(None) => continue,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse tool call: {:?}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
// Append new text to buffer
|
||||
self.buffer.push_str(chunk);
|
||||
let current_text = &self.buffer.clone();
|
||||
|
||||
// Check if current_text has tool_call
|
||||
let has_tool_start = self.has_tool_markers(current_text)
|
||||
|| (self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator));
|
||||
|
||||
if !has_tool_start {
|
||||
// Only clear buffer if we're sure no tool call is starting
|
||||
if helpers::ends_with_partial_token(&self.buffer, self.individual_tool_start_token)
|
||||
.is_none()
|
||||
{
|
||||
let normal_text = self.buffer.clone();
|
||||
self.buffer.clear();
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
} else {
|
||||
// Might be partial individual_tool_start_token, keep buffering
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// Build tool indices
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
// Determine start index for JSON parsing
|
||||
let start_idx = if let Some(pos) = current_text.find(self.individual_tool_start_token) {
|
||||
pos + self.individual_tool_start_token.len()
|
||||
} else if self.current_tool_id > 0 && current_text.starts_with(self.tool_call_separator) {
|
||||
self.tool_call_separator.len()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let mut result = helpers::handle_json_tool_streaming(
|
||||
current_text,
|
||||
start_idx,
|
||||
&mut self.partial_json,
|
||||
&tool_indices,
|
||||
&mut self.buffer,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
&mut self.prev_tool_call_arr,
|
||||
)?;
|
||||
|
||||
// Qwen-specific: Handle partial end tokens in normal text
|
||||
// After tool calls complete, normal text might contain partial "</tool_call>" tags
|
||||
if !result.normal_text.is_empty() {
|
||||
self.normal_text_buffer.push_str(&result.normal_text);
|
||||
|
||||
// Check if buffer contains complete end token (without leading newline)
|
||||
let end_token_without_newline = &self.individual_tool_end_token[1..]; // "</tool_call>"
|
||||
if self.normal_text_buffer.contains(end_token_without_newline) {
|
||||
// Complete end token found - clean it and return
|
||||
let cleaned_text = self
|
||||
.normal_text_buffer
|
||||
.replace(end_token_without_newline, "");
|
||||
self.normal_text_buffer.clear();
|
||||
result.normal_text = cleaned_text;
|
||||
} else {
|
||||
// Check if buffer might contain partial end token at the end
|
||||
if let Some(partial_match_len) = helpers::ends_with_partial_token(
|
||||
&self.normal_text_buffer,
|
||||
end_token_without_newline,
|
||||
) {
|
||||
// Keep potential partial match in buffer, return the rest
|
||||
let split_point = self.normal_text_buffer.len() - partial_match_len;
|
||||
result.normal_text = self.normal_text_buffer[..split_point].to_string();
|
||||
self.normal_text_buffer = self.normal_text_buffer[split_point..].to_string();
|
||||
} else {
|
||||
// No partial match, return all buffered text
|
||||
result.normal_text = self.normal_text_buffer.clone();
|
||||
self.normal_text_buffer.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains("<tool_call>")
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::tool_parser::types::ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
helpers::reset_parser_state(
|
||||
&mut self.buffer,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,587 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// Qwen Coder format parser for tool calls
|
||||
///
|
||||
/// Handles the Qwen Coder specific XML format:
|
||||
/// `<tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Tool Call Tags: `<tool_call>` and `</tool_call>` wrap each individual call
|
||||
/// - XML-style function declaration: `<function=name>`
|
||||
/// - XML-style parameters: `<parameter=key>value</parameter>`
|
||||
///
|
||||
/// Reference: https://huggingface.co/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8?chat_template=default
|
||||
pub struct QwenCoderParser {
|
||||
/// Regex for extracting tool calls in parse_complete
|
||||
extractor: Regex,
|
||||
|
||||
/// Buffer for accumulating incomplete patterns across chunks
|
||||
buffer: String,
|
||||
|
||||
/// Stores complete tool call info (name and arguments) for each tool being parsed
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
|
||||
/// Index of currently streaming tool call (-1 means no active tool)
|
||||
current_tool_id: i32,
|
||||
|
||||
/// Flag for whether current tool's name has been sent to client
|
||||
current_tool_name_sent: bool,
|
||||
|
||||
/// Tracks raw JSON string content streamed to client for each tool's arguments
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
|
||||
/// Token configuration
|
||||
tool_call_start_token: &'static str,
|
||||
tool_call_end_token: &'static str,
|
||||
|
||||
/// XML format streaming state
|
||||
in_tool_call: bool,
|
||||
current_function_name: String,
|
||||
current_parameters: serde_json::Map<String, Value>,
|
||||
|
||||
/// Precompiled regex patterns for XML format parsing
|
||||
xml_function_pattern: Regex,
|
||||
xml_param_pattern: Regex,
|
||||
}
|
||||
|
||||
/// Decode HTML entities in a string (equivalent to Python's html.unescape)
|
||||
///
|
||||
/// Handles common HTML entities like & < > " ' and numeric entities
|
||||
fn html_unescape(s: &str) -> String {
|
||||
let mut result = String::with_capacity(s.len());
|
||||
let mut chars = s.chars().peekable();
|
||||
|
||||
while let Some(c) = chars.next() {
|
||||
if c == '&' {
|
||||
let mut entity = String::new();
|
||||
let mut consumed_semicolon = false;
|
||||
while let Some(&next) = chars.peek() {
|
||||
if next == ';' {
|
||||
chars.next();
|
||||
consumed_semicolon = true;
|
||||
break;
|
||||
}
|
||||
if next.is_alphanumeric() || next == '#' {
|
||||
entity.push(chars.next().unwrap());
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let decoded = match entity.as_str() {
|
||||
"amp" => "&",
|
||||
"lt" => "<",
|
||||
"gt" => ">",
|
||||
"quot" => "\"",
|
||||
"apos" => "'",
|
||||
"nbsp" => "\u{00A0}",
|
||||
s if s.starts_with('#') => {
|
||||
let num_str = &s[1..];
|
||||
let code_point = if num_str.starts_with('x') || num_str.starts_with('X') {
|
||||
u32::from_str_radix(&num_str[1..], 16).ok()
|
||||
} else {
|
||||
num_str.parse::<u32>().ok()
|
||||
};
|
||||
if let Some(cp) = code_point {
|
||||
if let Some(ch) = char::from_u32(cp) {
|
||||
result.push(ch);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
// Invalid numeric entity, reconstruct original
|
||||
result.push('&');
|
||||
result.push_str(&entity);
|
||||
if consumed_semicolon {
|
||||
result.push(';');
|
||||
}
|
||||
continue;
|
||||
}
|
||||
_ => {
|
||||
// Unknown entity, reconstruct original
|
||||
result.push('&');
|
||||
result.push_str(&entity);
|
||||
if consumed_semicolon {
|
||||
result.push(';');
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
result.push_str(decoded);
|
||||
} else {
|
||||
result.push(c);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Parse a raw parameter value, similar to Python's _safe_val
|
||||
///
|
||||
/// 1. Decode HTML entities
|
||||
/// 2. Try to parse as JSON (numbers, booleans, null, objects, arrays)
|
||||
/// 3. Fall back to string if JSON parsing fails
|
||||
fn safe_val(raw: &str) -> Value {
|
||||
let unescaped = html_unescape(raw.trim());
|
||||
|
||||
// Try JSON parsing first
|
||||
if let Ok(v) = serde_json::from_str::<Value>(&unescaped) {
|
||||
return v;
|
||||
}
|
||||
|
||||
// Handle Python-style literals (True, False, None)
|
||||
match unescaped.as_str() {
|
||||
"True" => return Value::Bool(true),
|
||||
"False" => return Value::Bool(false),
|
||||
"None" => return Value::Null,
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Fall back to string
|
||||
Value::String(unescaped)
|
||||
}
|
||||
|
||||
impl QwenCoderParser {
|
||||
/// Create a new Qwen Coder parser
|
||||
pub fn new() -> Self {
|
||||
// Support XML format: <tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>
|
||||
let pattern = r"(?s)<tool_call>\s*(.*?)\s*</tool_call>";
|
||||
let extractor = Regex::new(pattern).expect("Valid regex pattern");
|
||||
|
||||
// Precompile XML format regex patterns for performance
|
||||
let xml_function_pattern =
|
||||
Regex::new(r"<function=([^>]+)>").expect("Valid XML function pattern");
|
||||
let xml_param_pattern = Regex::new(r"(?s)<parameter=([^>]+)>(.*?)</parameter>")
|
||||
.expect("Valid XML parameter pattern");
|
||||
|
||||
Self {
|
||||
extractor,
|
||||
buffer: String::new(),
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
current_tool_name_sent: false,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
tool_call_start_token: "<tool_call>",
|
||||
tool_call_end_token: "</tool_call>",
|
||||
in_tool_call: false,
|
||||
current_function_name: String::new(),
|
||||
current_parameters: serde_json::Map::new(),
|
||||
xml_function_pattern,
|
||||
xml_param_pattern,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse XML format tool call: <function=name><parameter=key>value</parameter></function>
|
||||
fn parse_xml_format(&self, content: &str) -> ParserResult<Option<ToolCall>> {
|
||||
let function_captures = self
|
||||
.xml_function_pattern
|
||||
.captures(content)
|
||||
.ok_or_else(|| ParserError::ParsingFailed("No function name found".to_string()))?;
|
||||
|
||||
let function_name = function_captures
|
||||
.get(1)
|
||||
.ok_or_else(|| ParserError::ParsingFailed("Function name capture failed".to_string()))?
|
||||
.as_str()
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
if function_name.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut parameters = serde_json::Map::new();
|
||||
|
||||
for cap in self.xml_param_pattern.captures_iter(content) {
|
||||
if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
|
||||
let key = key_match.as_str().trim().to_string();
|
||||
let value = value_match.as_str();
|
||||
let json_value = safe_val(value);
|
||||
parameters.insert(key, json_value);
|
||||
}
|
||||
}
|
||||
|
||||
let arguments = serde_json::to_string(¶meters)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: function_name,
|
||||
arguments,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
/// Parse and stream complete parameters from buffer
|
||||
/// Returns tool call items to emit (similar to Python's _parse_and_stream_parameters)
|
||||
fn parse_and_stream_parameters(&mut self) -> ParserResult<Vec<ToolCallItem>> {
|
||||
let mut calls: Vec<ToolCallItem> = vec![];
|
||||
|
||||
// Find all complete parameter patterns in buffer
|
||||
let mut new_params = serde_json::Map::new();
|
||||
for cap in self.xml_param_pattern.captures_iter(&self.buffer) {
|
||||
if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
|
||||
let key = key_match.as_str().trim().to_string();
|
||||
let value = value_match.as_str();
|
||||
let json_value = safe_val(value);
|
||||
new_params.insert(key, json_value);
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate parameter diff and stream updates
|
||||
if new_params != self.current_parameters {
|
||||
let current_args = &mut self.streamed_args_for_tool[self.current_tool_id as usize];
|
||||
|
||||
if self.current_parameters.is_empty() {
|
||||
// First parameter(s) - build JSON fragment (without closing brace)
|
||||
let mut items = Vec::new();
|
||||
for (key, value) in &new_params {
|
||||
let key_json =
|
||||
serde_json::to_string(key).unwrap_or_else(|_| format!("\"{}\"", key));
|
||||
let value_json = serde_json::to_string(value).unwrap_or_default();
|
||||
items.push(format!("{}: {}", key_json, value_json));
|
||||
}
|
||||
let json_fragment = format!("{{{}", items.join(", "));
|
||||
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: json_fragment.clone(),
|
||||
});
|
||||
*current_args = json_fragment;
|
||||
} else {
|
||||
// Additional parameters - add them incrementally
|
||||
let new_keys: Vec<_> = new_params
|
||||
.keys()
|
||||
.filter(|k| !self.current_parameters.contains_key(*k))
|
||||
.collect();
|
||||
|
||||
if !new_keys.is_empty() {
|
||||
let mut continuation_parts = Vec::new();
|
||||
for key in new_keys {
|
||||
if let Some(value) = new_params.get(key) {
|
||||
let key_json = serde_json::to_string(key)
|
||||
.unwrap_or_else(|_| format!("\"{}\"", key));
|
||||
let value_json = serde_json::to_string(value).unwrap_or_default();
|
||||
continuation_parts.push(format!("{}: {}", key_json, value_json));
|
||||
}
|
||||
}
|
||||
|
||||
let json_fragment = format!(", {}", continuation_parts.join(", "));
|
||||
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: json_fragment.clone(),
|
||||
});
|
||||
current_args.push_str(&json_fragment);
|
||||
}
|
||||
}
|
||||
|
||||
// Update current state
|
||||
self.current_parameters = new_params.clone();
|
||||
if let Some(tool_obj) =
|
||||
self.prev_tool_call_arr[self.current_tool_id as usize].as_object_mut()
|
||||
{
|
||||
tool_obj.insert("arguments".to_string(), Value::Object(new_params));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(calls)
|
||||
}
|
||||
|
||||
/// Reset streaming state for next tool call
|
||||
fn reset_streaming_state(&mut self) {
|
||||
self.in_tool_call = false;
|
||||
self.current_tool_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for QwenCoderParser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for QwenCoderParser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
// Check if text contains Qwen Coder format
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where the first tool call begins
|
||||
let idx = text.find(self.tool_call_start_token).unwrap();
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Extract tool calls
|
||||
let mut tools = Vec::new();
|
||||
for captures in self.extractor.captures_iter(text) {
|
||||
if let Some(content_str) = captures.get(1) {
|
||||
let content = content_str.as_str().trim();
|
||||
|
||||
match self.parse_xml_format(content) {
|
||||
Ok(Some(tool)) => tools.push(tool),
|
||||
Ok(None) => continue,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to parse XML tool call: {:?}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
|
||||
let mut normal_text = String::new();
|
||||
let mut calls: Vec<ToolCallItem> = vec![];
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
loop {
|
||||
// If we're not in a tool call and don't see a start token, return normal text
|
||||
if !self.in_tool_call && !self.buffer.contains(self.tool_call_start_token) {
|
||||
// Check for partial start token
|
||||
if helpers::ends_with_partial_token(&self.buffer, self.tool_call_start_token)
|
||||
.is_none()
|
||||
{
|
||||
normal_text.push_str(&self.buffer);
|
||||
self.buffer.clear();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
// Look for tool call start
|
||||
if !self.in_tool_call {
|
||||
if let Some(s) = self.buffer.find(self.tool_call_start_token) {
|
||||
normal_text.push_str(&self.buffer[..s]);
|
||||
self.buffer = self.buffer[s + self.tool_call_start_token.len()..].to_string();
|
||||
self.in_tool_call = true;
|
||||
self.current_tool_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
continue;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// We're in a tool call, try to parse function name if not sent yet
|
||||
if !self.current_tool_name_sent {
|
||||
if let Some(captures) = self.xml_function_pattern.captures(&self.buffer) {
|
||||
if let Some(name_match) = captures.get(1) {
|
||||
let function_name = name_match.as_str().trim().to_string();
|
||||
|
||||
// Validate function name
|
||||
if tool_indices.contains_key(&function_name) {
|
||||
self.current_function_name = function_name.clone();
|
||||
self.current_tool_name_sent = true;
|
||||
|
||||
// Initialize tool call tracking
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
}
|
||||
|
||||
// Ensure tracking arrays are large enough
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Store tool call info
|
||||
self.prev_tool_call_arr[self.current_tool_id as usize] = serde_json::json!({
|
||||
"name": function_name,
|
||||
"arguments": {}
|
||||
});
|
||||
|
||||
// Send tool name
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: Some(function_name),
|
||||
parameters: String::new(),
|
||||
});
|
||||
|
||||
// Remove processed function declaration from buffer
|
||||
self.buffer = self.buffer[captures.get(0).unwrap().end()..].to_string();
|
||||
continue;
|
||||
} else {
|
||||
// Invalid function name, reset state
|
||||
tracing::warn!("Invalid function name: {}", function_name);
|
||||
self.reset_streaming_state();
|
||||
normal_text.push_str(&self.buffer);
|
||||
self.buffer.clear();
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Function name not complete yet, wait for more text
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Parse parameters (only complete ones)
|
||||
if self.current_tool_name_sent {
|
||||
let param_calls = self.parse_and_stream_parameters()?;
|
||||
calls.extend(param_calls);
|
||||
|
||||
// Check if tool call is complete
|
||||
if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
|
||||
// Close JSON object if we have parameters
|
||||
let current_args = &self.streamed_args_for_tool[self.current_tool_id as usize];
|
||||
if !current_args.is_empty() {
|
||||
// Count braces to check if JSON is complete
|
||||
let open_braces = current_args.matches('{').count();
|
||||
let close_braces = current_args.matches('}').count();
|
||||
if open_braces > close_braces {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: "}".to_string(),
|
||||
});
|
||||
self.streamed_args_for_tool[self.current_tool_id as usize].push('}');
|
||||
}
|
||||
}
|
||||
|
||||
// Complete the tool call
|
||||
self.buffer =
|
||||
self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
|
||||
self.reset_streaming_state();
|
||||
self.current_tool_id += 1;
|
||||
continue;
|
||||
} else {
|
||||
// Tool call not complete yet, wait for more text
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult { normal_text, calls })
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains(self.tool_call_start_token)
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
helpers::reset_parser_state(
|
||||
&mut self.buffer,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.current_tool_id,
|
||||
&mut self.current_tool_name_sent,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
self.reset_streaming_state();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_html_unescape_basic() {
|
||||
assert_eq!(html_unescape("&"), "&");
|
||||
assert_eq!(html_unescape("<"), "<");
|
||||
assert_eq!(html_unescape(">"), ">");
|
||||
assert_eq!(html_unescape("""), "\"");
|
||||
assert_eq!(html_unescape("'"), "'");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_html_unescape_numeric() {
|
||||
assert_eq!(html_unescape("<"), "<");
|
||||
assert_eq!(html_unescape("<"), "<");
|
||||
assert_eq!(html_unescape("<"), "<");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_html_unescape_mixed() {
|
||||
assert_eq!(
|
||||
html_unescape("Hello & World <tag>"),
|
||||
"Hello & World <tag>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_html_unescape_unknown() {
|
||||
// Unknown entities with semicolon should be preserved as-is
|
||||
assert_eq!(html_unescape("&unknown;"), "&unknown;");
|
||||
// Unterminated entities should NOT have semicolon added
|
||||
assert_eq!(html_unescape("&foo bar"), "&foo bar");
|
||||
assert_eq!(html_unescape("&"), "&");
|
||||
assert_eq!(html_unescape("& "), "& ");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_val_json() {
|
||||
assert_eq!(safe_val("42"), Value::Number(42.into()));
|
||||
assert_eq!(safe_val("1.5"), serde_json::json!(1.5));
|
||||
assert_eq!(safe_val("true"), Value::Bool(true));
|
||||
assert_eq!(safe_val("false"), Value::Bool(false));
|
||||
assert_eq!(safe_val("null"), Value::Null);
|
||||
assert_eq!(
|
||||
safe_val(r#"{"key": "value"}"#),
|
||||
serde_json::json!({"key": "value"})
|
||||
);
|
||||
assert_eq!(safe_val(r#"[1, 2, 3]"#), serde_json::json!([1, 2, 3]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_val_python_literals() {
|
||||
assert_eq!(safe_val("True"), Value::Bool(true));
|
||||
assert_eq!(safe_val("False"), Value::Bool(false));
|
||||
assert_eq!(safe_val("None"), Value::Null);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_val_string_fallback() {
|
||||
assert_eq!(
|
||||
safe_val("hello world"),
|
||||
Value::String("hello world".to_string())
|
||||
);
|
||||
assert_eq!(safe_val(" spaces "), Value::String("spaces".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_val_html_entities() {
|
||||
assert_eq!(safe_val("<div>"), Value::String("<div>".to_string()));
|
||||
assert_eq!(
|
||||
safe_val("Tom & Jerry"),
|
||||
Value::String("Tom & Jerry".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,576 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
parsers::helpers,
|
||||
traits::ToolParser,
|
||||
types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
|
||||
},
|
||||
};
|
||||
|
||||
/// Step3 format parser for tool calls
|
||||
///
|
||||
/// Handles the Step3 specific format with steptml XML:
|
||||
/// `<｜tool_calls_begin｜><｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="{name}"><steptml:parameter name="{k}">{v}</steptml:parameter></steptml:invoke><｜tool_call_end｜><｜tool_calls_end｜>`
|
||||
///
|
||||
/// Features:
|
||||
/// - Unicode token delimiters
|
||||
/// - StepTML XML format for invocations
|
||||
/// - Support for multiple sequential tool calls
|
||||
pub struct Step3Parser {
|
||||
/// Regex for extracting tool call blocks
|
||||
tool_call_extractor: Regex,
|
||||
/// Regex for extracting steptml invocations
|
||||
invoke_extractor: Regex,
|
||||
/// Regex for extracting parameters
|
||||
param_extractor: Regex,
|
||||
|
||||
/// Buffer for accumulating chunks
|
||||
buffer: String,
|
||||
|
||||
/// Token configuration
|
||||
bot_token: &'static str,
|
||||
eot_token: &'static str,
|
||||
tool_call_begin: &'static str,
|
||||
tool_call_end: &'static str,
|
||||
tool_sep: &'static str,
|
||||
|
||||
/// Streaming state variables (mirrors Python's Step3Detector)
|
||||
in_tool_block: bool,
|
||||
tool_block_finished: bool,
|
||||
current_function_name: String,
|
||||
current_parameters: serde_json::Map<String, Value>,
|
||||
in_tool_call: bool,
|
||||
function_name_sent: bool,
|
||||
|
||||
/// Standard state machine fields
|
||||
prev_tool_call_arr: Vec<Value>,
|
||||
current_tool_id: i32,
|
||||
streamed_args_for_tool: Vec<String>,
|
||||
}
|
||||
|
||||
impl Step3Parser {
|
||||
/// Create a new Step3 parser
|
||||
pub fn new() -> Self {
|
||||
// Pattern for individual tool calls
|
||||
let tool_call_pattern = r"(?s)<｜tool_call_begin｜>.*?<｜tool_call_end｜>";
|
||||
let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Pattern for steptml invocations
|
||||
let invoke_pattern = r#"(?s)<steptml:invoke name="([^"]+)">(.+?)</steptml:invoke>"#;
|
||||
let invoke_extractor = Regex::new(invoke_pattern).expect("Valid regex pattern");
|
||||
|
||||
// Pattern for steptml parameters - using non-greedy match for values to handle < characters
|
||||
let param_pattern = r#"(?s)<steptml:parameter name="([^"]+)">(.+?)</steptml:parameter>"#;
|
||||
let param_extractor = Regex::new(param_pattern).expect("Valid regex pattern");
|
||||
|
||||
Self {
|
||||
tool_call_extractor,
|
||||
invoke_extractor,
|
||||
param_extractor,
|
||||
|
||||
buffer: String::new(),
|
||||
|
||||
bot_token: "<｜tool_calls_begin｜>",
|
||||
eot_token: "<｜tool_calls_end｜>",
|
||||
tool_call_begin: "<｜tool_call_begin｜>",
|
||||
tool_call_end: "<｜tool_call_end｜>",
|
||||
tool_sep: "<｜tool_sep｜>",
|
||||
|
||||
// Streaming state variables
|
||||
in_tool_block: false,
|
||||
tool_block_finished: false,
|
||||
current_function_name: String::new(),
|
||||
current_parameters: serde_json::Map::new(),
|
||||
in_tool_call: false,
|
||||
function_name_sent: false,
|
||||
|
||||
// Standard state machine fields
|
||||
prev_tool_call_arr: Vec::new(),
|
||||
current_tool_id: -1,
|
||||
streamed_args_for_tool: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset streaming state for the next tool call
|
||||
fn reset_streaming_state(&mut self) {
|
||||
self.in_tool_call = false;
|
||||
self.function_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
}
|
||||
|
||||
/// Parse partial tool call for streaming scenarios (mirrors Python's _parse_partial_tool_call)
|
||||
fn parse_partial_tool_call(
|
||||
&mut self,
|
||||
tool_indices: &HashMap<String, usize>,
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
let mut calls = Vec::new();
|
||||
|
||||
// Check if we have tool_sep (means we're past the type declaration)
|
||||
if !self.buffer.contains(self.tool_sep) {
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
});
|
||||
}
|
||||
|
||||
// Clone the buffer to avoid borrow conflicts
|
||||
let buffer_clone = self.buffer.clone();
|
||||
let parts: Vec<&str> = buffer_clone.splitn(2, self.tool_sep).collect();
|
||||
if parts.len() != 2 {
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
});
|
||||
}
|
||||
|
||||
let type_part = parts[0].trim();
|
||||
let invoke_part = parts[1];
|
||||
|
||||
// Check if it's a function type
|
||||
if type_part != "function" {
|
||||
// Invalid tool type, skip this tool call
|
||||
self.reset_streaming_state();
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
});
|
||||
}
|
||||
|
||||
// Try to extract function name if not sent yet
|
||||
if !self.function_name_sent {
|
||||
if let Some(captures) = self.invoke_extractor.captures(invoke_part) {
|
||||
let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Validate function name
|
||||
if tool_indices.contains_key(func_name) {
|
||||
self.current_function_name = func_name.to_string();
|
||||
self.function_name_sent = true;
|
||||
|
||||
// Initialize tool tracking
|
||||
if self.current_tool_id == -1 {
|
||||
self.current_tool_id = 0;
|
||||
}
|
||||
|
||||
// Ensure tracking arrays are large enough
|
||||
helpers::ensure_capacity(
|
||||
self.current_tool_id,
|
||||
&mut self.prev_tool_call_arr,
|
||||
&mut self.streamed_args_for_tool,
|
||||
);
|
||||
|
||||
// Store tool call info
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
self.prev_tool_call_arr[tool_id] = serde_json::json!({
|
||||
"name": func_name,
|
||||
"arguments": {},
|
||||
});
|
||||
|
||||
// Send tool name with empty parameters
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: Some(func_name.to_string()),
|
||||
parameters: String::new(),
|
||||
});
|
||||
} else {
|
||||
// Invalid function name
|
||||
tracing::debug!("Invalid function name: {}", func_name);
|
||||
self.reset_streaming_state();
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
});
|
||||
}
|
||||
} else {
|
||||
// Function name not complete yet
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Parse parameters incrementally
|
||||
if self.function_name_sent {
|
||||
// Extract all complete parameters
|
||||
let mut new_params = serde_json::Map::new();
|
||||
for capture in self.param_extractor.captures_iter(invoke_part) {
|
||||
let param_name = capture.get(1).map_or("", |m| m.as_str()).trim();
|
||||
let param_value_str = capture.get(2).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Try to parse the value as JSON first, fallback to string
|
||||
let param_value =
|
||||
if let Ok(json_val) = serde_json::from_str::<Value>(param_value_str) {
|
||||
json_val
|
||||
} else {
|
||||
// Try parsing as Python literal
|
||||
if param_value_str == "true" || param_value_str == "True" {
|
||||
Value::Bool(true)
|
||||
} else if param_value_str == "false" || param_value_str == "False" {
|
||||
Value::Bool(false)
|
||||
} else if param_value_str == "null" || param_value_str == "None" {
|
||||
Value::Null
|
||||
} else if let Ok(num) = param_value_str.parse::<i64>() {
|
||||
Value::Number(num.into())
|
||||
} else if let Ok(num) = param_value_str.parse::<f64>() {
|
||||
if let Some(n) = serde_json::Number::from_f64(num) {
|
||||
Value::Number(n)
|
||||
} else {
|
||||
Value::String(param_value_str.to_string())
|
||||
}
|
||||
} else {
|
||||
Value::String(param_value_str.to_string())
|
||||
}
|
||||
};
|
||||
|
||||
new_params.insert(param_name.to_string(), param_value);
|
||||
}
|
||||
|
||||
// Check if we have new parameters to stream
|
||||
if new_params != self.current_parameters {
|
||||
// Build the JSON content without the closing brace for streaming
|
||||
let diff = if self.current_parameters.is_empty() {
|
||||
// First parameters - send opening brace and content
|
||||
let params_content =
|
||||
serde_json::to_string(&new_params).unwrap_or_else(|_| "{}".to_string());
|
||||
if params_content.len() > 2 {
|
||||
// Send everything except the closing brace
|
||||
params_content[..params_content.len() - 1].to_string()
|
||||
} else {
|
||||
"{".to_string()
|
||||
}
|
||||
} else {
|
||||
// Subsequent parameters - calculate the incremental diff
|
||||
let old_json = serde_json::to_string(&self.current_parameters)
|
||||
.unwrap_or_else(|_| "{}".to_string());
|
||||
let new_json =
|
||||
serde_json::to_string(&new_params).unwrap_or_else(|_| "{}".to_string());
|
||||
|
||||
// Remove closing braces for comparison
|
||||
let old_without_brace = &old_json[..old_json.len() - 1];
|
||||
let new_without_brace = &new_json[..new_json.len() - 1];
|
||||
|
||||
// The new content should extend the old content
|
||||
new_without_brace
|
||||
.strip_prefix(old_without_brace)
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_default()
|
||||
};
|
||||
|
||||
if !diff.is_empty() {
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: diff.clone(),
|
||||
});
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.streamed_args_for_tool.len() {
|
||||
self.streamed_args_for_tool[tool_id].push_str(&diff);
|
||||
}
|
||||
}
|
||||
|
||||
// Update current state
|
||||
self.current_parameters = new_params.clone();
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.prev_tool_call_arr.len() {
|
||||
if let Some(obj) = self.prev_tool_call_arr[tool_id].as_object_mut() {
|
||||
obj.insert("arguments".to_string(), Value::Object(new_params));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if tool call is complete
|
||||
if self.buffer.contains(self.tool_call_end) {
|
||||
// Send closing brace if we've sent any parameters
|
||||
let tool_id = self.current_tool_id as usize;
|
||||
if tool_id < self.streamed_args_for_tool.len()
|
||||
&& !self.streamed_args_for_tool[tool_id].is_empty()
|
||||
{
|
||||
calls.push(ToolCallItem {
|
||||
tool_index: self.current_tool_id as usize,
|
||||
name: None,
|
||||
parameters: "}".to_string(),
|
||||
});
|
||||
self.streamed_args_for_tool[tool_id].push('}');
|
||||
}
|
||||
|
||||
// Find the end position
|
||||
if let Some(end_idx) = self.buffer.find(self.tool_call_end) {
|
||||
// Remove the processed tool call from buffer
|
||||
self.buffer = self.buffer[end_idx + self.tool_call_end.len()..].to_string();
|
||||
}
|
||||
|
||||
// Reset state for next tool call
|
||||
self.reset_streaming_state();
|
||||
self.current_tool_id += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult {
|
||||
normal_text: String::new(),
|
||||
calls,
|
||||
})
|
||||
}
|
||||
|
||||
/// Parse parameters from steptml format
|
||||
fn parse_steptml_parameters(
|
||||
&self,
|
||||
params_text: &str,
|
||||
) -> ParserResult<serde_json::Map<String, Value>> {
|
||||
let mut parameters = serde_json::Map::new();
|
||||
|
||||
for capture in self.param_extractor.captures_iter(params_text) {
|
||||
let param_name = capture.get(1).map_or("", |m| m.as_str()).trim();
|
||||
let param_value_str = capture.get(2).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Try to parse the value as JSON first, fallback to string
|
||||
let param_value = if let Ok(json_val) = serde_json::from_str::<Value>(param_value_str) {
|
||||
json_val
|
||||
} else {
|
||||
// Try parsing as Python literal
|
||||
if param_value_str == "true" || param_value_str == "True" {
|
||||
Value::Bool(true)
|
||||
} else if param_value_str == "false" || param_value_str == "False" {
|
||||
Value::Bool(false)
|
||||
} else if param_value_str == "null" || param_value_str == "None" {
|
||||
Value::Null
|
||||
} else if let Ok(num) = param_value_str.parse::<i64>() {
|
||||
Value::Number(num.into())
|
||||
} else if let Ok(num) = param_value_str.parse::<f64>() {
|
||||
if let Some(n) = serde_json::Number::from_f64(num) {
|
||||
Value::Number(n)
|
||||
} else {
|
||||
Value::String(param_value_str.to_string())
|
||||
}
|
||||
} else {
|
||||
Value::String(param_value_str.to_string())
|
||||
}
|
||||
};
|
||||
|
||||
parameters.insert(param_name.to_string(), param_value);
|
||||
}
|
||||
|
||||
Ok(parameters)
|
||||
}
|
||||
|
||||
/// Parse a single tool call block
|
||||
fn parse_tool_call(&self, block: &str) -> ParserResult<Option<ToolCall>> {
|
||||
// Check if it contains function marker and tool separator
|
||||
if !block.contains("function") || !block.contains("<｜tool_sep｜>") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Split by tool separator
|
||||
let parts: Vec<&str> = block.split("<｜tool_sep｜>").collect();
|
||||
if parts.len() != 2 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Check if it's a function type
|
||||
if !parts[0].contains("function") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let invoke_part = parts[1];
|
||||
|
||||
// Extract steptml invoke
|
||||
if let Some(captures) = self.invoke_extractor.captures(invoke_part) {
|
||||
let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
|
||||
|
||||
// Validate function name is not empty
|
||||
if func_name.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let params_text = captures.get(2).map_or("", |m| m.as_str());
|
||||
|
||||
// Parse parameters
|
||||
let parameters = self.parse_steptml_parameters(params_text)?;
|
||||
|
||||
let arguments_str = serde_json::to_string(¶meters)
|
||||
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
|
||||
|
||||
Ok(Some(ToolCall {
|
||||
function: FunctionCall {
|
||||
name: func_name.to_string(),
|
||||
arguments: arguments_str,
|
||||
},
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Step3Parser {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolParser for Step3Parser {
|
||||
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
|
||||
if !self.has_tool_markers(text) {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
// Find where tool calls begin
|
||||
let idx = text.find("<｜tool_calls_begin｜>").unwrap();
|
||||
let normal_text = text[..idx].to_string();
|
||||
|
||||
// Extract tool calls
|
||||
let mut tools = Vec::new();
|
||||
for mat in self.tool_call_extractor.find_iter(text) {
|
||||
match self.parse_tool_call(mat.as_str()) {
|
||||
Ok(Some(tool)) => tools.push(tool),
|
||||
Ok(None) => continue,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse tool call: {}", e);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were successfully parsed despite having markers, return entire text as fallback
|
||||
if tools.is_empty() {
|
||||
return Ok((text.to_string(), vec![]));
|
||||
}
|
||||
|
||||
Ok((normal_text, tools))
|
||||
}
|
||||
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult> {
|
||||
self.buffer.push_str(chunk);
|
||||
|
||||
// Build tool indices for validation
|
||||
let tool_indices = helpers::get_tool_indices(tools);
|
||||
|
||||
// Stage 1: If we've finished the tool block, everything is normal text
|
||||
if self.tool_block_finished {
|
||||
let normal_text = std::mem::take(&mut self.buffer);
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Stage 2: Check if tool block hasn't started yet
|
||||
if !self.in_tool_block {
|
||||
if self.buffer.contains(self.bot_token) {
|
||||
let idx = self.buffer.find(self.bot_token).unwrap();
|
||||
let normal_text = self.buffer[..idx].to_string();
|
||||
self.buffer = self.buffer[idx + self.bot_token.len()..].to_string();
|
||||
self.in_tool_block = true;
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
} else {
|
||||
// Check if we might have a partial bot_token
|
||||
if helpers::ends_with_partial_token(&self.buffer, self.bot_token).is_some() {
|
||||
return Ok(StreamingParseResult::default()); // Wait for more text
|
||||
} else {
|
||||
let normal_text = std::mem::take(&mut self.buffer);
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text,
|
||||
calls: vec![],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// We're inside the tool block
|
||||
let mut calls = Vec::new();
|
||||
|
||||
// Stage 3: Check if tool block is ending
|
||||
if self.buffer.contains(self.eot_token) {
|
||||
let idx = self.buffer.find(self.eot_token).unwrap();
|
||||
|
||||
// If we're in the middle of a tool call, we need to handle it
|
||||
if self.in_tool_call {
|
||||
// The buffer before eot_token might contain the end of the current tool call
|
||||
let before_eot = &self.buffer[..idx];
|
||||
if before_eot.contains(self.tool_call_end) {
|
||||
// Parse this final tool call
|
||||
let result = self.parse_partial_tool_call(&tool_indices)?;
|
||||
calls.extend(result.calls);
|
||||
} else {
|
||||
// Incomplete tool call - log warning
|
||||
tracing::warn!("Tool block ended with incomplete tool call");
|
||||
}
|
||||
}
|
||||
|
||||
let remaining = self.buffer[idx + self.eot_token.len()..].to_string();
|
||||
self.buffer.clear();
|
||||
self.tool_block_finished = true;
|
||||
|
||||
// Reset any partial tool call state
|
||||
self.reset_streaming_state();
|
||||
|
||||
return Ok(StreamingParseResult {
|
||||
normal_text: remaining,
|
||||
calls,
|
||||
});
|
||||
}
|
||||
|
||||
// Stage 4: Check if we're in a tool call or need to start one
|
||||
if !self.in_tool_call {
|
||||
if self.buffer.contains(self.tool_call_begin) {
|
||||
let idx = self.buffer.find(self.tool_call_begin).unwrap();
|
||||
// Remove any content before tool call begin (shouldn't happen but be safe)
|
||||
self.buffer = self.buffer[idx + self.tool_call_begin.len()..].to_string();
|
||||
self.in_tool_call = true;
|
||||
self.function_name_sent = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
// Fall through to parse the partial tool call
|
||||
} else {
|
||||
// Wait for tool call to begin
|
||||
return Ok(StreamingParseResult::default());
|
||||
}
|
||||
}
|
||||
|
||||
// Stage 5: Parse partial tool call
|
||||
if self.in_tool_call {
|
||||
return self.parse_partial_tool_call(&tool_indices);
|
||||
}
|
||||
|
||||
Ok(StreamingParseResult::default())
|
||||
}
|
||||
|
||||
fn has_tool_markers(&self, text: &str) -> bool {
|
||||
text.contains(self.bot_token)
|
||||
}
|
||||
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
|
||||
helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
// Reset standard state
|
||||
self.buffer.clear();
|
||||
self.prev_tool_call_arr.clear();
|
||||
self.current_tool_id = -1;
|
||||
self.streamed_args_for_tool.clear();
|
||||
|
||||
// Reset Step3-specific fields
|
||||
self.in_tool_block = false;
|
||||
self.tool_block_finished = false;
|
||||
self.current_function_name.clear();
|
||||
self.current_parameters.clear();
|
||||
self.in_tool_call = false;
|
||||
self.function_name_sent = false;
|
||||
}
|
||||
}
|
||||
@@ -1,533 +0,0 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::tool_parser::{
|
||||
errors::{ParserError, ParserResult},
|
||||
traits::PartialJsonParser,
|
||||
};
|
||||
|
||||
/// Parser for incomplete JSON
|
||||
pub struct PartialJson {
|
||||
/// Maximum depth for nested structures
|
||||
max_depth: usize,
|
||||
/// Whether to allow incomplete values
|
||||
allow_incomplete: bool,
|
||||
}
|
||||
|
||||
impl PartialJson {
|
||||
/// Create a new partial JSON parser
|
||||
pub fn new(max_depth: usize, allow_incomplete: bool) -> Self {
|
||||
Self {
|
||||
max_depth,
|
||||
allow_incomplete,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse potentially incomplete JSON, returning parsed value and consumed bytes
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `input` - The JSON string to parse
|
||||
/// * `allow_partial_strings` - When false, incomplete strings cause parsing to stop
|
||||
/// (matches Python's Allow.ALL & ~Allow.STR behavior)
|
||||
pub fn parse_value(
|
||||
&self,
|
||||
input: &str,
|
||||
allow_partial_strings: bool,
|
||||
) -> ParserResult<(Value, usize)> {
|
||||
let mut parser = Parser::new(
|
||||
input,
|
||||
self.max_depth,
|
||||
self.allow_incomplete,
|
||||
allow_partial_strings,
|
||||
);
|
||||
let value = parser.parse_value(0)?;
|
||||
Ok((value, parser.position))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PartialJson {
|
||||
fn default() -> Self {
|
||||
Self::new(32, true)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialJsonParser for PartialJson {
|
||||
fn parse(&self, input: &str) -> ParserResult<(Value, usize)> {
|
||||
// Default to allowing partial strings
|
||||
self.parse_value(input, true)
|
||||
}
|
||||
|
||||
fn is_complete(&self, input: &str) -> bool {
|
||||
// Try to parse as complete JSON
|
||||
serde_json::from_str::<Value>(input).is_ok()
|
||||
}
|
||||
|
||||
fn max_depth(&self) -> usize {
|
||||
self.max_depth
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal parser state
|
||||
struct Parser<'a> {
|
||||
chars: std::iter::Peekable<std::str::Chars<'a>>,
|
||||
position: usize,
|
||||
max_depth: usize,
|
||||
allow_incomplete: bool,
|
||||
allow_partial_strings: bool,
|
||||
}
|
||||
|
||||
impl<'a> Parser<'a> {
|
||||
fn new(
|
||||
input: &'a str,
|
||||
max_depth: usize,
|
||||
allow_incomplete: bool,
|
||||
allow_partial_strings: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
chars: input.chars().peekable(),
|
||||
position: 0,
|
||||
max_depth,
|
||||
allow_incomplete,
|
||||
allow_partial_strings,
|
||||
}
|
||||
}
|
||||
|
||||
fn peek(&mut self) -> Option<char> {
|
||||
self.chars.peek().copied()
|
||||
}
|
||||
|
||||
fn advance(&mut self) {
|
||||
if self.chars.next().is_some() {
|
||||
self.position += 1;
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_whitespace(&mut self) {
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_whitespace() {
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_value(&mut self, depth: usize) -> ParserResult<Value> {
|
||||
if depth > self.max_depth {
|
||||
return Err(ParserError::DepthExceeded(self.max_depth));
|
||||
}
|
||||
|
||||
self.skip_whitespace();
|
||||
|
||||
match self.peek() {
|
||||
Some('{') => self.parse_object(depth + 1),
|
||||
Some('[') => self.parse_array(depth + 1),
|
||||
Some('"') => self.parse_string(),
|
||||
Some('t') | Some('f') => self.parse_bool(),
|
||||
Some('n') => self.parse_null(),
|
||||
Some(c) if c == '-' || c.is_ascii_digit() => self.parse_number(),
|
||||
_ => {
|
||||
if self.allow_incomplete {
|
||||
Ok(Value::Null)
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed("Unexpected character".into()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_object(&mut self, depth: usize) -> ParserResult<Value> {
|
||||
if depth > self.max_depth {
|
||||
return Err(ParserError::DepthExceeded(self.max_depth));
|
||||
}
|
||||
|
||||
let mut object = Map::new();
|
||||
|
||||
// Consume '{'
|
||||
self.advance();
|
||||
self.skip_whitespace();
|
||||
|
||||
// Check for empty object
|
||||
if self.peek() == Some('}') {
|
||||
self.advance();
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
|
||||
loop {
|
||||
// Parse key
|
||||
let key = match self.parse_string() {
|
||||
Ok(Value::String(s)) => s,
|
||||
Err(_) if self.allow_incomplete => {
|
||||
// Incomplete object
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
_ => return Err(ParserError::ParsingFailed("Expected string key".into())),
|
||||
};
|
||||
|
||||
self.skip_whitespace();
|
||||
|
||||
// Expect ':'
|
||||
if self.peek() != Some(':') {
|
||||
if self.allow_incomplete {
|
||||
// Add null value for incomplete pair
|
||||
object.insert(key, Value::Null);
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
return Err(ParserError::ParsingFailed("Expected ':'".into()));
|
||||
}
|
||||
self.advance();
|
||||
self.skip_whitespace();
|
||||
|
||||
// Parse value (keep same depth - we already incremented in parse_object)
|
||||
let value = match self.parse_value(depth) {
|
||||
Ok(v) => v,
|
||||
Err(_) if self.allow_incomplete => {
|
||||
// When allow_partial_strings is false, don't add the key with Null
|
||||
// Just return the object without this incomplete key-value pair
|
||||
// This matches Python's behavior: Allow.ALL & ~Allow.STR
|
||||
if self.allow_partial_strings {
|
||||
// Add null for incomplete value
|
||||
object.insert(key, Value::Null);
|
||||
}
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
object.insert(key, value);
|
||||
self.skip_whitespace();
|
||||
|
||||
match self.peek() {
|
||||
Some(',') => {
|
||||
self.advance();
|
||||
self.skip_whitespace();
|
||||
// Check for trailing comma
|
||||
if self.peek() == Some('}') {
|
||||
self.advance();
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
}
|
||||
Some('}') => {
|
||||
self.advance();
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
None if self.allow_incomplete => {
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
_ => {
|
||||
if self.allow_incomplete {
|
||||
return Ok(Value::Object(object));
|
||||
}
|
||||
return Err(ParserError::ParsingFailed("Expected ',' or '}'".into()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_array(&mut self, depth: usize) -> ParserResult<Value> {
|
||||
if depth > self.max_depth {
|
||||
return Err(ParserError::DepthExceeded(self.max_depth));
|
||||
}
|
||||
|
||||
let mut array = Vec::new();
|
||||
|
||||
// Consume '['
|
||||
self.advance();
|
||||
self.skip_whitespace();
|
||||
|
||||
// Check for empty array
|
||||
if self.peek() == Some(']') {
|
||||
self.advance();
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
|
||||
loop {
|
||||
// Parse value (keep same depth - we already incremented in parse_object)
|
||||
let value = match self.parse_value(depth) {
|
||||
Ok(v) => v,
|
||||
Err(_) if self.allow_incomplete => {
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
array.push(value);
|
||||
self.skip_whitespace();
|
||||
|
||||
match self.peek() {
|
||||
Some(',') => {
|
||||
self.advance();
|
||||
self.skip_whitespace();
|
||||
// Check for trailing comma
|
||||
if self.peek() == Some(']') {
|
||||
self.advance();
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
}
|
||||
Some(']') => {
|
||||
self.advance();
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
None if self.allow_incomplete => {
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
_ => {
|
||||
if self.allow_incomplete {
|
||||
return Ok(Value::Array(array));
|
||||
}
|
||||
return Err(ParserError::ParsingFailed("Expected ',' or ']'".into()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string(&mut self) -> ParserResult<Value> {
|
||||
if self.peek() != Some('"') {
|
||||
return Err(ParserError::ParsingFailed("Expected '\"'".into()));
|
||||
}
|
||||
|
||||
// Consume opening quote
|
||||
self.advance();
|
||||
|
||||
let mut string = String::new();
|
||||
let mut escaped = false;
|
||||
|
||||
while let Some(ch) = self.peek() {
|
||||
if escaped {
|
||||
// Handle escape sequences
|
||||
let escaped_char = match ch {
|
||||
'"' | '\\' | '/' => ch,
|
||||
'b' => '\u{0008}',
|
||||
'f' => '\u{000C}',
|
||||
'n' => '\n',
|
||||
'r' => '\r',
|
||||
't' => '\t',
|
||||
'u' => {
|
||||
// Unicode escape
|
||||
self.advance();
|
||||
let hex = self.parse_unicode_escape()?;
|
||||
string.push(hex);
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
_ => ch, // Invalid escape, but be lenient
|
||||
};
|
||||
string.push(escaped_char);
|
||||
escaped = false;
|
||||
} else if ch == '\\' {
|
||||
escaped = true;
|
||||
} else if ch == '"' {
|
||||
// End of string
|
||||
self.advance();
|
||||
return Ok(Value::String(string));
|
||||
} else {
|
||||
string.push(ch);
|
||||
}
|
||||
self.advance();
|
||||
}
|
||||
|
||||
// Incomplete string
|
||||
if self.allow_incomplete && self.allow_partial_strings {
|
||||
Ok(Value::String(string))
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed("Unterminated string".into()))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_unicode_escape(&mut self) -> ParserResult<char> {
|
||||
let mut hex = String::new();
|
||||
for _ in 0..4 {
|
||||
if let Some(ch) = self.peek() {
|
||||
if ch.is_ascii_hexdigit() {
|
||||
hex.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if hex.len() == 4 {
|
||||
u32::from_str_radix(&hex, 16)
|
||||
.ok()
|
||||
.and_then(char::from_u32)
|
||||
.ok_or_else(|| ParserError::ParsingFailed("Invalid unicode escape".into()))
|
||||
} else if self.allow_incomplete {
|
||||
Ok('\u{FFFD}') // Replacement character
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed(
|
||||
"Incomplete unicode escape".into(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_number(&mut self) -> ParserResult<Value> {
|
||||
let mut number = String::new();
|
||||
|
||||
// Handle negative sign
|
||||
if self.peek() == Some('-') {
|
||||
number.push('-');
|
||||
self.advance();
|
||||
}
|
||||
|
||||
// Parse integer part
|
||||
if self.peek() == Some('0') {
|
||||
number.push('0');
|
||||
self.advance();
|
||||
} else {
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_ascii_digit() {
|
||||
number.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse decimal part
|
||||
if self.peek() == Some('.') {
|
||||
number.push('.');
|
||||
self.advance();
|
||||
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_ascii_digit() {
|
||||
number.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse exponent
|
||||
if let Some(ch) = self.peek() {
|
||||
if ch == 'e' || ch == 'E' {
|
||||
number.push(ch);
|
||||
self.advance();
|
||||
|
||||
if let Some(sign) = self.peek() {
|
||||
if sign == '+' || sign == '-' {
|
||||
number.push(sign);
|
||||
self.advance();
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_ascii_digit() {
|
||||
number.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try to parse as integer first, then as float
|
||||
if let Ok(n) = number.parse::<i64>() {
|
||||
Ok(Value::Number(serde_json::Number::from(n)))
|
||||
} else if let Ok(n) = number.parse::<f64>() {
|
||||
Ok(Value::Number(
|
||||
serde_json::Number::from_f64(n).unwrap_or_else(|| serde_json::Number::from(0)),
|
||||
))
|
||||
} else if self.allow_incomplete {
|
||||
Ok(Value::Number(serde_json::Number::from(0)))
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed("Invalid number".into()))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_bool(&mut self) -> ParserResult<Value> {
|
||||
let mut word = String::new();
|
||||
|
||||
// Peek at upcoming characters to validate it looks like a boolean
|
||||
let mut temp_chars = self.chars.clone();
|
||||
while let Some(&ch) = temp_chars.peek() {
|
||||
if ch.is_alphabetic() && word.len() < 5 {
|
||||
// "false" is 5 chars
|
||||
word.push(ch);
|
||||
temp_chars.next();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if it's a valid boolean prefix
|
||||
let is_valid = word == "true"
|
||||
|| word == "false"
|
||||
|| (self.allow_incomplete && ("true".starts_with(&word) || "false".starts_with(&word)));
|
||||
|
||||
if !is_valid {
|
||||
return Err(ParserError::ParsingFailed("Invalid boolean".into()));
|
||||
}
|
||||
|
||||
// Now actually consume the characters
|
||||
word.clear();
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_alphabetic() {
|
||||
word.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
match word.as_str() {
|
||||
"true" => Ok(Value::Bool(true)),
|
||||
"false" => Ok(Value::Bool(false)),
|
||||
partial if self.allow_incomplete => {
|
||||
if "true".starts_with(partial) {
|
||||
Ok(Value::Bool(true))
|
||||
} else if "false".starts_with(partial) {
|
||||
Ok(Value::Bool(false))
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed("Invalid boolean".into()))
|
||||
}
|
||||
}
|
||||
_ => Err(ParserError::ParsingFailed("Invalid boolean".into())),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_null(&mut self) -> ParserResult<Value> {
|
||||
let mut word = String::new();
|
||||
|
||||
// Peek at upcoming characters to validate it looks like "null"
|
||||
let mut temp_chars = self.chars.clone();
|
||||
while let Some(&ch) = temp_chars.peek() {
|
||||
if ch.is_alphabetic() && word.len() < 4 {
|
||||
// "null" is 4 chars
|
||||
word.push(ch);
|
||||
temp_chars.next();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if it's a valid null prefix
|
||||
let is_valid = word == "null" || (self.allow_incomplete && "null".starts_with(&word));
|
||||
|
||||
if !is_valid {
|
||||
return Err(ParserError::ParsingFailed("Invalid null".into()));
|
||||
}
|
||||
|
||||
// Now actually consume the characters
|
||||
word.clear();
|
||||
while let Some(ch) = self.peek() {
|
||||
if ch.is_alphabetic() {
|
||||
word.push(ch);
|
||||
self.advance();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if word == "null" || (self.allow_incomplete && "null".starts_with(&word)) {
|
||||
Ok(Value::Null)
|
||||
} else {
|
||||
Err(ParserError::ParsingFailed("Invalid null".into()))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,740 +0,0 @@
|
||||
use super::*;
|
||||
use crate::{
|
||||
protocols::common::{Function, Tool},
|
||||
tool_parser::{
|
||||
parsers::{JsonParser, QwenCoderParser},
|
||||
partial_json::PartialJson,
|
||||
traits::ToolParser,
|
||||
},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_parser_factory() {
|
||||
let factory = ParserFactory::new();
|
||||
|
||||
// Test that we can get a pooled parser
|
||||
let pooled_parser = factory.get_pooled("gpt-4");
|
||||
let parser = pooled_parser.lock().await;
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "arguments": {}}"#));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_parser_factory_model_mapping() {
|
||||
let factory = ParserFactory::new();
|
||||
|
||||
// Test model mapping
|
||||
factory.registry().map_model("test-model", "json");
|
||||
|
||||
// Get parser for the test model
|
||||
let pooled_parser = factory.get_pooled("test-model");
|
||||
let parser = pooled_parser.lock().await;
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "arguments": {}}"#));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_call_serialization() {
|
||||
let tool_call = ToolCall {
|
||||
function: FunctionCall {
|
||||
name: "search".to_string(),
|
||||
arguments: r#"{"query": "rust programming"}"#.to_string(),
|
||||
},
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&tool_call).unwrap();
|
||||
assert!(json.contains("search"));
|
||||
assert!(json.contains("rust programming"));
|
||||
|
||||
let parsed: ToolCall = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.function.name, "search");
|
||||
assert_eq!(
|
||||
parsed.function.arguments,
|
||||
r#"{"query": "rust programming"}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_json_parser() {
|
||||
let parser = PartialJson::default();
|
||||
|
||||
let input = r#"{"name": "test", "value": 42}"#;
|
||||
let (value, consumed) = parser.parse_value(input, true).unwrap();
|
||||
assert_eq!(value["name"], "test");
|
||||
assert_eq!(value["value"], 42);
|
||||
assert_eq!(consumed, input.len());
|
||||
|
||||
let input = r#"{"name": "test", "value": "#;
|
||||
let (value, _consumed) = parser.parse_value(input, true).unwrap();
|
||||
assert_eq!(value["name"], "test");
|
||||
assert!(value["value"].is_null());
|
||||
|
||||
let input = r#"{"name": "tes"#;
|
||||
let (value, _consumed) = parser.parse_value(input, true).unwrap();
|
||||
assert_eq!(value["name"], "tes");
|
||||
|
||||
let input = r#"[1, 2, "#;
|
||||
let (value, _consumed) = parser.parse_value(input, true).unwrap();
|
||||
assert!(value.is_array());
|
||||
assert_eq!(value[0], 1);
|
||||
assert_eq!(value[1], 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_json_depth_limit() {
|
||||
// max_depth of 3 allows nesting up to 3 levels
|
||||
// Set allow_incomplete to false to get errors instead of partial results
|
||||
let parser = PartialJson::new(3, false);
|
||||
|
||||
// This should work (simple object)
|
||||
let input = r#"{"a": 1}"#;
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// This should work (nested to depth 3)
|
||||
let input = r#"{"a": {"b": {"c": 1}}}"#;
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// This should fail (nested to depth 4, exceeds limit)
|
||||
let input = r#"{"a": {"b": {"c": {"d": 1}}}}"#;
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// NOTE: test_stream_result_variants removed - StreamResult enum replaced by StreamingParseResult
|
||||
|
||||
#[test]
|
||||
fn test_partial_tool_call() {
|
||||
let mut partial = PartialToolCall {
|
||||
name: None,
|
||||
arguments_buffer: String::new(),
|
||||
start_position: 0,
|
||||
name_sent: false,
|
||||
streamed_args: String::new(),
|
||||
};
|
||||
|
||||
// Set name
|
||||
partial.name = Some("test_function".to_string());
|
||||
assert_eq!(partial.name.as_ref().unwrap(), "test_function");
|
||||
|
||||
// Append arguments
|
||||
partial.arguments_buffer.push_str(r#"{"key": "value"}"#);
|
||||
assert_eq!(partial.arguments_buffer, r#"{"key": "value"}"#);
|
||||
|
||||
// Update streaming state
|
||||
partial.name_sent = true;
|
||||
partial.streamed_args = r#"{"key": "#.to_string();
|
||||
assert!(partial.name_sent);
|
||||
assert_eq!(partial.streamed_args, r#"{"key": "#);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_complete_single() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "get_weather", "arguments": {"location": "San Francisco", "units": "celsius"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert!(tools[0].function.arguments.contains("San Francisco"));
|
||||
assert!(tools[0].function.arguments.contains("celsius"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_complete_array() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
let input = r#"[
|
||||
{"name": "get_weather", "arguments": {"location": "SF"}},
|
||||
{"name": "get_news", "arguments": {"query": "technology"}}
|
||||
]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert_eq!(tools[1].function.name, "get_news");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_with_parameters() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "calculate", "parameters": {"x": 10, "y": 20, "operation": "add"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "calculate");
|
||||
assert!(tools[0].function.arguments.contains("10"));
|
||||
assert!(tools[0].function.arguments.contains("20"));
|
||||
assert!(tools[0].function.arguments.contains("add"));
|
||||
}
|
||||
|
||||
// Tests removed - TokenConfig no longer supported in JsonParser
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiline_json_array() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
let input = r#"[
|
||||
{
|
||||
"name": "function1",
|
||||
"arguments": {
|
||||
"param1": "value1",
|
||||
"param2": 42
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "function2",
|
||||
"parameters": {
|
||||
"data": [1, 2, 3],
|
||||
"flag": false
|
||||
}
|
||||
}
|
||||
]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "function1");
|
||||
assert_eq!(tools[1].function.name, "function2");
|
||||
assert!(tools[0].function.arguments.contains("value1"));
|
||||
assert!(tools[1].function.arguments.contains("[1,2,3]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_json_parser_format_detection() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Should detect valid tool call formats
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "arguments": {}}"#));
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "parameters": {"x": 1}}"#));
|
||||
assert!(parser.has_tool_markers(r#"[{"name": "test"}]"#));
|
||||
|
||||
// Should not detect non-tool formats
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_factory_with_json_parser() {
|
||||
let factory = ParserFactory::new();
|
||||
|
||||
// Should get JSON parser for OpenAI models
|
||||
let pooled_parser = factory.get_pooled("gpt-4-turbo");
|
||||
let parser = pooled_parser.lock().await;
|
||||
|
||||
let input = r#"{"name": "test", "arguments": {"x": 1}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_invalid_input() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Invalid JSON should return empty results
|
||||
assert_eq!(parser.parse_complete("not json").await.unwrap().1.len(), 0);
|
||||
assert_eq!(parser.parse_complete("{invalid}").await.unwrap().1.len(), 0);
|
||||
assert_eq!(parser.parse_complete("").await.unwrap().1.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_empty_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Tool call with no arguments
|
||||
let input = r#"{"name": "get_time"}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_time");
|
||||
assert_eq!(tools[0].function.arguments, "{}");
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod failure_cases {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_malformed_tool_missing_name() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Missing name field
|
||||
let input = r#"{"arguments": {"x": 1}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should return empty for tool without name");
|
||||
|
||||
// Empty name
|
||||
let input = r#"{"name": "", "arguments": {"x": 1}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1, "Should accept empty name string");
|
||||
assert_eq!(tools[0].function.name, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_arguments_json() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Arguments is a string instead of object
|
||||
let input = r#"{"name": "test", "arguments": "not an object"}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
// Should serialize the string as JSON
|
||||
assert!(tools[0].function.arguments.contains("not an object"));
|
||||
|
||||
// Arguments is a number
|
||||
let input = r#"{"name": "test", "arguments": 42}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.arguments, "42");
|
||||
|
||||
// Arguments is null
|
||||
let input = r#"{"name": "test", "arguments": null}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.arguments, "null");
|
||||
}
|
||||
|
||||
// Test removed - wrapper token functionality moved to specific parsers
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_json_structures() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Trailing comma
|
||||
let input = r#"{"name": "test", "arguments": {"x": 1,}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should reject JSON with trailing comma");
|
||||
|
||||
// Missing quotes on keys
|
||||
let input = r#"{name: "test", arguments: {}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should reject invalid JSON syntax");
|
||||
|
||||
// Unclosed object
|
||||
let input = r#"{"name": "test", "arguments": {"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should reject incomplete JSON");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod edge_cases {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unicode_in_names_and_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Unicode in function name
|
||||
let input = r#"{"name": "获取天气", "arguments": {"location": "北京"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "获取天气");
|
||||
assert!(tools[0].function.arguments.contains("北京"));
|
||||
|
||||
// Emoji in arguments
|
||||
let input = r#"{"name": "send_message", "arguments": {"text": "Hello 👋 World 🌍"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("👋"));
|
||||
assert!(tools[0].function.arguments.contains("🌍"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_escaped_characters() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Escaped quotes in arguments
|
||||
let input = r#"{"name": "echo", "arguments": {"text": "He said \"hello\""}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains(r#"\"hello\""#));
|
||||
|
||||
// Escaped backslashes
|
||||
let input = r#"{"name": "path", "arguments": {"dir": "C:\\Users\\test"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("\\\\"));
|
||||
|
||||
// Newlines and tabs
|
||||
let input = r#"{"name": "format", "arguments": {"text": "line1\nline2\ttabbed"}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("\\n"));
|
||||
assert!(tools[0].function.arguments.contains("\\t"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_very_large_payloads() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Large arguments object
|
||||
let mut large_args = r#"{"name": "process", "arguments": {"#.to_string();
|
||||
for i in 0..1000 {
|
||||
large_args.push_str(&format!(r#""field_{}": "value_{}","#, i, i));
|
||||
}
|
||||
large_args.push_str(r#""final": "value"}}"#);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&large_args).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process");
|
||||
assert!(tools[0].function.arguments.contains("field_999"));
|
||||
|
||||
// Large array of tool calls
|
||||
let mut large_array = "[".to_string();
|
||||
for i in 0..100 {
|
||||
if i > 0 {
|
||||
large_array.push(',');
|
||||
}
|
||||
large_array.push_str(&format!(r#"{{"name": "func_{}", "arguments": {{}}}}"#, i));
|
||||
}
|
||||
large_array.push(']');
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&large_array).await.unwrap();
|
||||
assert_eq!(tools.len(), 100);
|
||||
assert_eq!(tools[99].function.name, "func_99");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mixed_array_tools_and_non_tools() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Array with both tool calls and non-tool objects
|
||||
let input = r#"[
|
||||
{"name": "tool1", "arguments": {}},
|
||||
{"not_a_tool": "just_data"},
|
||||
{"name": "tool2", "parameters": {"x": 1}},
|
||||
{"key": "value", "another": "field"}
|
||||
]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2, "Should only parse valid tool calls");
|
||||
assert_eq!(tools[0].function.name, "tool1");
|
||||
assert_eq!(tools[1].function.name, "tool2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_duplicate_keys_in_json() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// JSON with duplicate keys (last one wins in most parsers)
|
||||
let input = r#"{"name": "first", "name": "second", "arguments": {"x": 1, "x": 2}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(
|
||||
tools[0].function.name, "second",
|
||||
"Last duplicate key should win"
|
||||
);
|
||||
assert!(
|
||||
tools[0].function.arguments.contains("2"),
|
||||
"Last duplicate value should win"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_null_values_in_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Null values in arguments
|
||||
let input = r#"{"name": "test", "arguments": {"required": "value", "optional": null}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("null"));
|
||||
|
||||
// Array with null
|
||||
let input = r#"{"name": "test", "arguments": {"items": [1, null, "three"]}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("null"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_special_json_values() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Boolean values
|
||||
let input = r#"{"name": "toggle", "arguments": {"enabled": true, "disabled": false}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("true"));
|
||||
assert!(tools[0].function.arguments.contains("false"));
|
||||
|
||||
// Numbers (including float and negative)
|
||||
let input = r#"{"name": "calc", "arguments": {"int": 42, "float": 3.14, "negative": -17}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("42"));
|
||||
assert!(tools[0].function.arguments.contains("3.14"));
|
||||
assert!(tools[0].function.arguments.contains("-17"));
|
||||
|
||||
// Empty arrays and objects
|
||||
let input = r#"{"name": "test", "arguments": {"empty_arr": [], "empty_obj": {}}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("[]"));
|
||||
assert!(tools[0].function.arguments.contains("{}"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_function_field_alternative() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Using "function" instead of "name"
|
||||
let input = r#"{"function": "test_func", "arguments": {"x": 1}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test_func");
|
||||
|
||||
// Both "name" and "function" present (name should take precedence)
|
||||
let input = r#"{"name": "primary", "function": "secondary", "arguments": {}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "primary");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_whitespace_handling() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Extra whitespace everywhere
|
||||
let input = r#" {
|
||||
"name" : "test" ,
|
||||
"arguments" : {
|
||||
"key" : "value"
|
||||
}
|
||||
} "#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
|
||||
// Minified JSON (no whitespace)
|
||||
let input = r#"{"name":"compact","arguments":{"a":1,"b":2}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "compact");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod stress_tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deeply_nested_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Deeply nested structure
|
||||
let input = r#"{
|
||||
"name": "nested",
|
||||
"arguments": {
|
||||
"level1": {
|
||||
"level2": {
|
||||
"level3": {
|
||||
"level4": {
|
||||
"level5": {
|
||||
"value": "deep"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert!(tools[0].function.arguments.contains("deep"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_parser_usage() {
|
||||
let parser = std::sync::Arc::new(JsonParser::new());
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
for i in 0..10 {
|
||||
let parser_clone = parser.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let input = format!(r#"{{"name": "func_{}", "arguments": {{}}}}"#, i);
|
||||
let (_normal_text, tools) = parser_clone.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, format!("func_{}", i));
|
||||
});
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod qwen_coder_tests {
|
||||
use super::*;
|
||||
|
||||
fn create_test_tools() -> Vec<Tool> {
|
||||
vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather information".to_string()),
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"units": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
}]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_incremental_parameter_streaming() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = [
|
||||
"<tool_call>",
|
||||
r#"<function=get_weather>"#,
|
||||
r#"<parameter=city>Paris</parameter>"#,
|
||||
r#"<parameter=units>metric</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut all_calls = Vec::new();
|
||||
|
||||
// Process each chunk
|
||||
for (i, chunk) in chunks.iter().enumerate() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
println!("Chunk {}: {:?}", i, chunk);
|
||||
println!(" Calls: {:?}", result.calls);
|
||||
println!(" Normal text: {:?}", result.normal_text);
|
||||
|
||||
for call in &result.calls {
|
||||
all_calls.push(call.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Verify the final result
|
||||
// We should have:
|
||||
// 1. Tool name call (from chunk 2)
|
||||
// 2. First parameter call: {"city": "Paris"}
|
||||
// 3. Second parameter call: , "units": "metric"
|
||||
// Final result should be: {"city": "Paris", "units": "metric"}
|
||||
|
||||
assert!(!all_calls.is_empty(), "Should have at least one call");
|
||||
|
||||
// Check that we have the tool name
|
||||
let name_call = all_calls.iter().find(|c| c.name.is_some());
|
||||
assert!(name_call.is_some(), "Should have tool name call");
|
||||
assert_eq!(name_call.unwrap().name.as_ref().unwrap(), "get_weather");
|
||||
|
||||
// Check parameter calls
|
||||
let param_calls: Vec<_> = all_calls.iter().filter(|c| c.name.is_none()).collect();
|
||||
assert!(!param_calls.is_empty(), "Should have parameter calls");
|
||||
|
||||
// Verify final arguments format by concatenating all parameter fragments
|
||||
let params_str: String = param_calls.iter().map(|c| c.parameters.as_str()).collect();
|
||||
println!("Final streamed args: {}", params_str);
|
||||
|
||||
// Should contain both city and units parameters
|
||||
assert!(params_str.contains("city"), "Should contain city parameter");
|
||||
assert!(params_str.contains("Paris"), "Should contain Paris value");
|
||||
assert!(
|
||||
params_str.contains("units"),
|
||||
"Should contain units parameter"
|
||||
);
|
||||
assert!(params_str.contains("metric"), "Should contain metric value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_incremental_parameter_streaming_with_partial_values() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Test with parameter values that arrive in multiple chunks
|
||||
// This tests the buffering logic for partial XML tags
|
||||
let chunks = [
|
||||
"<tool_call><function=get_weather>",
|
||||
r#"<parameter=city>Paris</parameter>"#,
|
||||
r#"<parameter=units>metric</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut all_calls = Vec::new();
|
||||
|
||||
// Process each chunk
|
||||
for (i, chunk) in chunks.iter().enumerate() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
println!("Chunk {}: {:?}", i, chunk);
|
||||
println!(" Calls: {:?}", result.calls);
|
||||
|
||||
for call in &result.calls {
|
||||
all_calls.push(call.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Verify we got the tool name
|
||||
let name_call = all_calls.iter().find(|c| c.name.is_some());
|
||||
assert!(name_call.is_some(), "Should have tool name call");
|
||||
assert_eq!(name_call.unwrap().name.as_ref().unwrap(), "get_weather");
|
||||
|
||||
// Verify we got parameter calls
|
||||
let param_calls: Vec<_> = all_calls.iter().filter(|c| c.name.is_none()).collect();
|
||||
assert!(!param_calls.is_empty(), "Should have parameter calls");
|
||||
|
||||
// Verify final arguments by concatenating all parameter fragments
|
||||
let params_str: String = param_calls.iter().map(|c| c.parameters.as_str()).collect();
|
||||
assert!(params_str.contains("city"), "Should contain city parameter");
|
||||
assert!(params_str.contains("Paris"), "Should contain Paris value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_nested_json_parameter() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: Function {
|
||||
name: "test_function".to_string(),
|
||||
description: None,
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {"type": "object"}
|
||||
}
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
r#"<function=test_function>"#,
|
||||
r#"<parameter=nested>{"key": "value"}</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut all_calls = Vec::new();
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
all_calls.push(call);
|
||||
}
|
||||
}
|
||||
|
||||
// Verify nested JSON is parsed correctly via collected calls
|
||||
let param_calls: Vec<_> = all_calls.iter().filter(|c| c.name.is_none()).collect();
|
||||
assert!(!param_calls.is_empty(), "Should have parameter calls");
|
||||
|
||||
// The first parameter call should contain the nested JSON
|
||||
let params_str: String = param_calls.iter().map(|c| c.parameters.as_str()).collect();
|
||||
assert!(
|
||||
params_str.contains("nested"),
|
||||
"Should contain nested parameter"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,76 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::{
|
||||
protocols::common::Tool,
|
||||
tool_parser::{
|
||||
errors::ParserResult,
|
||||
types::{StreamingParseResult, ToolCall},
|
||||
},
|
||||
};
|
||||
|
||||
/// Core trait for all tool parsers
|
||||
#[async_trait]
|
||||
pub trait ToolParser: Send + Sync {
|
||||
/// Parse complete tool calls from final output
|
||||
/// Returns (remaining_normal_text, tool_calls) tuple
|
||||
async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec<ToolCall>)>;
|
||||
|
||||
/// Parse tool calls from model output (streaming)
|
||||
/// Parsers now maintain internal state, so self is mutable
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `chunk` - New text chunk from model output
|
||||
/// * `tools` - List of available tools for validation
|
||||
async fn parse_incremental(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult>;
|
||||
|
||||
/// Check if text contains tool calls in this parser's format
|
||||
fn has_tool_markers(&self, text: &str) -> bool;
|
||||
|
||||
/// Optionally expose a token-aware parser implementation.
|
||||
/// Default returns `None`, meaning the parser only supports text input.
|
||||
fn as_token_parser(&self) -> Option<&dyn TokenToolParser> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Get unstreamed tool call arguments
|
||||
/// Returns tool call items for arguments that have been parsed but not yet streamed
|
||||
fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::tool_parser::types::ToolCallItem>> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Reset the parser state for reuse across requests.
|
||||
/// This should clear all buffers and reset state to initial values.
|
||||
fn reset(&mut self) {
|
||||
// Default no-op implementation
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for partial JSON parsing
|
||||
pub trait PartialJsonParser: Send + Sync {
|
||||
/// Parse potentially incomplete JSON
|
||||
fn parse(&self, input: &str) -> ParserResult<(serde_json::Value, usize)>;
|
||||
|
||||
/// Check if JSON is complete
|
||||
fn is_complete(&self, input: &str) -> bool;
|
||||
|
||||
/// Get the maximum parsing depth
|
||||
fn max_depth(&self) -> usize;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait TokenToolParser: ToolParser {
|
||||
/// Parse complete tool calls when provided with raw token IDs.
|
||||
async fn parse_complete_tokens(&self, tokens: &[u32]) -> ParserResult<(String, Vec<ToolCall>)>;
|
||||
|
||||
/// Streaming parser entrypoint for token chunks.
|
||||
/// Parsers maintain internal state, so self is mutable
|
||||
async fn parse_incremental_tokens(
|
||||
&mut self,
|
||||
tokens: &[u32],
|
||||
tools: &[Tool],
|
||||
) -> ParserResult<StreamingParseResult>;
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Parsed tool call from model output
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct ToolCall {
|
||||
/// Function call details
|
||||
pub function: FunctionCall,
|
||||
}
|
||||
|
||||
/// Function call within a tool call
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct FunctionCall {
|
||||
/// Name of the function to call
|
||||
pub name: String,
|
||||
/// Arguments as JSON string
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// Simple partial tool call for streaming
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PartialToolCall {
|
||||
/// Tool name (if parsed)
|
||||
pub name: Option<String>,
|
||||
/// Buffer for accumulating arguments
|
||||
pub arguments_buffer: String,
|
||||
/// Start position in the input buffer
|
||||
pub start_position: usize,
|
||||
/// Whether the name has been sent (for streaming)
|
||||
pub name_sent: bool,
|
||||
/// Arguments already streamed
|
||||
pub streamed_args: String,
|
||||
}
|
||||
|
||||
/// Result of streaming parse operation (matches Python StreamingParseResult)
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct StreamingParseResult {
|
||||
/// Normal text that's not part of tool calls
|
||||
pub normal_text: String,
|
||||
/// Tool call items parsed from the chunk
|
||||
pub calls: Vec<ToolCallItem>,
|
||||
}
|
||||
|
||||
/// Simple encapsulation of parsed tool call for streaming (matches Python ToolCallItem)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolCallItem {
|
||||
/// Tool index in the array
|
||||
pub tool_index: usize,
|
||||
/// Tool name (only present on first chunk)
|
||||
pub name: Option<String>,
|
||||
/// Incremental JSON arguments
|
||||
pub parameters: String,
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
//! Tool parser integration tests
|
||||
|
||||
pub mod tool_parser_deepseek;
|
||||
pub mod tool_parser_edge_cases;
|
||||
pub mod tool_parser_fallback;
|
||||
pub mod tool_parser_glm47_moe;
|
||||
pub mod tool_parser_glm4_moe;
|
||||
pub mod tool_parser_json;
|
||||
pub mod tool_parser_kimik2;
|
||||
pub mod tool_parser_llama;
|
||||
pub mod tool_parser_minimax_m2;
|
||||
pub mod tool_parser_mistral;
|
||||
pub mod tool_parser_mixed_edge_cases;
|
||||
pub mod tool_parser_partial_json;
|
||||
pub mod tool_parser_pythonic;
|
||||
pub mod tool_parser_qwen;
|
||||
pub mod tool_parser_qwen_coder;
|
||||
pub mod tool_parser_step3;
|
||||
@@ -1,160 +0,0 @@
|
||||
//! DeepSeek V3 Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{DeepSeekParser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_complete_parsing() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
let input = r#"Let me help you with that.
|
||||
<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
|
||||
```json
|
||||
{"location": "Tokyo", "units": "celsius"}
|
||||
```<｜tool▁call▁end｜><｜tool▁calls▁end｜>
|
||||
The weather in Tokyo is..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me help you with that.\n");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["location"], "Tokyo");
|
||||
assert_eq!(args["units"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_multiple_tools() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
let input = r#"<｜tool▁calls▁begin｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>search
|
||||
```json
|
||||
{"query": "rust programming"}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>translate
|
||||
```json
|
||||
{"text": "Hello World", "to": "ja"}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁calls▁end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_streaming() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = DeepSeekParser::new();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>",
|
||||
"function<｜tool▁sep｜>get_weather\n",
|
||||
"```json\n",
|
||||
r#"{"location": "#,
|
||||
r#""Beijing", "#,
|
||||
r#""units": "metric"}"#,
|
||||
"\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_nested_json() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
let input = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>process
|
||||
```json
|
||||
{
|
||||
"data": {
|
||||
"nested": {
|
||||
"deep": [1, 2, 3]
|
||||
}
|
||||
}
|
||||
}
|
||||
```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["data"]["nested"]["deep"].is_array());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deepseek_format_detection() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
// Should detect DeepSeek format
|
||||
assert!(parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
|
||||
assert!(parser.has_tool_markers("text with <｜tool▁calls▁begin｜> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<tool_call>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_malformed_json_handling() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
// Malformed JSON should be skipped
|
||||
let input = r#"<｜tool▁calls▁begin｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>broken
|
||||
```json
|
||||
{invalid json}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>valid
|
||||
```json
|
||||
{"key": "value"}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁calls▁end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
// Only the valid tool call should be parsed
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "valid");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_tool_calls() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
let input = r#"<｜tool▁calls▁begin｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
|
||||
```json
|
||||
{"location": "Tokyo"}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather
|
||||
```json
|
||||
{"location": "Paris"}
|
||||
```<｜tool▁call▁end｜>
|
||||
<｜tool▁calls▁end｜><｜end▁of▁sentence｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert_eq!(tools[1].function.name, "get_weather");
|
||||
}
|
||||
@@ -1,348 +0,0 @@
|
||||
//! Edge Cases and Error Handling Tests
|
||||
//!
|
||||
//! Tests for malformed input, edge cases, and error recovery
|
||||
|
||||
use smg::tool_parser::{JsonParser, MistralParser, PythonicParser, QwenParser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_input() {
|
||||
// Test that all parsers handle empty input correctly
|
||||
let json_parser = JsonParser::new();
|
||||
let (_normal_text, tools) = json_parser.parse_complete("").await.unwrap();
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"JSON parser should return empty for empty input"
|
||||
);
|
||||
|
||||
let mistral_parser = MistralParser::new();
|
||||
let (_normal_text, tools) = mistral_parser.parse_complete("").await.unwrap();
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Mistral parser should return empty for empty input"
|
||||
);
|
||||
|
||||
let qwen_parser = QwenParser::new();
|
||||
let (_normal_text, tools) = qwen_parser.parse_complete("").await.unwrap();
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Qwen parser should return empty for empty input"
|
||||
);
|
||||
|
||||
let pythonic_parser = PythonicParser::new();
|
||||
let (_normal_text, tools) = pythonic_parser.parse_complete("").await.unwrap();
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Pythonic parser should return empty for empty input"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_plain_text_no_tools() {
|
||||
let plain_text = "This is just a regular response with no tool calls whatsoever.";
|
||||
|
||||
let json_parser = JsonParser::new();
|
||||
assert_eq!(
|
||||
json_parser
|
||||
.parse_complete(plain_text)
|
||||
.await
|
||||
.unwrap()
|
||||
.1
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
|
||||
let mistral_parser = MistralParser::new();
|
||||
assert_eq!(
|
||||
mistral_parser
|
||||
.parse_complete(plain_text)
|
||||
.await
|
||||
.unwrap()
|
||||
.1
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
|
||||
let qwen_parser = QwenParser::new();
|
||||
assert_eq!(
|
||||
qwen_parser
|
||||
.parse_complete(plain_text)
|
||||
.await
|
||||
.unwrap()
|
||||
.1
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
|
||||
let pythonic_parser = PythonicParser::new();
|
||||
assert_eq!(
|
||||
pythonic_parser
|
||||
.parse_complete(plain_text)
|
||||
.await
|
||||
.unwrap()
|
||||
.1
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_incomplete_json() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let incomplete_cases = vec![
|
||||
r#"{"name": "test""#, // Missing closing brace
|
||||
r#"{"name": "test", "arguments":"#, // Incomplete arguments
|
||||
r#"{"name": "test", "arguments": {"#, // Incomplete nested object
|
||||
];
|
||||
|
||||
for input in incomplete_cases {
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Should not parse incomplete JSON: {}",
|
||||
input
|
||||
);
|
||||
}
|
||||
|
||||
// This case might actually parse because [{"name": "test"}] is complete
|
||||
// The trailing comma suggests more items but the first item is valid
|
||||
let _result = json_parser
|
||||
.parse_complete(r#"[{"name": "test"},"#)
|
||||
.await
|
||||
.unwrap();
|
||||
// This could parse the first element or return empty - implementation dependent
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_malformed_mistral() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
let malformed_cases = vec![
|
||||
"[TOOL_CALLS]", // Missing array
|
||||
"[TOOL_CALLS] {", // Not an array
|
||||
"[TOOL_CALLS] [", // Incomplete array
|
||||
"[TOOL_CALLS] [{]", // Invalid JSON in array
|
||||
"[TOOL_CALLS] [{\"name\": }]", // Invalid value
|
||||
];
|
||||
|
||||
for input in malformed_cases {
|
||||
// Parser might return error or empty vec for malformed input
|
||||
if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Should not parse malformed Mistral: {}",
|
||||
input
|
||||
);
|
||||
}
|
||||
// Error is also acceptable for malformed input
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_required_fields() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
// Missing name field
|
||||
let input = r#"{"arguments": {"x": 1}}"#;
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should not parse without name field");
|
||||
|
||||
// Name is not a string
|
||||
let input = r#"{"name": 123, "arguments": {}}"#;
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0, "Should not parse with non-string name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_very_long_strings() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let long_string = "x".repeat(10000);
|
||||
let input = format!(
|
||||
r#"{{"name": "test", "arguments": {{"data": "{}"}}}}"#,
|
||||
long_string
|
||||
);
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["data"].as_str().unwrap().len(), 10000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unicode_edge_cases() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
// Various Unicode characters including emojis, CJK, RTL text
|
||||
let input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍 مرحبا עולם"}}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Hello 世界 🌍 مرحبا עולם");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_nested_brackets_in_strings() {
|
||||
let mistral_parser = MistralParser::new();
|
||||
let input = r#"[TOOL_CALLS] [{"name": "echo", "arguments": {"text": "Array: [1, 2, 3]"}}]"#;
|
||||
let (_normal_text, tools) = mistral_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Array: [1, 2, 3]");
|
||||
|
||||
let pythonic_parser = PythonicParser::new();
|
||||
let input = r#"[echo(text="List: [a, b, c]")]"#;
|
||||
let (_normal_text, tools) = pythonic_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "List: [a, b, c]");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_formats_in_text() {
|
||||
let json_parser = JsonParser::new();
|
||||
let input = r#"
|
||||
Here's some text with [TOOL_CALLS] that shouldn't trigger.
|
||||
{"name": "actual_tool", "arguments": {}}
|
||||
And some more text with <tool_call> tags.
|
||||
"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "actual_tool");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_escaped_characters() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "write", "arguments": {"content": "Line 1\nLine 2\r\nLine 3\tTabbed\\Backslash\"Quote"}}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
let content = args["content"].as_str().unwrap();
|
||||
assert!(content.contains('\n'));
|
||||
assert!(content.contains('\t'));
|
||||
assert!(content.contains('\\'));
|
||||
assert!(content.contains('"'));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_numeric_edge_cases() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let input = r#"{
|
||||
"name": "calculate",
|
||||
"arguments": {
|
||||
"int": 42,
|
||||
"float": 123.456,
|
||||
"scientific": 1.23e-4,
|
||||
"negative": -999,
|
||||
"zero": 0,
|
||||
"large": 9007199254740991
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["int"], 42);
|
||||
assert_eq!(args["float"], 123.456);
|
||||
assert_eq!(args["scientific"], 0.000123);
|
||||
assert_eq!(args["negative"], -999);
|
||||
assert_eq!(args["zero"], 0);
|
||||
assert_eq!(args["large"], 9007199254740991i64);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_null_and_boolean_values() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let input = r#"{
|
||||
"name": "configure",
|
||||
"arguments": {
|
||||
"enabled": true,
|
||||
"disabled": false,
|
||||
"optional": null
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["disabled"], false);
|
||||
assert_eq!(args["optional"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_partial_token_at_buffer_boundary() {
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Send exactly "<tool" which is a 5-character prefix of "<tool_call>\n"
|
||||
let result = parser.parse_incremental("<tool", &tools).await.unwrap();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should be incomplete for partial tag"
|
||||
);
|
||||
|
||||
// Complete the token
|
||||
let result = parser
|
||||
.parse_incremental(
|
||||
"_call>\n{\"name\": \"test\", \"arguments\": {}}\n</tool_call>",
|
||||
&tools,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should successfully parse after completing
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "test");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_exact_prefix_lengths() {
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let test_cases = vec![
|
||||
("<", 1), // 1-char prefix
|
||||
("<t", 2), // 2-char prefix
|
||||
("<tool", 5), // 5-char prefix (the main bug case)
|
||||
("<tool_call", 10), // 10-char prefix
|
||||
("<tool_call>", 11), // 11-char prefix (full start without \n)
|
||||
];
|
||||
|
||||
for (prefix, expected_len) in test_cases {
|
||||
let result = parser.parse_incremental(prefix, &tools).await.unwrap();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Prefix '{}' (len {}) should be incomplete",
|
||||
prefix,
|
||||
expected_len
|
||||
);
|
||||
// Buffer is now internal to parser - can't assert on it
|
||||
}
|
||||
}
|
||||
@@ -1,272 +0,0 @@
|
||||
//! Tests for tool parser fallback behavior
|
||||
//!
|
||||
//! When tool call parsing fails, the original text should be preserved as normal text
|
||||
//! rather than being lost. This ensures graceful degradation.
|
||||
|
||||
use smg::tool_parser::{
|
||||
DeepSeekParser, JsonParser, LlamaParser, MistralParser, QwenParser, ToolParser,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parser_invalid_json_returns_as_normal_text() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Malformed JSON should be returned as normal text (note: commas may be processed)
|
||||
let input = r#"{"name": "test", "arguments": invalid json here}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
r#"{"name": "test", "arguments": invalid json here}"#
|
||||
);
|
||||
|
||||
// Plain text with no JSON structure should be returned as normal text
|
||||
let input = "This is just plain text that should not be parsed as a tool call";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
|
||||
// Text that looks like it might have JSON but doesn't should be returned as normal text
|
||||
let input = "The user said: {something} but it's not valid JSON";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_parser_invalid_format_returns_as_normal_text() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Missing closing tag
|
||||
let input = r#"<tool_call>
|
||||
{"name": "test", "arguments": {}}
|
||||
This text is missing the closing tag"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve original text when no valid tools found
|
||||
|
||||
// Malformed JSON inside valid tags
|
||||
let input = r#"<tool_call>
|
||||
{"name": "test", "arguments": invalid}
|
||||
</tool_call>"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
// When JSON parsing fails but tags are present, it should preserve the original text
|
||||
assert_eq!(normal_text, input);
|
||||
|
||||
// Plain text without any tool markers
|
||||
let input = "This is a regular response without any tool calls.";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return original text when no markers found
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_parser_invalid_format_returns_as_normal_text() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
// Invalid JSON after python_tag
|
||||
let input = r#"<|python_tag|>{"name": "test", "arguments": invalid}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve original text when parsing fails
|
||||
|
||||
// Plain text without markers or JSON
|
||||
let input = "Just explaining something without any function calls.";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return original text
|
||||
|
||||
// Text with python_tag but completely invalid content
|
||||
let input = r#"Here's my response <|python_tag|>not even close to JSON"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve everything when parsing fails
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_parser_invalid_format_returns_as_normal_text() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
// Missing closing bracket
|
||||
let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve original text when parsing fails
|
||||
|
||||
// Invalid JSON in tool calls section
|
||||
let input = r#"[TOOL_CALLS] [{"name": invalid json}]"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve original text when parsing fails
|
||||
|
||||
// Plain text
|
||||
let input = "No tool calls here, just regular text.";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return original text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_parser_invalid_format_returns_as_normal_text() {
|
||||
let parser = DeepSeekParser::new();
|
||||
|
||||
// Invalid JSON in tool call
|
||||
let input = r#"Some text<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>test
|
||||
```json
|
||||
{"name": "test", "arguments": malformed}
|
||||
```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should preserve original text when parsing fails
|
||||
|
||||
// Missing function marker
|
||||
let input = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>notfunction<｜tool▁sep｜>test
|
||||
```json
|
||||
{"x": 1}
|
||||
```<｜tool▁call▁end｜><｜tool▁calls▁end｜>"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return original text when parsing fails
|
||||
|
||||
// No tool markers at all
|
||||
let input = "Regular response without any special markers.";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return original text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mixed_valid_and_invalid_content() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Text with one valid tool call and one invalid
|
||||
let input = r#"Let me help you with that.
|
||||
<tool_call>
|
||||
{"name": "valid_tool", "arguments": {"x": 1}}
|
||||
</tool_call>
|
||||
And here's another one:
|
||||
<tool_call>
|
||||
{"name": "invalid_tool", "arguments": malformed}
|
||||
</tool_call>
|
||||
That's all!"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1); // Should extract the valid tool
|
||||
assert_eq!(tools[0].function.name, "valid_tool");
|
||||
// Normal text should contain text before the first tool call
|
||||
assert_eq!(normal_text, "Let me help you with that.\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_partial_tool_markers() {
|
||||
// Test cases where tool markers are incomplete or cut off
|
||||
|
||||
let parser = QwenParser::new();
|
||||
let input = "<tool_call>\nThis looks like it might be a tool call but it's not";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
|
||||
let parser = MistralParser::new();
|
||||
let input = "[TOOL_CALLS] But then nothing follows...";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
|
||||
let parser = LlamaParser::new();
|
||||
let input = "Starting a response <|python_tag|> but no JSON";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_escaped_json_like_content() {
|
||||
// Test that JSON-like content in regular text doesn't get parsed as tools
|
||||
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"The user typed: {"name": "example"} but this is just quoted text"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
// JsonParser should extract the valid JSON and return normal text
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "example");
|
||||
assert_eq!(normal_text, "The user typed: but this is just quoted text");
|
||||
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"The syntax is: <tool_call>
|
||||
{"name": "example"}
|
||||
</tool_call> - that's how you format it"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
// This actually contains valid tool call syntax, so it should parse
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "example");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unicode_and_special_chars_in_failed_parsing() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Unicode in malformed tool calls
|
||||
let input = r#"<tool_call>
|
||||
{"name": "测试", "arguments": 🚀 invalid}
|
||||
</tool_call>"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
// Should handle Unicode properly in the fallback text - malformed content should be preserved
|
||||
assert_eq!(normal_text, input);
|
||||
|
||||
// Special characters that might confuse parsers
|
||||
let input = r#"Response: <tool_call>{"name": "test\n\t", "arguments": {"]}"}</tool_call>"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
// This might or might not parse depending on JSON handling of escape sequences
|
||||
if tools.is_empty() {
|
||||
assert!(!normal_text.is_empty() || normal_text == input);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_very_long_invalid_input() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Generate a very long string that looks like it might be JSON but isn't
|
||||
let mut input = String::from("{\"name\": \"test\", \"arguments\": {");
|
||||
for i in 0..1000 {
|
||||
input.push_str(&format!("\"field{}\": \"value{}\", ", i, i));
|
||||
}
|
||||
input.push_str("\"final\": incomplete"); // Don't close the JSON properly
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Invalid JSON should be returned as normal text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_almost_valid_tool_calls() {
|
||||
// Test tool calls that are almost valid but have small issues
|
||||
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Missing closing quote should be returned as normal text
|
||||
let input = r#"{"name": "test", "arguments": {"key": "value}}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
r#"{"name": "test", "arguments": {"key": "value}}"#
|
||||
);
|
||||
|
||||
// Extra comma
|
||||
let input = r#"{"name": "test", "arguments": {},}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
// Some JSON parsers might accept trailing commas
|
||||
if tools.is_empty() {
|
||||
assert_eq!(normal_text, r#"{"name": "test", "arguments": {},}"#);
|
||||
}
|
||||
|
||||
// Wrong quote types
|
||||
let input = r#"{'name': 'test', 'arguments': {}}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0); // Standard JSON requires double quotes
|
||||
assert_eq!(normal_text, r#"{'name': 'test', 'arguments': {}}"#);
|
||||
}
|
||||
@@ -1,131 +0,0 @@
|
||||
//! GLM-4.7 MoE Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{Glm4MoeParser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm47_complete_parsing() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
let input = r#"Let me search for that.
|
||||
<tool_call>get_weather<arg_key>city</arg_key><arg_value>Beijing</arg_value><arg_key>date</arg_key><arg_value>2024-12-25</arg_value></tool_call>
|
||||
The weather will be..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me search for that.\n");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Beijing");
|
||||
assert_eq!(args["date"], "2024-12-25");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm47_multiple_tools() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
let input = r#"<tool_call>search<arg_key>query</arg_key><arg_value>rust tutorials</arg_value></tool_call><tool_call>translate<arg_key>text</arg_key><arg_value>Hello World</arg_value><arg_key>target_lang</arg_key><arg_value>zh</arg_value></tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm47_type_conversion() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
let input = r#"<tool_call>process<arg_key>count</arg_key><arg_value>42</arg_value><arg_key>rate</arg_key><arg_value>1.5</arg_value><arg_key>enabled</arg_key><arg_value>true</arg_value><arg_key>data</arg_key><arg_value>null</arg_value><arg_key>text</arg_key><arg_value>string value</arg_value></tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["count"], 42);
|
||||
assert_eq!(args["rate"], 1.5);
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["data"], serde_json::Value::Null);
|
||||
assert_eq!(args["text"], "string value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm47_streaming() {
|
||||
let mut parser = Glm4MoeParser::glm47();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
"get_weather",
|
||||
"<arg_key>city</arg_key>",
|
||||
"<arg_value>Shanghai</arg_value>",
|
||||
"<arg_key>units</arg_key>",
|
||||
"<arg_value>celsius</arg_value>",
|
||||
"</tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_glm47_format_detection() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
// Should detect GLM-4 format
|
||||
assert!(parser.has_tool_markers("<tool_call>"));
|
||||
assert!(parser.has_tool_markers("text with <tool_call> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_python_literals() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
let input = r#"<tool_call>test_func<arg_key>bool_true</arg_key><arg_value>True</arg_value><arg_key>bool_false</arg_key><arg_value>False</arg_value><arg_key>none_val</arg_key><arg_value>None</arg_value></tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["bool_true"], true);
|
||||
assert_eq!(args["bool_false"], false);
|
||||
assert_eq!(args["none_val"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm47_nested_json_in_arg_values() {
|
||||
let parser = Glm4MoeParser::glm47();
|
||||
|
||||
let input = r#"<tool_call>process<arg_key>data</arg_key><arg_value>{"nested": {"key": "value"}}</arg_value><arg_key>list</arg_key><arg_value>[1, 2, 3]</arg_value></tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["data"].is_object());
|
||||
assert!(args["list"].is_array());
|
||||
}
|
||||
@@ -1,168 +0,0 @@
|
||||
//! GLM-4 MoE Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{Glm4MoeParser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_complete_parsing() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
let input = r#"Let me search for that.
|
||||
<tool_call>get_weather
|
||||
<arg_key>city</arg_key>
|
||||
<arg_value>Beijing</arg_value>
|
||||
<arg_key>date</arg_key>
|
||||
<arg_value>2024-12-25</arg_value>
|
||||
</tool_call>
|
||||
The weather will be..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me search for that.\n");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Beijing");
|
||||
assert_eq!(args["date"], "2024-12-25");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_multiple_tools() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
let input = r#"<tool_call>search
|
||||
<arg_key>query</arg_key>
|
||||
<arg_value>rust tutorials</arg_value>
|
||||
</tool_call>
|
||||
<tool_call>translate
|
||||
<arg_key>text</arg_key>
|
||||
<arg_value>Hello World</arg_value>
|
||||
<arg_key>target_lang</arg_key>
|
||||
<arg_value>zh</arg_value>
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_type_conversion() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
let input = r#"<tool_call>process
|
||||
<arg_key>count</arg_key>
|
||||
<arg_value>42</arg_value>
|
||||
<arg_key>rate</arg_key>
|
||||
<arg_value>1.5</arg_value>
|
||||
<arg_key>enabled</arg_key>
|
||||
<arg_value>true</arg_value>
|
||||
<arg_key>data</arg_key>
|
||||
<arg_value>null</arg_value>
|
||||
<arg_key>text</arg_key>
|
||||
<arg_value>string value</arg_value>
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["count"], 42);
|
||||
assert_eq!(args["rate"], 1.5);
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["data"], serde_json::Value::Null);
|
||||
assert_eq!(args["text"], "string value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_streaming() {
|
||||
let mut parser = Glm4MoeParser::glm45();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
"get_weather\n",
|
||||
"<arg_key>city</arg_key>\n",
|
||||
"<arg_value>Shanghai</arg_value>\n",
|
||||
"<arg_key>units</arg_key>\n",
|
||||
"<arg_value>celsius</arg_value>\n",
|
||||
"</tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_glm4_format_detection() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
// Should detect GLM-4 format
|
||||
assert!(parser.has_tool_markers("<tool_call>"));
|
||||
assert!(parser.has_tool_markers("text with <tool_call> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_python_literals() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
let input = r#"<tool_call>test_func
|
||||
<arg_key>bool_true</arg_key>
|
||||
<arg_value>True</arg_value>
|
||||
<arg_key>bool_false</arg_key>
|
||||
<arg_value>False</arg_value>
|
||||
<arg_key>none_val</arg_key>
|
||||
<arg_value>None</arg_value>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["bool_true"], true);
|
||||
assert_eq!(args["bool_false"], false);
|
||||
assert_eq!(args["none_val"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_nested_json_in_arg_values() {
|
||||
let parser = Glm4MoeParser::glm45();
|
||||
|
||||
let input = r#"<tool_call>process
|
||||
<arg_key>data</arg_key>
|
||||
<arg_value>{"nested": {"key": "value"}}</arg_value>
|
||||
<arg_key>list</arg_key>
|
||||
<arg_value>[1, 2, 3]</arg_value>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["data"].is_object());
|
||||
assert!(args["list"].is_array());
|
||||
}
|
||||
@@ -1,716 +0,0 @@
|
||||
//! JSON Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the JSON parser which handles OpenAI, Claude, and generic JSON formats
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{JsonParser, ToolParser};
|
||||
|
||||
use crate::common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_simple_json_tool_call() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["location"], "San Francisco");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_array_of_tools() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"Hello, here are the results: [
|
||||
{"name": "get_weather", "arguments": {"location": "SF"}},
|
||||
{"name": "search", "arguments": {"query": "news"}}
|
||||
]"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "Hello, here are the results: ");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert_eq!(tools[1].function.name, "search");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_parameters_key() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "calculate");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["x"], 10);
|
||||
assert_eq!(args["y"], 20);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_extraction_from_text() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"I'll help you with that. {"name": "search", "arguments": {"query": "rust"}} Let me search for that."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"I'll help you with that. Let me search for that."
|
||||
);
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_nested_objects() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{
|
||||
"name": "update_config",
|
||||
"arguments": {
|
||||
"settings": {
|
||||
"theme": "dark",
|
||||
"language": "en",
|
||||
"notifications": {
|
||||
"email": true,
|
||||
"push": false
|
||||
}
|
||||
}
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "update_config");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["settings"]["theme"], "dark");
|
||||
assert_eq!(args["settings"]["notifications"]["email"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_special_characters() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "echo", "arguments": {"text": "Line 1\nLine 2\tTabbed", "path": "C:\\Users\\test"}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Line 1\nLine 2\tTabbed");
|
||||
assert_eq!(args["path"], "C:\\Users\\test");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_unicode() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍", "emoji": "😊"}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Hello 世界 🌍");
|
||||
assert_eq!(args["emoji"], "😊");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_empty_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "ping", "arguments": {}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args, json!({}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_invalid_format() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Missing closing brace
|
||||
let input = r#"{"name": "test", "arguments": {"key": "value""#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"{\"name\": \"test\", \"arguments\": {\"key\": \"value\""
|
||||
);
|
||||
|
||||
// Not JSON at all
|
||||
let input = "This is just plain text";
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_format_detection() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "arguments": {}}"#));
|
||||
assert!(parser.has_tool_markers(r#"[{"name": "test"}]"#));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
// Streaming tests for JSON array format
|
||||
#[tokio::test]
|
||||
async fn test_json_array_streaming_required_mode() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test that simulates the exact streaming pattern from required mode
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
// Define test tools
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
// Simulate the EXACT chunks from the debug log
|
||||
let chunks = vec![
|
||||
"[{",
|
||||
" \"",
|
||||
"name",
|
||||
"\":",
|
||||
" \"",
|
||||
"get",
|
||||
"_weather",
|
||||
"\",",
|
||||
" \"",
|
||||
"parameters",
|
||||
"\":",
|
||||
" {",
|
||||
" \"",
|
||||
"city",
|
||||
"\":",
|
||||
" \"",
|
||||
"Paris",
|
||||
"\"",
|
||||
" }",
|
||||
" }]",
|
||||
];
|
||||
|
||||
let mut all_results = Vec::new();
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_results.extend(result.calls);
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// We should have gotten tool call chunks
|
||||
assert!(
|
||||
!all_results.is_empty(),
|
||||
"Should have emitted tool call chunks"
|
||||
);
|
||||
|
||||
// Should not have emitted any normal text (including the closing ])
|
||||
assert_eq!(
|
||||
all_normal_text, "",
|
||||
"Should not emit normal text for JSON array format"
|
||||
);
|
||||
|
||||
// Check that we got the function name
|
||||
let has_name = all_results
|
||||
.iter()
|
||||
.any(|item| item.name.as_ref().is_some_and(|n| n == "get_weather"));
|
||||
assert!(has_name, "Should have emitted function name");
|
||||
|
||||
// Check that we got the parameters
|
||||
let has_params = all_results.iter().any(|item| !item.parameters.is_empty());
|
||||
assert!(has_params, "Should have emitted parameters");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_array_multiple_tools_streaming() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test with multiple tools in array
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let tools = vec![
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_news".to_string(),
|
||||
description: Some("Get news".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
// Split into smaller, more realistic chunks
|
||||
let chunks = vec![
|
||||
"[{",
|
||||
"\"name\":",
|
||||
"\"get_weather\"",
|
||||
",\"parameters\":",
|
||||
"{\"city\":",
|
||||
"\"SF\"}",
|
||||
"}",
|
||||
",",
|
||||
"{\"name\":",
|
||||
"\"get_news\"",
|
||||
",\"parameters\":",
|
||||
"{\"topic\":",
|
||||
"\"tech\"}",
|
||||
"}]",
|
||||
];
|
||||
|
||||
let mut all_results = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_results.extend(result.calls);
|
||||
}
|
||||
|
||||
// Should have gotten tool calls for both functions
|
||||
let has_weather = all_results
|
||||
.iter()
|
||||
.any(|item| item.name.as_ref().is_some_and(|n| n == "get_weather"));
|
||||
let has_news = all_results
|
||||
.iter()
|
||||
.any(|item| item.name.as_ref().is_some_and(|n| n == "get_news"));
|
||||
|
||||
assert!(has_weather, "Should have get_weather tool call");
|
||||
assert!(has_news, "Should have get_news tool call");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_array_closing_bracket_separate_chunk() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test case where the closing ] comes as a separate chunk
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
// Closing ] as separate chunk, followed by normal text
|
||||
let chunks = vec![
|
||||
"[{",
|
||||
"\"",
|
||||
"name",
|
||||
"\":",
|
||||
"\"",
|
||||
"get",
|
||||
"_weather",
|
||||
"\",",
|
||||
"\"",
|
||||
"parameters",
|
||||
"\":",
|
||||
"{",
|
||||
"\"",
|
||||
"city",
|
||||
"\":",
|
||||
"\"",
|
||||
"Paris",
|
||||
"\"",
|
||||
"}",
|
||||
"}",
|
||||
"]",
|
||||
" Here's",
|
||||
" the",
|
||||
" weather",
|
||||
" info",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// Should emit only the third chunk as normal text, NOT the ]
|
||||
assert_eq!(
|
||||
all_normal_text, " Here's the weather info",
|
||||
"Should emit only normal text without ], got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_single_object_with_trailing_text() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test single object format (no array) with trailing text
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"{",
|
||||
"\"",
|
||||
"name",
|
||||
"\":",
|
||||
"\"",
|
||||
"get_weather",
|
||||
"\",",
|
||||
"\"",
|
||||
"parameters",
|
||||
"\":",
|
||||
"{",
|
||||
"\"city",
|
||||
"\":",
|
||||
"\"Paris",
|
||||
"\"}",
|
||||
"}",
|
||||
" Here's",
|
||||
" the",
|
||||
" weather",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// Should emit the trailing text as normal_text (no ] to strip for single object)
|
||||
assert_eq!(
|
||||
all_normal_text, " Here's the weather",
|
||||
"Should emit normal text for single object format, got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_single_object_with_bracket_in_text() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test that ] in normal text is NOT stripped for single object format
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"{",
|
||||
"\"name",
|
||||
"\":",
|
||||
"\"get_weather",
|
||||
"\",",
|
||||
"\"parameters",
|
||||
"\":",
|
||||
"{",
|
||||
"\"city",
|
||||
"\":",
|
||||
"\"Paris",
|
||||
"\"}",
|
||||
"}",
|
||||
"]",
|
||||
" Here's",
|
||||
" the",
|
||||
" weather",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// For single object format, ] should NOT be stripped (it's part of normal text)
|
||||
assert_eq!(
|
||||
all_normal_text, "] Here's the weather",
|
||||
"Should preserve ] in normal text for single object format, got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_array_bracket_in_text_after_tools() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test that ] in normal text AFTER array tools is preserved
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: serde_json::json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"[",
|
||||
"{",
|
||||
"\"name",
|
||||
"\":",
|
||||
"\"get_weather",
|
||||
"\",",
|
||||
"\"parameters",
|
||||
"\":",
|
||||
"{",
|
||||
"\"city",
|
||||
"\":",
|
||||
"\"Paris",
|
||||
"\"}",
|
||||
"}",
|
||||
"]",
|
||||
" Array",
|
||||
" notation:",
|
||||
" arr",
|
||||
"[",
|
||||
"0",
|
||||
"]",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// Should preserve ] in normal text after array tools complete
|
||||
assert_eq!(
|
||||
all_normal_text, " Array notation: arr[0]",
|
||||
"Should preserve ] in normal text after array tools, got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
// =============================================================================
|
||||
// REALISTIC STREAMING TESTS
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_bug_incomplete_tool_name_string() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
// This exact sequence triggered the bug:
|
||||
// Parser receives {"name": " and must NOT parse it as empty name
|
||||
let chunks = vec![
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"name"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#, // ← Critical moment: parser has {"name": "
|
||||
// At this point, partial_json should NOT allow incomplete strings
|
||||
// when current_tool_name_sent=false
|
||||
r#"search"#, // Use valid tool name from create_test_tools()
|
||||
r#"""#,
|
||||
r#", "#,
|
||||
r#"""#,
|
||||
r#"arguments"#,
|
||||
r#"""#,
|
||||
r#": {"#,
|
||||
r#"""#,
|
||||
r#"query"#,
|
||||
r#"""#,
|
||||
r#": "#,
|
||||
r#"""#,
|
||||
r#"rust programming"#,
|
||||
r#"""#,
|
||||
r#"}}"#,
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
let mut saw_empty_name = false;
|
||||
|
||||
for chunk in chunks.iter() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = &call.name {
|
||||
if name.is_empty() {
|
||||
saw_empty_name = true;
|
||||
}
|
||||
if name == "search" {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
!saw_empty_name,
|
||||
"Parser should NEVER return empty tool name"
|
||||
);
|
||||
assert!(got_tool_name, "Should have parsed tool name correctly");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_realistic_chunks_simple_tool() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "get_weather", "arguments": {"city": "Paris"}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 10, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_strategic_chunks_with_quotes() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "search", "arguments": {"query": "rust programming"}}"#;
|
||||
let chunks = create_strategic_chunks(input);
|
||||
|
||||
// Strategic chunks break after quotes and colons
|
||||
assert!(chunks.iter().any(|c| c.ends_with('"')));
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_incremental_arguments_streaming() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "search", "arguments": {"query": "test", "limit": 10}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
let mut tool_name_sent = false;
|
||||
let mut got_arguments = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
tool_name_sent = true;
|
||||
}
|
||||
if tool_name_sent && !call.parameters.is_empty() {
|
||||
got_arguments = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(tool_name_sent, "Should have sent tool name");
|
||||
assert!(got_arguments, "Should have sent arguments");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_very_long_url_in_arguments() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
// Simulate long URL arriving in many chunks
|
||||
let long_url = "https://example.com/very/long/path/".to_string() + &"segment/".repeat(50);
|
||||
let input = format!(
|
||||
r#"{{"name": "search", "arguments": {{"query": "{}"}}}}"#,
|
||||
long_url
|
||||
);
|
||||
let chunks = create_realistic_chunks(&input);
|
||||
|
||||
assert!(chunks.len() > 100, "Long URL should create many chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_unicode() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "search", "arguments": {"query": "Hello 世界 🌍"}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed with unicode");
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
//! Kimi K2 Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{KimiK2Parser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kimik2_complete_parsing() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"Let me help you with that.
|
||||
<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"location": "Tokyo", "units": "celsius"}<|tool_call_end|>
|
||||
<|tool_calls_section_end|>
|
||||
The weather in Tokyo is..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me help you with that.\n");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["location"], "Tokyo");
|
||||
assert_eq!(args["units"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kimik2_multiple_tools() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>{"query": "rust tutorials"}<|tool_call_end|>
|
||||
<|tool_call_begin|>functions.translate:1<|tool_call_argument_begin|>{"text": "Hello", "to": "ja"}<|tool_call_end|>
|
||||
<|tool_calls_section_end|>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kimik2_with_whitespace() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {"key": "value", "num": 42} <|tool_call_end|>
|
||||
<|tool_calls_section_end|>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["key"], "value");
|
||||
assert_eq!(args["num"], 42);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kimik2_streaming() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = KimiK2Parser::new();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<|tool_calls_section_begin|>\n",
|
||||
"<|tool_call_begin|>functions.",
|
||||
"calculate:0",
|
||||
"<|tool_call_argument_begin|>",
|
||||
r#"{"x": 10, "#,
|
||||
r#""y": 20}"#,
|
||||
"<|tool_call_end|>\n",
|
||||
"<|tool_calls_section_end|>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "calculate");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kimik2_format_detection() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
// Should detect Kimi K2 format
|
||||
assert!(parser.has_tool_markers("<|tool_calls_section_begin|>"));
|
||||
assert!(parser.has_tool_markers("text with <|tool_calls_section_begin|> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<tool_call>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_kimik2_sequential_indices() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|>functions.first:0<|tool_call_argument_begin|>{"param": "a"}<|tool_call_end|>
|
||||
<|tool_call_begin|>functions.second:1<|tool_call_argument_begin|>{"param": "b"}<|tool_call_end|>
|
||||
<|tool_call_begin|>functions.third:2<|tool_call_argument_begin|>{"param": "c"}<|tool_call_end|>
|
||||
<|tool_calls_section_end|>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 3);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "first");
|
||||
assert_eq!(tools[1].function.name, "second");
|
||||
assert_eq!(tools[2].function.name, "third");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_function_index_extraction() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"Text before tool calls.
|
||||
<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>{"query": "rust"}<|tool_call_end|>
|
||||
<|tool_call_begin|>functions.calc:1<|tool_call_argument_begin|>{"x": 10}<|tool_call_end|>
|
||||
<|tool_calls_section_end|>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "Text before tool calls.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "calc");
|
||||
// TODO: Verify indices are preserved: 0 and 1
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_namespace_extraction() {
|
||||
let parser = KimiK2Parser::new();
|
||||
|
||||
let input = r#"<|tool_calls_section_begin|>
|
||||
<|tool_call_begin|>api.tools.search:0<|tool_call_argument_begin|>{"q": "test"}<|tool_call_end|>
|
||||
<|tool_calls_section_end|>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "api.tools.search"); // Includes full namespace
|
||||
}
|
||||
@@ -1,454 +0,0 @@
|
||||
//! Llama Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the Llama parser which handles <|python_tag|> format and plain JSON
|
||||
|
||||
use smg::tool_parser::{LlamaParser, ToolParser};
|
||||
|
||||
use crate::common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_python_tag_format() {
|
||||
let parser = LlamaParser::new();
|
||||
let input = r#"Here are some results: <|python_tag|>{"name": "search", "parameters": {"query": "weather"}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(normal_text, "Here are some results: ");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["query"], "weather");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_with_semicolon_separation() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
let input = r#"<|python_tag|>{"name": "tool1", "parameters": {}};{"name": "tool2", "parameters": {"y": 2}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "tool1");
|
||||
assert_eq!(tools[1].function.name, "tool2");
|
||||
assert_eq!(normal_text, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_no_tool_calls() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
let input = "This is just plain text with no tool calls";
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_plain_json_fallback() {
|
||||
let parser = LlamaParser::new();
|
||||
let input = r#"{"name": "calculate", "parameters": {"x": 5, "y": 10}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "calculate");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["x"], 5);
|
||||
assert_eq!(args["y"], 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_with_text_before() {
|
||||
let parser = LlamaParser::new();
|
||||
let input = r#"Let me help you with that. <|python_tag|>{"name": "get_time", "parameters": {"timezone": "UTC"}}"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me help you with that. ");
|
||||
assert_eq!(tools[0].function.name, "get_time");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["timezone"], "UTC");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_with_nested_json() {
|
||||
let parser = LlamaParser::new();
|
||||
let input = r#"<|python_tag|>{
|
||||
"name": "update_settings",
|
||||
"parameters": {
|
||||
"preferences": {
|
||||
"theme": "dark",
|
||||
"language": "en"
|
||||
},
|
||||
"notifications": true
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "update_settings");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["preferences"]["theme"], "dark");
|
||||
assert_eq!(args["notifications"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_empty_arguments() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
// With python_tag
|
||||
let input = r#"<|python_tag|>{"name": "ping", "parameters": {}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
|
||||
// Plain JSON
|
||||
let input = r#"{"name": "ping", "parameters": {}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_format_detection() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
assert!(parser.has_tool_markers(r#"<|python_tag|>{"name": "test"}"#));
|
||||
assert!(parser.has_tool_markers(r#"{"name": "test", "parameters": {}}"#));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_invalid_json_after_tag() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
let input = r#"<|python_tag|>{"name": invalid}"#;
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, "<|python_tag|>{\"name\": invalid}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_real_world_output() {
|
||||
let parser = LlamaParser::new();
|
||||
|
||||
// Actual output from Llama 3.2 model - simplified for testing
|
||||
let input = r#"I'll search for that information for you.
|
||||
|
||||
<|python_tag|>{"name": "web_search", "parameters": {"query": "Llama 3.2 model capabilities", "num_results": 5, "search_type": "recent"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "web_search");
|
||||
|
||||
let formatted_input = r#"<|python_tag|>{
|
||||
"name": "get_current_time",
|
||||
"parameters": {
|
||||
"timezone": "America/New_York",
|
||||
"format": "ISO8601"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools2) = parser.parse_complete(formatted_input).await.unwrap();
|
||||
assert_eq!(tools2.len(), 1);
|
||||
assert_eq!(tools2[0].function.name, "get_current_time");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_single_json() {
|
||||
let parser = LlamaParser::new();
|
||||
let text = r#"{"name": "get_weather", "parameters": {"city": "Paris"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Paris");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_json_with_separator() {
|
||||
let parser = LlamaParser::new();
|
||||
let text = r#"<|python_tag|>{"name": "get_weather", "parameters": {"city": "Paris"}};{"name": "get_tourist_attractions", "parameters": {"city": "Paris"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
// Note: Current implementation may only parse the first one due to semicolon handling
|
||||
assert!(!tools.is_empty());
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_trailing_text() {
|
||||
let parser = LlamaParser::new();
|
||||
// Valid JSON with trailing text - LlamaParser doesn't support this mixed format
|
||||
let text = r#"{"name": "get_weather", "parameters": {}} Some follow-up text"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
// LlamaParser expects pure JSON or <|python_tag|> format, not JSON with trailing text
|
||||
// So this returns as normal text
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, text);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_then_valid_json() {
|
||||
let parser = LlamaParser::new();
|
||||
let text =
|
||||
r#"{"name": "get_weather", "parameters": {{"name": "get_weather", "parameters": {}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
// Should parse at least one valid JSON
|
||||
if !tools.is_empty() {
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_plain_text_only() {
|
||||
let parser = LlamaParser::new();
|
||||
let text = "This is just plain explanation text.";
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_with_python_tag_prefix() {
|
||||
let parser = LlamaParser::new();
|
||||
let text = r#"Some intro. <|python_tag|>{"name": "get_weather", "parameters": {}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
}
|
||||
|
||||
// STREAMING TESTS
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_simple() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
// Send complete JSON at once
|
||||
let full_json = r#"<|python_tag|>{"name": "search", "parameters": {"query": "weather"}}"#;
|
||||
|
||||
let result = parser.parse_incremental(full_json, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Expected tool call for complete JSON input"
|
||||
);
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_partial() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
// Stream in chunks
|
||||
let chunks = vec![
|
||||
r#"<|python"#,
|
||||
r#"_tag|>{"name": "#,
|
||||
r#""calculate", "#,
|
||||
r#""parameters": {"x": 10}"#,
|
||||
r#"}"#,
|
||||
];
|
||||
|
||||
let mut got_complete = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "calculate");
|
||||
got_complete = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_complete, "Should have completed parsing");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_plain_json() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
// Stream plain JSON without python_tag
|
||||
let chunks = vec![
|
||||
r#"{"name": "#,
|
||||
r#""search", "#,
|
||||
r#""parameters": "#,
|
||||
r#"{"query": "#,
|
||||
r#""test"}}"#,
|
||||
];
|
||||
|
||||
let mut got_complete = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "search");
|
||||
got_complete = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_complete, "Should have completed parsing");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_with_text_before() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
let chunks = vec![
|
||||
r#"Let me help you. "#,
|
||||
r#"<|python_tag|>"#,
|
||||
r#"{"name": "get_time","#,
|
||||
r#" "parameters": {"#,
|
||||
r#""timezone": "UTC"}}"#,
|
||||
];
|
||||
|
||||
let mut got_complete = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "get_time");
|
||||
got_complete = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_complete, "Should have completed parsing");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_multiple_tools() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
let text =
|
||||
r#"<|python_tag|>{"name": "func1", "parameters": {}};{"name": "func2", "parameters": {}}"#;
|
||||
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
// Should get first tool complete
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Expected first tool to be complete"
|
||||
);
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "func1");
|
||||
}
|
||||
|
||||
// Process remaining buffer to get second tool
|
||||
let result2 = parser.parse_incremental("", &tools).await.unwrap();
|
||||
if !result2.calls.is_empty() {
|
||||
if let Some(name) = &result2.calls[0].name {
|
||||
assert_eq!(name, "func2");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_multiple_tools_chunked() {
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First chunk - incomplete first JSON
|
||||
let chunk1 = r#"<|python_tag|>{"name": "get_weather", "parameters""#;
|
||||
let result1 = parser.parse_incremental(chunk1, &tools).await.unwrap();
|
||||
if !result1.calls.is_empty() {
|
||||
if let Some(name) = &result1.calls[0].name {
|
||||
assert_eq!(name, "get_weather");
|
||||
}
|
||||
}
|
||||
|
||||
// Second chunk - complete first JSON and separator
|
||||
let chunk2 = r#": {"city": "Paris"}};{"name": "#;
|
||||
let result2 = parser.parse_incremental(chunk2, &tools).await.unwrap();
|
||||
|
||||
// Should get parameters for first tool (name already sent in result1)
|
||||
if !result2.calls.is_empty() {
|
||||
let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["city"], "Paris");
|
||||
}
|
||||
|
||||
let chunk3 = r#""get_time", "parameters": {"timezone": "UTC"}}"#;
|
||||
let result3 = parser.parse_incremental(chunk3, &tools).await.unwrap();
|
||||
if !result3.calls.is_empty() {
|
||||
if let Some(name) = &result3.calls[0].name {
|
||||
assert_eq!(name, "get_time");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// REALISTIC STREAMING TESTS
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_realistic_chunks_with_python_tag() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
let input = r#"<|python_tag|>{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 15, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "calculate");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_python_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
// Python tag itself arrives in small chunks
|
||||
let chunks = vec![
|
||||
"<|p", "yth", "on_", "tag", "|>{", r#"""#, "na", r#"me""#, ": ", r#"""#, "sea", "rch",
|
||||
r#"""#, ", ", r#"""#, "par", "ame", "ter", "s", r#"""#, ": {", r#"""#, "q", r#"""#, ": ",
|
||||
r#"""#, "tes", "t", r#"""#, "}}",
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "search");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
@@ -1,779 +0,0 @@
|
||||
//! MiniMax M2 Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{MinimaxM2Parser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_complete_parsing() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"Let me search for that.
|
||||
<minimax:tool_call>
|
||||
<invoke name="get_weather">
|
||||
<parameter name="city">Beijing</parameter>
|
||||
<parameter name="date">2024-12-25</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>
|
||||
The weather will be..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me search for that.\n");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Beijing");
|
||||
assert_eq!(args["date"], "2024-12-25");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_multiple_tools() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="search">
|
||||
<parameter name="query">rust tutorials</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>
|
||||
<minimax:tool_call>
|
||||
<invoke name="translate">
|
||||
<parameter name="text">Hello World</parameter>
|
||||
<parameter name="target_lang">zh</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_type_conversion() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="count">42</parameter>
|
||||
<parameter name="rate">1.5</parameter>
|
||||
<parameter name="enabled">true</parameter>
|
||||
<parameter name="data">null</parameter>
|
||||
<parameter name="text">string value</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["count"], 42);
|
||||
assert_eq!(args["rate"], 1.5);
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["data"], serde_json::Value::Null);
|
||||
assert_eq!(args["text"], "string value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_streaming_basic() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<minimax:tool_call>",
|
||||
r#"<invoke name="get_weather">"#,
|
||||
r#"<parameter name="city">Shanghai</parameter>"#,
|
||||
r#"<parameter name="units">celsius</parameter>"#,
|
||||
"</invoke>",
|
||||
"</minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
let mut found_params = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
if !call.parameters.is_empty() {
|
||||
found_params = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
assert!(found_params, "Should have streamed parameters");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_minimax_format_detection() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
// Should detect MiniMax format
|
||||
assert!(parser.has_tool_markers("<minimax:tool_call>"));
|
||||
assert!(parser.has_tool_markers("text with <minimax:tool_call> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("<tool_call>")); // GLM4 format
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<｜tool▁calls▁begin｜>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_python_literals() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="test_func">
|
||||
<parameter name="bool_true">True</parameter>
|
||||
<parameter name="bool_false">False</parameter>
|
||||
<parameter name="none_val">None</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["bool_true"], true);
|
||||
assert_eq!(args["bool_false"], false);
|
||||
assert_eq!(args["none_val"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_nested_json_in_parameters() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="data">{"nested": {"key": "value"}}</parameter>
|
||||
<parameter name="list">[1, 2, 3]</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// JSON-like strings are kept as strings, not parsed as JSON
|
||||
// This matches the behavior of other parsers like GLM4 MOE
|
||||
assert!(args["data"].is_string());
|
||||
assert_eq!(args["data"], r#"{"nested": {"key": "value"}}"#);
|
||||
assert!(args["list"].is_string());
|
||||
assert_eq!(args["list"], "[1, 2, 3]");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_xml_entities() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="html"><div>content</div></parameter>
|
||||
<parameter name="text">Quote: "hello"</parameter>
|
||||
<parameter name="code">if (a && b) { }</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["html"], "<div>content</div>");
|
||||
assert_eq!(args["text"], "Quote: \"hello\"");
|
||||
assert_eq!(args["code"], "if (a && b) { }");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_streaming_partial_tags() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Chunks split mid-tag
|
||||
let chunks = vec![
|
||||
"<minimax:tool_c",
|
||||
"all><invoke na",
|
||||
r#"me="get_weather"><param"#,
|
||||
r#"eter name="city">Bei"#,
|
||||
"jing</parameter></inv",
|
||||
"oke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
let mut buffer = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
buffer.push_str(&result.normal_text);
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
found_name,
|
||||
"Should have parsed function name from partial chunks"
|
||||
);
|
||||
assert_eq!(buffer, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_streaming_incremental_json() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"<minimax:tool_call>",
|
||||
r#"<invoke name="get_weather">"#,
|
||||
r#"<parameter name="city">Paris</parameter>"#,
|
||||
r#"<parameter name="units">metric</parameter>"#,
|
||||
"</invoke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut json_fragments = Vec::new();
|
||||
let mut found_function = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(_name) = call.name {
|
||||
found_function = true;
|
||||
}
|
||||
if !call.parameters.is_empty() {
|
||||
json_fragments.push(call.parameters.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_function);
|
||||
|
||||
// Verify JSON was built incrementally
|
||||
assert!(!json_fragments.is_empty());
|
||||
|
||||
// First fragment should start with opening brace
|
||||
if let Some(first) = json_fragments.first() {
|
||||
assert!(
|
||||
first.starts_with('{'),
|
||||
"First JSON fragment should start with '{{': {}",
|
||||
first
|
||||
);
|
||||
}
|
||||
|
||||
// Last fragment should be closing brace
|
||||
if let Some(last) = json_fragments.last() {
|
||||
assert!(
|
||||
last.contains('}'),
|
||||
"Last JSON fragment should contain '}}': {}",
|
||||
last
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_multiple_tools_boundary() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Tool boundary at chunk boundary
|
||||
let chunks = vec![
|
||||
r#"<minimax:tool_call><invoke name="get_weather"><parameter name="city">Tokyo</parameter></invoke></minimax:tool_call>"#,
|
||||
r#"<minimax:tool_call><invoke name="search"><parameter name="query">weather forecast</parameter></invoke></minimax:tool_call>"#,
|
||||
];
|
||||
|
||||
let mut tool_names = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
tool_names.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(tool_names.len(), 2);
|
||||
assert_eq!(tool_names[0], "get_weather");
|
||||
assert_eq!(tool_names[1], "search");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_invalid_function_name() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"<minimax:tool_call>",
|
||||
r#"<invoke name="invalid_function">"#,
|
||||
r#"<parameter name="param">value</parameter>"#,
|
||||
"</invoke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut found_invalid = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
// Invalid function should be skipped
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
if name == "invalid_function" {
|
||||
found_invalid = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(!found_invalid, "Invalid function should not be parsed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_empty_parameters() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="simple_func">
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "simple_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args, serde_json::json!({}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_multiline_parameter_values() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="multiline">line1
|
||||
line2
|
||||
line3</parameter>
|
||||
<parameter name="unicode">你好世界 🌍</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["multiline"], "line1\nline2\nline3");
|
||||
assert_eq!(args["unicode"], "你好世界 🌍");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_nested_xml_like_content() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="template"><html><body>Hello</body></html></parameter>
|
||||
<parameter name="config">{"key": "<value>nested</value>"}</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["template"], "<html><body>Hello</body></html>");
|
||||
|
||||
// The nested JSON with XML-like content
|
||||
let config =
|
||||
serde_json::from_str::<serde_json::Value>(args["config"].as_str().unwrap()).unwrap();
|
||||
assert_eq!(config["key"], "<value>nested</value>");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_streaming_state_reset() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First tool
|
||||
let chunks1 = vec![
|
||||
r#"<minimax:tool_call><invoke name="get_weather">"#,
|
||||
r#"<parameter name="city">London</parameter>"#,
|
||||
"</invoke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
for chunk in chunks1 {
|
||||
parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
}
|
||||
|
||||
// Second tool - state should be reset
|
||||
let chunks2 = vec![
|
||||
r#"<minimax:tool_call><invoke name="search">"#,
|
||||
r#"<parameter name="query">rust</parameter>"#,
|
||||
"</invoke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut second_tool_name = None;
|
||||
for chunk in chunks2 {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
second_tool_name = Some(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(second_tool_name, Some("search".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_many_parameters() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let mut params_xml = String::new();
|
||||
for i in 1..=20 {
|
||||
params_xml.push_str(&format!(
|
||||
r#"<parameter name="param{}">value{}</parameter>
|
||||
"#,
|
||||
i, i
|
||||
));
|
||||
}
|
||||
|
||||
let input = format!(
|
||||
r#"<minimax:tool_call>
|
||||
<invoke name="complex_func">
|
||||
{}
|
||||
</invoke>
|
||||
</minimax:tool_call>"#,
|
||||
params_xml
|
||||
);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "complex_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
|
||||
// Verify all 20 parameters are parsed
|
||||
for i in 1..=20 {
|
||||
let key = format!("param{}", i);
|
||||
let expected_value = format!("value{}", i);
|
||||
assert_eq!(args[key], expected_value);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_character_by_character_streaming() {
|
||||
// Test character-by-character streaming to simulate real-world streaming
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let complete_text = r#"Let me help you. <minimax:tool_call>
|
||||
<invoke name="get_weather">
|
||||
<parameter name="city">Seattle</parameter>
|
||||
<parameter name="units">celsius</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call> Here are the results."#;
|
||||
|
||||
let mut content_collected = String::new();
|
||||
let mut tool_name_found = false;
|
||||
let mut parameters_found = Vec::new();
|
||||
|
||||
// Stream character by character - feed only one character at a time
|
||||
for i in 0..complete_text.len() {
|
||||
let delta = &complete_text[i..i + 1];
|
||||
let result = parser.parse_incremental(delta, &tools).await.unwrap();
|
||||
content_collected.push_str(&result.normal_text);
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
tool_name_found = true;
|
||||
}
|
||||
if !call.parameters.is_empty() && !parameters_found.contains(&call.parameters) {
|
||||
parameters_found.push(call.parameters.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
tool_name_found,
|
||||
"Should find tool name during character-by-character streaming"
|
||||
);
|
||||
assert!(
|
||||
!parameters_found.is_empty(),
|
||||
"Should find parameters during streaming"
|
||||
);
|
||||
|
||||
// Should have initial content and final content
|
||||
assert!(content_collected.contains("Let me help you."));
|
||||
assert!(content_collected.contains("Here are the results."));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_content_before_and_after_tool_calls() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"I'll analyze the weather for you now.
|
||||
<minimax:tool_call>
|
||||
<invoke name="get_weather">
|
||||
<parameter name="city">Boston</parameter>
|
||||
<parameter name="state">MA</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>
|
||||
Based on the analysis, here's what I found."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Verify tool extraction
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
// Verify content preservation (only text before tool call is returned)
|
||||
assert!(normal_text.contains("I'll analyze the weather for you now."));
|
||||
// Text after tool call is not included in parse_complete
|
||||
assert!(!normal_text.contains("Based on the analysis, here's what I found."));
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Boston");
|
||||
assert_eq!(args["state"], "MA");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_incomplete_tool_call() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
// Incomplete tool call - missing closing tag
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="get_weather">
|
||||
<parameter name="city">Chicago</parameter>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Should not extract incomplete tool calls
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return as normal text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_malformed_invoke_tag() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
// Malformed invoke tag - missing name attribute
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke>
|
||||
<parameter name="city">Miami</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Should not extract tool calls with malformed invoke tags
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_streaming_with_invalid_function_progressive() {
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Progressive chunks building an invalid function call
|
||||
let chunks = vec![
|
||||
"<minimax:tool_call>",
|
||||
r#"<invoke name="invalid_function">"#,
|
||||
r#"<parameter name="test">value</parameter>"#,
|
||||
"</invoke>",
|
||||
"</minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
let mut found_valid_tool = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
// Should not get here for invalid function
|
||||
if tools.iter().any(|t| t.function.name == name) {
|
||||
found_valid_tool = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
!found_valid_tool,
|
||||
"Invalid function should not be parsed as tool call"
|
||||
);
|
||||
// The invalid tool call should be returned as normal text
|
||||
assert!(all_normal_text.contains("invalid_function"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_rapid_streaming_bursts() {
|
||||
// Test handling of rapid streaming bursts (multiple chunks at once)
|
||||
let mut parser = MinimaxM2Parser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"<minimax:tool_call><invoke name=\"search\"><parameter name=\"query\">",
|
||||
"rust programming",
|
||||
"</parameter></invoke></minimax:tool_call>",
|
||||
];
|
||||
|
||||
let mut found_function = false;
|
||||
let mut parameters = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "search");
|
||||
found_function = true;
|
||||
}
|
||||
if !call.parameters.is_empty() {
|
||||
parameters.push(call.parameters.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_function);
|
||||
|
||||
// Verify that parameters were streamed correctly
|
||||
let final_params = parameters.join("");
|
||||
assert!(final_params.contains("rust programming"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_special_characters_in_values() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="text">Special chars: @#$%^&*()</parameter>
|
||||
<parameter name="emoji">🦀 Rust 🚀</parameter>
|
||||
<parameter name="quotes">"double" and 'single' quotes</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Special chars: @#$%^&*()");
|
||||
assert_eq!(args["emoji"], "🦀 Rust 🚀");
|
||||
assert_eq!(args["quotes"], "\"double\" and 'single' quotes");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_whitespace_handling() {
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
// Test with various whitespace scenarios
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="trimmed"> spaces around </parameter>
|
||||
<parameter name="newlines">
|
||||
Line 1
|
||||
Line 2
|
||||
</parameter>
|
||||
<parameter name="tabs"> tab separated </parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// Values should preserve internal whitespace but may trim edges based on parser design
|
||||
assert!(args["newlines"].as_str().unwrap().contains("Line 1"));
|
||||
assert!(args["newlines"].as_str().unwrap().contains("Line 2"));
|
||||
assert_eq!(args["tabs"], "\ttab\tseparated\t");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_no_tools() {
|
||||
// Test input with no tool calls at all
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"This is just a normal response without any tool calls.
|
||||
I can provide information directly without using any tools.
|
||||
Even if I mention function names like get_weather or search,
|
||||
they are not actual tool calls unless properly formatted."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// No tools should be extracted
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Should not extract any tools from plain text"
|
||||
);
|
||||
|
||||
// All content should be returned as normal text
|
||||
assert_eq!(
|
||||
normal_text, input,
|
||||
"All content should be returned as normal text when no tools present"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minimax_invalid_json_in_parameters() {
|
||||
// Test handling of invalid JSON in parameter values
|
||||
let parser = MinimaxM2Parser::new();
|
||||
|
||||
let input = r#"<minimax:tool_call>
|
||||
<invoke name="process">
|
||||
<parameter name="valid">{"key": "value"}</parameter>
|
||||
<parameter name="invalid">{invalid json: no quotes}</parameter>
|
||||
<parameter name="broken">[1, 2, unclosed</parameter>
|
||||
<parameter name="mixed">Some text {"partial": json} more text</parameter>
|
||||
</invoke>
|
||||
</minimax:tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Tool should still be extracted despite invalid JSON in parameters
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
1,
|
||||
"Should extract tool even with invalid JSON in parameters"
|
||||
);
|
||||
assert_eq!(tools[0].function.name, "process");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
|
||||
// Parameters are stored as strings, not parsed as JSON
|
||||
// Even invalid JSON should be preserved as string values
|
||||
assert!(args["valid"].is_string());
|
||||
assert_eq!(args["valid"], r#"{"key": "value"}"#);
|
||||
|
||||
assert!(args["invalid"].is_string());
|
||||
assert_eq!(args["invalid"], "{invalid json: no quotes}");
|
||||
|
||||
assert!(args["broken"].is_string());
|
||||
assert_eq!(args["broken"], "[1, 2, unclosed");
|
||||
|
||||
assert!(args["mixed"].is_string());
|
||||
assert_eq!(args["mixed"], r#"Some text {"partial": json} more text"#);
|
||||
|
||||
assert_eq!(normal_text, "");
|
||||
}
|
||||
@@ -1,274 +0,0 @@
|
||||
//! Mistral Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the Mistral parser which handles [TOOL_CALLS] format
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{MistralParser, ToolParser};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_single_tool() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"Let me search for that.
|
||||
[TOOL_CALLS] [{"name": "search_web", "arguments": {"query": "latest news", "max_results": 5}}]"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me search for that.\n");
|
||||
assert_eq!(tools[0].function.name, "search_web");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["query"], "latest news");
|
||||
assert_eq!(args["max_results"], 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_multiple_tools() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"I'll help you with both tasks.
|
||||
[TOOL_CALLS] [
|
||||
{"name": "get_weather", "arguments": {"city": "Tokyo", "units": "celsius"}},
|
||||
{"name": "search_news", "arguments": {"query": "AI developments", "limit": 10}}
|
||||
]"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "I'll help you with both tasks.\n");
|
||||
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args0["city"], "Tokyo");
|
||||
|
||||
assert_eq!(tools[1].function.name, "search_news");
|
||||
let args1: serde_json::Value = serde_json::from_str(&tools[1].function.arguments).unwrap();
|
||||
assert_eq!(args1["query"], "AI developments");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_nested_json() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"Processing complex data.
|
||||
[TOOL_CALLS] [{"name": "process_data", "arguments": {"config": {"nested": {"value": [1, 2, 3]}}, "enabled": true}}]"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Processing complex data.\n");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["config"]["nested"]["value"], json!([1, 2, 3]));
|
||||
assert_eq!(args["enabled"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_with_text_after() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}}]
|
||||
|
||||
And here's some text after the tool call that should be ignored."#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_empty_arguments() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"[TOOL_CALLS] [{"name": "ping", "arguments": {}}]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_with_brackets_in_strings() {
|
||||
let parser = MistralParser::new();
|
||||
let input = r#"[TOOL_CALLS] [{"name": "echo", "arguments": {"text": "Array notation: arr[0] = value[1]"}}]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Array notation: arr[0] = value[1]");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_format_detection() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
assert!(parser.has_tool_markers("[TOOL_CALLS] ["));
|
||||
assert!(parser.has_tool_markers("Some text [TOOL_CALLS] ["));
|
||||
assert!(!parser.has_tool_markers("Just plain text"));
|
||||
assert!(!parser.has_tool_markers("[{\"name\": \"test\"}]")); // JSON array without TOOL_CALLS
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_malformed_json() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
// Missing closing bracket
|
||||
let input = r#"[TOOL_CALLS] [{"name": "test", "arguments": {}"#;
|
||||
if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
// Error is also acceptable for malformed input
|
||||
|
||||
// Invalid JSON inside
|
||||
let input = r#"[TOOL_CALLS] [{"name": invalid}]"#;
|
||||
if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
// Error is also acceptable for malformed input
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_real_world_output() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
// Actual output from Mistral model
|
||||
let input = r#"I'll search for information about Rust programming and check the weather in San Francisco.
|
||||
|
||||
[TOOL_CALLS] [
|
||||
{
|
||||
"name": "web_search",
|
||||
"arguments": {
|
||||
"query": "Rust programming language features 2024",
|
||||
"max_results": 3,
|
||||
"include_snippets": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "get_weather",
|
||||
"arguments": {
|
||||
"location": "San Francisco, CA",
|
||||
"units": "fahrenheit",
|
||||
"include_forecast": false
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
Let me execute these searches for you."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "I'll search for information about Rust programming and check the weather in San Francisco.\n\n");
|
||||
assert_eq!(tools[0].function.name, "web_search");
|
||||
assert_eq!(tools[1].function.name, "get_weather");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_streaming_closing_bracket() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test that closing ] is stripped for Mistral array format
|
||||
let mut parser = MistralParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"[TOOL_CALLS] ",
|
||||
"[{",
|
||||
"\"",
|
||||
"name",
|
||||
"\":",
|
||||
"\"",
|
||||
"get",
|
||||
"_weather",
|
||||
"\",",
|
||||
"\"",
|
||||
"arguments",
|
||||
"\":",
|
||||
"{",
|
||||
"\"",
|
||||
"city",
|
||||
"\":",
|
||||
"\"",
|
||||
"Paris",
|
||||
"\"",
|
||||
"}",
|
||||
"}",
|
||||
"]",
|
||||
" Here's",
|
||||
" the weather",
|
||||
" info",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// Should emit only the third chunk as normal text, NOT the ]
|
||||
assert_eq!(
|
||||
all_normal_text, " Here's the weather info",
|
||||
"Should not emit ] for Mistral array format, got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_streaming_bracket_in_text_after_tools() {
|
||||
use smg::protocols::common::Tool;
|
||||
|
||||
// Test that ] in normal text AFTER tool calls is preserved
|
||||
let mut parser = MistralParser::new();
|
||||
|
||||
let tools = vec![Tool {
|
||||
tool_type: "function".to_string(),
|
||||
function: smg::protocols::common::Function {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: json!({}),
|
||||
strict: None,
|
||||
},
|
||||
}];
|
||||
|
||||
let chunks = vec![
|
||||
"[TOOL_CALLS] ",
|
||||
"[",
|
||||
"{",
|
||||
"\"name",
|
||||
"\":",
|
||||
"\"get_weather",
|
||||
"\",",
|
||||
"\"arguments",
|
||||
"\":",
|
||||
"{\"",
|
||||
"city",
|
||||
"\":",
|
||||
"\"Paris",
|
||||
"\"}",
|
||||
"}",
|
||||
"]",
|
||||
" Array",
|
||||
" notation:",
|
||||
" arr",
|
||||
"[",
|
||||
"0",
|
||||
"]",
|
||||
];
|
||||
|
||||
let mut all_normal_text = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
all_normal_text.push_str(&result.normal_text);
|
||||
}
|
||||
|
||||
// Should preserve ] in normal text after tools complete
|
||||
assert_eq!(
|
||||
all_normal_text, " Array notation: arr[0]",
|
||||
"Should preserve ] in normal text after tools, got: '{}'",
|
||||
all_normal_text
|
||||
);
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
//! Mixed Format and Additional Edge Case Tests
|
||||
//!
|
||||
//! Tests for edge cases across parsers and mixed format scenarios
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{
|
||||
JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
|
||||
};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mixed_formats_in_text() {
|
||||
let json_parser = JsonParser::new();
|
||||
let input = r#"
|
||||
Some text with [TOOL_CALLS] marker that shouldn't trigger.
|
||||
Also has <tool_call> tags and [function()] syntax.
|
||||
But here's the actual JSON: {"name": "test", "arguments": {}}
|
||||
"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test");
|
||||
|
||||
// Mistral parser should ignore JSON and other formats
|
||||
let mistral_parser = MistralParser::new();
|
||||
let input = r#"
|
||||
{"name": "fake"} [function()] <tool_call>
|
||||
[TOOL_CALLS] [{"name": "real", "arguments": {}}]
|
||||
"#;
|
||||
|
||||
let (_normal_text, tools) = mistral_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "real");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_format_markers_in_string_content() {
|
||||
let pythonic_parser = PythonicParser::new();
|
||||
let input = r#"[echo(text="Use [TOOL_CALLS] and <tool_call> in text")]"#;
|
||||
|
||||
let (_normal_text, tools) = pythonic_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Use [TOOL_CALLS] and <tool_call> in text");
|
||||
|
||||
let qwen_parser = QwenParser::new();
|
||||
let input = r#"<tool_call>
|
||||
{"name": "log", "arguments": {"msg": "Found [function()] pattern"}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = qwen_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["msg"], "Found [function()] pattern");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_deeply_nested_json_structures() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let input = r#"{
|
||||
"name": "deep_process",
|
||||
"arguments": {
|
||||
"level1": {
|
||||
"level2": {
|
||||
"level3": {
|
||||
"level4": {
|
||||
"level5": {
|
||||
"data": [1, 2, [3, [4, 5]]]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "deep_process");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["level1"]["level2"]["level3"]["level4"]["level5"]["data"].is_array());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_sequential_calls_different_formats() {
|
||||
// Simulate a scenario where different parts of text have different formats
|
||||
// (though each parser will only recognize its own format)
|
||||
|
||||
let llama_parser = LlamaParser::new();
|
||||
|
||||
// Llama parser currently only returns the first tool found
|
||||
let input = r#"First call: <|python_tag|>{"name": "call1", "arguments": {}}"#;
|
||||
|
||||
let (_normal_text, tools) = llama_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "call1");
|
||||
|
||||
let input2 = r#"{"name": "call2", "arguments": {"x": 1}}"#;
|
||||
let (_normal_text2, tools2) = llama_parser.parse_complete(input2).await.unwrap();
|
||||
assert_eq!(tools2.len(), 1);
|
||||
assert_eq!(tools2[0].function.name, "call2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_and_whitespace_variations() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
// Various whitespace scenarios
|
||||
let cases = vec![
|
||||
r#" {"name":"compact","arguments":{}} "#,
|
||||
r#"
|
||||
|
||||
{"name": "spaced", "arguments": {}}
|
||||
|
||||
"#,
|
||||
r#" {"name": "tabbed", "arguments": {}} "#, // tabs
|
||||
];
|
||||
|
||||
for input in cases {
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1, "Should parse regardless of whitespace");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_special_json_values() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
let input = r#"{
|
||||
"name": "test_special",
|
||||
"arguments": {
|
||||
"float_e": 1.23e10,
|
||||
"float_neg_e": 1.23e-10,
|
||||
"hex_like": "0x1234",
|
||||
"very_long_num": 99999999999999999999,
|
||||
"special_strings": ["", " ", "\u0000", "\u001f"],
|
||||
"escaped": "\\n\\r\\t\\\"\\\\",
|
||||
"unicode": "\u4e2d\u6587"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "test_special");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["special_strings"].is_array());
|
||||
assert!(args["escaped"].is_string());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parser_recovery_after_invalid_input() {
|
||||
let mut parser = JsonParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Send invalid JSON first
|
||||
let _ = parser.parse_incremental(r#"{"broken": "#, &tools).await;
|
||||
|
||||
// Create a new parser instance for clean state
|
||||
let mut parser2 = JsonParser::new();
|
||||
let result = parser2
|
||||
.parse_incremental(r#"{"name": "valid", "arguments": {}}"#, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "valid");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_boundary_cases_for_extraction() {
|
||||
let json_parser = JsonParser::new();
|
||||
|
||||
// JSON at the very beginning
|
||||
let input = r#"{"name": "start", "arguments": {}} and then text"#;
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "start");
|
||||
|
||||
// JSON at the very end
|
||||
let input = r#"Some text first {"name": "end", "arguments": {}}"#;
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "end");
|
||||
|
||||
// Multiple JSON objects in text (should find first valid one)
|
||||
let input =
|
||||
r#"Text {"name": "first", "arguments": {}} more {"name": "second", "arguments": {}}"#;
|
||||
let (_normal_text, tools) = json_parser.parse_complete(input).await.unwrap();
|
||||
assert!(!tools.is_empty());
|
||||
assert_eq!(tools[0].function.name, "first");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_edge_cases() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
// Function name with underscores and numbers
|
||||
let input = r#"[func_name_2(param_1="value")]"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "func_name_2");
|
||||
|
||||
// Empty string argument
|
||||
let input = r#"[process(text="")]"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_with_pretty_json() {
|
||||
let parser = MistralParser::new();
|
||||
|
||||
// Pretty-printed JSON in Mistral format
|
||||
let input = r#"[TOOL_CALLS] [
|
||||
{
|
||||
"name": "formatted",
|
||||
"arguments": {
|
||||
"nested": {
|
||||
"key": "value"
|
||||
},
|
||||
"array": [
|
||||
1,
|
||||
2,
|
||||
3
|
||||
]
|
||||
}
|
||||
}
|
||||
]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "formatted");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["nested"]["key"], "value");
|
||||
assert_eq!(args["array"], json!([1, 2, 3]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_with_cdata_like_content() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Note: QwenParser expects exactly "<tool_call>\n" with the newline
|
||||
let input = r#"<tool_call>
|
||||
{"name": "process", "arguments": {"xml": "<![CDATA[some data]]>"}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["xml"], "<![CDATA[some data]]>");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extremely_long_function_names() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let long_name = "very_long_function_name_that_might_appear_in_generated_code_somewhere";
|
||||
let input = format!(r#"[{}(param="value")]"#, long_name);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, long_name);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_with_duplicate_keys() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// JSON with duplicate keys (last one should win per JSON spec)
|
||||
let input = r#"{"name": "test", "arguments": {"key": "first", "key": "second"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// JSON parsers typically keep the last value for duplicate keys
|
||||
assert_eq!(args["key"], "second");
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
//! Partial JSON Parser Tests
|
||||
//!
|
||||
//! Tests for the partial JSON parser with allow_partial_strings flag behavior
|
||||
|
||||
use smg::tool_parser::partial_json::PartialJson;
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_disallows_incomplete_strings() {
|
||||
// Test case from the bug report: {"name": "
|
||||
// With allow_partial_strings=false, should return {} (stop before incomplete string)
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should parse just the opening brace and stop at the incomplete string
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
|
||||
// Should have empty object (stopped before parsing incomplete "name" key)
|
||||
assert!(
|
||||
obj_map.is_empty() || !obj_map.contains_key("name"),
|
||||
"Should not parse incomplete string key, got: {:?}",
|
||||
obj_map
|
||||
);
|
||||
|
||||
// Should consume characters up to the incomplete string
|
||||
assert!(consumed <= input.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_allows_incomplete_strings() {
|
||||
// Test case: {"name": "
|
||||
// With allow_partial_strings=true, should parse the incomplete string
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should parse the object with incomplete string value
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
|
||||
// With allow_partial_strings=true, should parse "name" key with empty string value
|
||||
assert!(
|
||||
obj_map.contains_key("name"),
|
||||
"Should parse incomplete string with allow_partial_strings=true"
|
||||
);
|
||||
|
||||
assert_eq!(consumed, input.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_complete_json() {
|
||||
// Test case: {"name": "test"}
|
||||
// Both flags should parse complete JSON the same way
|
||||
let input = r#"{"name": "test"}"#;
|
||||
|
||||
let parser = PartialJson::new(32, true);
|
||||
let result1 = parser.parse_value(input, false);
|
||||
assert!(result1.is_ok());
|
||||
let (obj1, consumed1) = result1.unwrap();
|
||||
|
||||
let result2 = parser.parse_value(input, true);
|
||||
assert!(result2.is_ok());
|
||||
let (obj2, consumed2) = result2.unwrap();
|
||||
|
||||
// Both should parse the same complete JSON
|
||||
assert_eq!(obj1, obj2);
|
||||
assert_eq!(consumed1, consumed2);
|
||||
assert_eq!(consumed1, input.len());
|
||||
|
||||
// Check the parsed value
|
||||
assert!(obj1.is_object());
|
||||
let obj_map = obj1.as_object().unwrap();
|
||||
assert_eq!(obj_map.get("name").and_then(|v| v.as_str()), Some("test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backward_compatibility_default() {
|
||||
// Test that default PartialJson still allows partial strings (backward compatible)
|
||||
let parser = PartialJson::default();
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, _) = result.unwrap();
|
||||
assert!(obj.is_object());
|
||||
|
||||
// Default behavior should allow partial strings
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
assert!(
|
||||
obj_map.contains_key("name"),
|
||||
"Default should allow partial strings for backward compatibility"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_in_nested_object() {
|
||||
// Test case: {"tool": {"name": "
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"tool": {"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, _) = result.unwrap();
|
||||
assert!(obj.is_object());
|
||||
|
||||
// With allow_partial_strings=false, should stop before incomplete nested string
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
if let Some(tool) = obj_map.get("tool") {
|
||||
if let Some(tool_map) = tool.as_object() {
|
||||
assert!(
|
||||
!tool_map.contains_key("name")
|
||||
|| tool_map.get("name").and_then(|v| v.as_str()).is_none(),
|
||||
"Should not parse incomplete nested string"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug_fix_exact_scenario() {
|
||||
// This test verifies the exact bug scenario from the issue:
|
||||
// buffer = "{\"name\": \""
|
||||
// flags = Allow.ALL & ~Allow.STR
|
||||
// Python returns: Parsed object: {}, consumed length: 10
|
||||
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should return empty object (not {"name": null} or {"name": ""})
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
assert!(
|
||||
obj_map.is_empty(),
|
||||
"Expected empty object, got: {:?}. This matches Python behavior with Allow.ALL & ~Allow.STR",
|
||||
obj_map
|
||||
);
|
||||
|
||||
// Should consume all characters (10 bytes)
|
||||
assert_eq!(consumed, 10, "Should consume all 10 characters");
|
||||
}
|
||||
@@ -1,517 +0,0 @@
|
||||
//! Pythonic Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the Pythonic parser which handles Python function call syntax
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{PythonicParser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_single_function() {
|
||||
let parser = PythonicParser::new();
|
||||
let input = r#"[get_weather(city="London", units="celsius")]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "London");
|
||||
assert_eq!(args["units"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_multiple_functions() {
|
||||
let parser = PythonicParser::new();
|
||||
let input =
|
||||
r#"[search_web(query="Rust programming", max_results=5), get_time(timezone="UTC")]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "search_web");
|
||||
assert_eq!(tools[1].function.name, "get_time");
|
||||
|
||||
let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args0["query"], "Rust programming");
|
||||
assert_eq!(args0["max_results"], 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_with_python_literals() {
|
||||
let parser = PythonicParser::new();
|
||||
let input = r#"[configure(enabled=True, disabled=False, optional=None)]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["disabled"], false);
|
||||
assert_eq!(args["optional"], json!(null));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_with_lists_and_dicts() {
|
||||
let parser = PythonicParser::new();
|
||||
let input =
|
||||
r#"[process_data(items=[1, 2, 3], config={"key": "value", "nested": {"deep": True}})]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["items"], json!([1, 2, 3]));
|
||||
assert_eq!(args["config"]["key"], "value");
|
||||
assert_eq!(args["config"]["nested"]["deep"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_with_special_tokens() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
// Llama 4 sometimes outputs these tokens
|
||||
let input = r#"<|python_start|>[calculate(x=10, y=20)]<|python_end|>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "calculate");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["x"], 10);
|
||||
assert_eq!(args["y"], 20);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_with_nested_parentheses() {
|
||||
let parser = PythonicParser::new();
|
||||
let input = r#"[math_eval(expression="(2 + 3) * (4 - 1)", round_to=2)]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["expression"], "(2 + 3) * (4 - 1)");
|
||||
assert_eq!(args["round_to"], 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_with_escaped_quotes() {
|
||||
let parser = PythonicParser::new();
|
||||
let input = r#"[echo(text="She said \"Hello\" to him")]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "She said \"Hello\" to him");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_empty_arguments() {
|
||||
let parser = PythonicParser::new();
|
||||
let input = r#"[ping()]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args, json!({}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_format_detection() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
assert!(!parser.has_tool_markers("[function_name(")); // Incomplete
|
||||
assert!(parser.has_tool_markers("[get_weather(city=\"NYC\")]"));
|
||||
assert!(!parser.has_tool_markers("Just plain text"));
|
||||
assert!(!parser.has_tool_markers("{\"name\": \"test\"}")); // JSON
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_invalid_syntax() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
// Missing closing bracket
|
||||
let input = r#"[function(arg=value"#;
|
||||
if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
// Error is also acceptable for invalid syntax
|
||||
|
||||
// Invalid Python syntax - empty parameter name
|
||||
// Note: The parser currently accepts this invalid syntax and returns a result
|
||||
// This is a known limitation of the current implementation
|
||||
let input = r#"[function(=value)]"#;
|
||||
if let Ok((_normal_text, tools)) = parser.parse_complete(input).await {
|
||||
// The parser incorrectly accepts this, returning 1 result
|
||||
// We'll accept this behavior for now but note it's not ideal
|
||||
assert!(tools.len() <= 1, "Should parse at most one function");
|
||||
}
|
||||
// Error would be the correct behavior
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_real_world_llama4() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
// Actual output from Llama 4 model
|
||||
let input = r#"I'll help you with multiple tasks. Let me search for information and perform calculations.
|
||||
|
||||
[web_search(query="latest Rust features", max_results=3, safe_search=True),
|
||||
calculate(expression="42 * 3.14159", precision=2),
|
||||
get_weather(city="San Francisco", units="fahrenheit", include_forecast=False)]
|
||||
|
||||
These functions will provide the information you need."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 3);
|
||||
assert_eq!(normal_text, "I'll help you with multiple tasks. Let me search for information and perform calculations.\n\n\n\nThese functions will provide the information you need.");
|
||||
assert_eq!(tools[0].function.name, "web_search");
|
||||
assert_eq!(tools[1].function.name, "calculate");
|
||||
assert_eq!(tools[2].function.name, "get_weather");
|
||||
|
||||
let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args0["query"], "latest Rust features");
|
||||
assert_eq!(args0["safe_search"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_nested_brackets_in_lists() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let input = r#"[process_matrix(data=[[1, 2], [3, 4]], labels=["row[0]", "row[1]"])]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process_matrix");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["data"], json!([[1, 2], [3, 4]]));
|
||||
assert_eq!(args["labels"], json!(["row[0]", "row[1]"]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_nested_brackets_in_dicts() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let input =
|
||||
r#"[analyze(config={"patterns": ["[a-z]+", "[0-9]+"], "nested": {"list": [1, [2, 3]]}})]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "analyze");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["config"]["patterns"], json!(["[a-z]+", "[0-9]+"]));
|
||||
assert_eq!(args["config"]["nested"]["list"], json!([1, [2, 3]]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_mixed_quotes() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let input = r#"[format_text(single='Hello', double="World", mixed="It's \"quoted\"")]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "format_text");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["single"], "Hello");
|
||||
assert_eq!(args["double"], "World");
|
||||
assert_eq!(args["mixed"], "It's \"quoted\"");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_complex_nesting() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let input = r#"[transform(
|
||||
matrix=[[1, [2, 3]], [4, [5, [6, 7]]]],
|
||||
operations=[{"type": "scale", "factor": [2, 3]}, {"type": "rotate", "angle": 90}],
|
||||
metadata={"tags": ["nested[0]", "nested[1]"], "config": {"depth": [1, 2, 3]}}
|
||||
)]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "transform");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["matrix"].is_array());
|
||||
assert!(args["operations"].is_array());
|
||||
assert_eq!(args["operations"][0]["type"], "scale");
|
||||
assert_eq!(args["metadata"]["config"]["depth"], json!([1, 2, 3]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_no_brackets() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "This is just normal text without any tool calls.";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
// Expected - no tool calls found
|
||||
assert!(result.calls.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_complete_tool_call() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "Here's a tool call: [get_weather(location='New York', unit='celsius')]";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse complete tool call");
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "New York");
|
||||
assert_eq!(args["unit"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_text_before_tool_call() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "This is some text before [get_weather(location='London')]";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse tool call");
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "London");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_partial_tool_call() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First chunk with opening bracket but no closing bracket
|
||||
let text1 = "Let me check the weather: [get_weather(location=";
|
||||
let result1 = parser.parse_incremental(text1, &tools).await.unwrap();
|
||||
|
||||
// First chunk should be incomplete
|
||||
assert!(
|
||||
result1.calls.is_empty(),
|
||||
"First chunk should not return tool call"
|
||||
);
|
||||
|
||||
// Second chunk completing the tool call
|
||||
let text2 = "'Paris')]";
|
||||
let result2 = parser.parse_incremental(text2, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result2.calls.is_empty(),
|
||||
"Second chunk should complete tool call"
|
||||
);
|
||||
assert_eq!(result2.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "Paris");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_bracket_without_text_before() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "[search(query='python programming')]";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse tool call");
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["query"], "python programming");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_text_after_tool_call() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First chunk with complete tool call and some text after
|
||||
let text = "[get_weather(location='Tokyo')] Here's the forecast:";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse tool call");
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
// Text after tool call is handled by parser internally
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_multiple_tool_calls() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "[get_weather(location='Berlin'), search(query='restaurants')]";
|
||||
|
||||
// Current implementation may handle this as a single parse
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
// The parser should handle multiple tools in one bracket pair
|
||||
// This test is flexible about the implementation behavior
|
||||
if !result.calls.is_empty() {
|
||||
// Parser found at least one tool
|
||||
assert!(result.calls[0].name.is_some());
|
||||
}
|
||||
// Also acceptable if parser returns empty waiting for more context
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_opening_bracket_only() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "Let's try this: [";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
// Should be incomplete - no complete tool call
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should not return tool call for partial bracket"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_nested_brackets() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = "[get_weather(location='New York', unit='celsius', data=[1, 2, 3])]";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Should parse tool call with nested brackets"
|
||||
);
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "New York");
|
||||
assert_eq!(args["unit"], "celsius");
|
||||
assert_eq!(args["data"], json!([1, 2, 3]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_nested_brackets_dict() {
|
||||
let mut parser = PythonicParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text = r#"[search(query='test', config={'options': [1, 2], 'nested': {'key': 'value'}})]"#;
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Should parse tool call with nested dict"
|
||||
);
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "search");
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["query"], "test");
|
||||
assert_eq!(args["config"]["options"], json!([1, 2]));
|
||||
assert_eq!(args["config"]["nested"]["key"], "value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_multiple_tools_with_nested_brackets() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let text =
|
||||
"[get_weather(location='Paris', data=[10, 20]), search(query='test', filters=['a', 'b'])]";
|
||||
let result = parser.parse_incremental(text, &tools).await.unwrap();
|
||||
|
||||
// Should parse tools successfully
|
||||
if !result.calls.is_empty() {
|
||||
// At least gets the first tool
|
||||
assert!(result.calls[0].name.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_partial_nested_brackets() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First chunk with nested brackets but incomplete
|
||||
let text1 = "Here's a call: [get_weather(location='Tokyo', data=[1, 2";
|
||||
let result1 = parser.parse_incremental(text1, &tools).await.unwrap();
|
||||
|
||||
// First chunk should be incomplete
|
||||
assert!(result1.calls.is_empty(), "First chunk should not complete");
|
||||
|
||||
// Second chunk completing the nested brackets
|
||||
let text2 = ", 3])]";
|
||||
let result2 = parser.parse_incremental(text2, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result2.calls.is_empty(),
|
||||
"Second chunk should complete tool call"
|
||||
);
|
||||
assert_eq!(result2.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
let args: serde_json::Value = serde_json::from_str(&result2.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "Tokyo");
|
||||
assert_eq!(args["data"], json!([1, 2, 3]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_streaming_with_python_start_and_end_token() {
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"Here's a call: ",
|
||||
"<|python_",
|
||||
"start|>[get_weather(location=",
|
||||
"'Tokyo', data=[1, 2",
|
||||
", 3])]<|python_end|>",
|
||||
];
|
||||
|
||||
let mut got_tool = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "get_weather");
|
||||
let args: serde_json::Value =
|
||||
serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["location"], "Tokyo");
|
||||
assert_eq!(args["data"], json!([1, 2, 3]));
|
||||
got_tool = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool, "Should have parsed the tool call");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_detect_and_parse_with_python_start_and_end_token() {
|
||||
let parser = PythonicParser::new();
|
||||
|
||||
let text = "User wants to get the weather in Mars. <|python_start|>[get_weather(location='Mars', unit='celsius')]<|python_end|> In this way we will get the weather in Mars.";
|
||||
let (_normal_text, tools) = parser.parse_complete(text).await.unwrap();
|
||||
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["location"], "Mars");
|
||||
assert_eq!(args["unit"], "celsius");
|
||||
}
|
||||
@@ -1,306 +0,0 @@
|
||||
//! Qwen Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the Qwen parser which handles <tool_call>...</tool_call> format
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{QwenParser, ToolParser};
|
||||
|
||||
use crate::common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_single_tool() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"<tool_call>
|
||||
{"name": "get_weather", "arguments": {"city": "Beijing", "units": "celsius"}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Beijing");
|
||||
assert_eq!(args["units"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_multiple_sequential_tools() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"Let me help you with that.
|
||||
<tool_call>
|
||||
{"name": "search", "arguments": {"query": "Qwen model"}}
|
||||
</tool_call>
|
||||
<tool_call>
|
||||
{"name": "translate", "arguments": {"text": "Hello", "to": "zh"}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "Let me help you with that.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_pretty_printed_json() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"<tool_call>
|
||||
{
|
||||
"name": "create_document",
|
||||
"arguments": {
|
||||
"title": "Test Document",
|
||||
"content": "This is a test",
|
||||
"metadata": {
|
||||
"author": "Qwen",
|
||||
"tags": ["test", "example"]
|
||||
}
|
||||
}
|
||||
}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "create_document");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["metadata"]["author"], "Qwen");
|
||||
assert_eq!(args["metadata"]["tags"], json!(["test", "example"]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_with_text_between() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"First, let me search for information.
|
||||
<tool_call>
|
||||
{"name": "search", "arguments": {"query": "test"}}
|
||||
</tool_call>
|
||||
|
||||
Now I'll translate something.
|
||||
|
||||
<tool_call>
|
||||
{"name": "translate", "arguments": {"text": "world", "to": "es"}}
|
||||
</tool_call>
|
||||
Done!"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "First, let me search for information.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_empty_arguments() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"<tool_call>
|
||||
{"name": "get_time", "arguments": {}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_time");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_with_newlines_in_strings() {
|
||||
let parser = QwenParser::new();
|
||||
let input = r#"<tool_call>
|
||||
{"name": "write_file", "arguments": {"content": "Line 1\nLine 2\nLine 3", "path": "/tmp/test.txt"}}
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["content"], "Line 1\nLine 2\nLine 3");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_format_detection() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
assert!(parser.has_tool_markers("<tool_call>"));
|
||||
assert!(parser.has_tool_markers("Some text <tool_call>\n{"));
|
||||
assert!(!parser.has_tool_markers("Just plain text"));
|
||||
assert!(!parser.has_tool_markers("{\"name\": \"test\"}")); // Plain JSON
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_incomplete_tags() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Missing closing tag
|
||||
let input = r#"<tool_call>
|
||||
{"name": "test", "arguments": {}}"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
|
||||
// Missing opening tag
|
||||
let input = r#"{"name": "test", "arguments": {}}
|
||||
</tool_call>"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_real_world_output() {
|
||||
let parser = QwenParser::new();
|
||||
|
||||
// Actual output from Qwen model
|
||||
let input = r#"I'll help you search for information and perform calculations.
|
||||
|
||||
<tool_call>
|
||||
{
|
||||
"name": "web_search",
|
||||
"arguments": {
|
||||
"query": "quantum computing breakthroughs 2024",
|
||||
"language": "en",
|
||||
"region": "us",
|
||||
"safe_search": true
|
||||
}
|
||||
}
|
||||
</tool_call>
|
||||
|
||||
Let me also calculate something for you:
|
||||
|
||||
<tool_call>
|
||||
{
|
||||
"name": "calculator",
|
||||
"arguments": {
|
||||
"expression": "sqrt(144) + 3^2",
|
||||
"precision": 2
|
||||
}
|
||||
}
|
||||
</tool_call>
|
||||
|
||||
These tools will provide the information you need."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"I'll help you search for information and perform calculations.\n\n"
|
||||
);
|
||||
assert_eq!(tools[0].function.name, "web_search");
|
||||
assert_eq!(tools[1].function.name, "calculator");
|
||||
|
||||
let args0: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args0["query"], "quantum computing breakthroughs 2024");
|
||||
assert_eq!(args0["safe_search"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_buffer_drain_optimization() {
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First chunk - incomplete tool call
|
||||
let chunk1 = "<tool_call>\n{\"name\": \"test1\", ";
|
||||
let _result = parser.parse_incremental(chunk1, &tools).await.unwrap();
|
||||
// The important thing is buffer accumulation works
|
||||
|
||||
// Complete first tool and start second
|
||||
let chunk2 = "\"arguments\": {}}\n</tool_call><tool_call>\n{\"name\": \"test2\", ";
|
||||
let result = parser.parse_incremental(chunk2, &tools).await.unwrap();
|
||||
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(_name) = &result.calls[0].name {
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "test1");
|
||||
// After consuming the first tool, buffer is managed internally
|
||||
}
|
||||
}
|
||||
|
||||
// Complete the second tool
|
||||
let chunk3 = "\"arguments\": {\"x\": 1}}\n</tool_call>";
|
||||
let result = parser.parse_incremental(chunk3, &tools).await.unwrap();
|
||||
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(_name) = &result.calls[0].name {
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "test2");
|
||||
// Buffer is managed internally
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_buffer_efficiency_with_multiple_tools() {
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Send multiple complete tools at once
|
||||
let input = r#"<tool_call>
|
||||
{"name": "tool1", "arguments": {"a": 1}}
|
||||
</tool_call><tool_call>
|
||||
{"name": "tool2", "arguments": {"b": 2}}
|
||||
</tool_call><tool_call>
|
||||
{"name": "tool3", "arguments": {"c": 3}}
|
||||
</tool_call>"#;
|
||||
|
||||
// This should efficiently process tools using drain() without creating new strings
|
||||
let result = parser.parse_incremental(input, &tools).await.unwrap();
|
||||
|
||||
// In Phase 2, this will likely parse only the first tool
|
||||
// The important thing is that drain() doesn't cause any issues
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert!(["tool1", "tool2", "tool3"].contains(&name.as_str()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// REALISTIC STREAMING TESTS
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_realistic_chunks_with_xml_tags() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let input = "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>";
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 20, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_xml_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let chunks = vec![
|
||||
"<to", "ol_", "cal", "l>\n", "{", r#"""#, "na", "me", r#"""#, ": ", r#"""#, "tra", "nsl",
|
||||
"ate", r#"""#, ", ", r#"""#, "arg", "ume", "nts", r#"""#, ": {", r#"""#, "tex", "t",
|
||||
r#"""#, ": ", r#"""#, "hel", "lo", r#"""#, "}}\n", "</t", "ool", "_ca", "ll>",
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "translate");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
@@ -1,958 +0,0 @@
|
||||
//! Qwen Coder Parser Integration Tests
|
||||
//!
|
||||
//! Tests for the Qwen Coder parser which handles XML format:
|
||||
//! <tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>
|
||||
|
||||
use serde_json::json;
|
||||
use smg::tool_parser::{parsers::QwenCoderParser, traits::ToolParser};
|
||||
|
||||
use crate::common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_single_tool() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Beijing</parameter>
|
||||
<parameter=units>celsius</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Beijing");
|
||||
assert_eq!(args["units"], "celsius");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_multiple_sequential_tools() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"Let me help you with that.
|
||||
<tool_call>
|
||||
<function=search>
|
||||
<parameter=query>Qwen model</parameter>
|
||||
</function>
|
||||
</tool_call>
|
||||
<tool_call>
|
||||
<function=translate>
|
||||
<parameter=text>Hello</parameter>
|
||||
<parameter=to>zh</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "Let me help you with that.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
assert_eq!(tools[1].function.name, "translate");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_nested_json_in_parameters() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"<tool_call>
|
||||
<function=process_data>
|
||||
<parameter=config>{"nested": {"value": [1, 2, 3]}}</parameter>
|
||||
<parameter=enabled>true</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process_data");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// JSON values should be parsed
|
||||
assert_eq!(args["config"]["nested"]["value"], json!([1, 2, 3]));
|
||||
assert_eq!(args["enabled"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_string_parameters() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=text>Hello World</parameter>
|
||||
<parameter=number>42</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Hello World");
|
||||
// JSON numbers should be parsed as numbers (consistent with Python's json.loads)
|
||||
assert_eq!(args["number"], 42);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_empty_arguments() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"<tool_call>
|
||||
<function=get_time>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_time");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args, json!({}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_multiline_parameter_values() {
|
||||
let parser = QwenCoderParser::new();
|
||||
let input = r#"<tool_call>
|
||||
<function=write_file>
|
||||
<parameter=content>Line 1
|
||||
Line 2
|
||||
Line 3</parameter>
|
||||
<parameter=path>/tmp/test.txt</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["content"], "Line 1\nLine 2\nLine 3");
|
||||
assert_eq!(args["path"], "/tmp/test.txt");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_format_detection() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
assert!(parser.has_tool_markers("<tool_call>"));
|
||||
assert!(parser.has_tool_markers("Some text <tool_call>"));
|
||||
assert!(!parser.has_tool_markers("Just plain text"));
|
||||
assert!(!parser.has_tool_markers("<function=test>")); // Without tool_call tags
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_incomplete_tags() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Missing closing tag
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Beijing</parameter>"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
|
||||
// Missing opening tag
|
||||
let input = r#"<parameter=city>Beijing</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_streaming_basic() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
r#"<function=get_weather>"#,
|
||||
r#"<parameter=city>Shanghai</parameter>"#,
|
||||
r#"<parameter=units>celsius</parameter>"#,
|
||||
"</function>",
|
||||
"</tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
let mut found_params = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
if !call.parameters.is_empty() {
|
||||
found_params = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_name, "Should have found tool name during streaming");
|
||||
assert!(found_params, "Should have streamed parameters");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_streaming_incremental_json() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
r#"<function=get_weather>"#,
|
||||
r#"<parameter=city>Paris</parameter>"#,
|
||||
r#"<parameter=units>metric</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut json_fragments = Vec::new();
|
||||
let mut found_function = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(_name) = call.name {
|
||||
found_function = true;
|
||||
}
|
||||
if !call.parameters.is_empty() {
|
||||
json_fragments.push(call.parameters.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_function);
|
||||
|
||||
// Verify JSON was built incrementally
|
||||
assert!(!json_fragments.is_empty());
|
||||
|
||||
// First fragment should start with opening brace
|
||||
if let Some(first) = json_fragments.first() {
|
||||
assert!(
|
||||
first.starts_with('{'),
|
||||
"First JSON fragment should start with '{{': {}",
|
||||
first
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_streaming_partial_tags() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Chunks split mid-tag
|
||||
let chunks = vec![
|
||||
"<tool_c",
|
||||
"all><function=",
|
||||
r#"get_weather><param"#,
|
||||
r#"eter=city>Bei"#,
|
||||
"jing</parameter></func",
|
||||
"tion></tool_call>",
|
||||
];
|
||||
|
||||
let mut found_name = false;
|
||||
let mut buffer = String::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
buffer.push_str(&result.normal_text);
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
found_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
found_name,
|
||||
"Should have parsed function name from partial chunks"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_multiple_tools_boundary() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Tool boundary at chunk boundary
|
||||
let chunks = vec![
|
||||
r#"<tool_call><function=get_weather><parameter=city>Tokyo</parameter></function></tool_call>"#,
|
||||
r#"<tool_call><function=search><parameter=query>weather forecast</parameter></function></tool_call>"#,
|
||||
];
|
||||
|
||||
let mut tool_names = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
tool_names.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(tool_names.len(), 2);
|
||||
assert_eq!(tool_names[0], "get_weather");
|
||||
assert_eq!(tool_names[1], "search");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_invalid_function_name() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let chunks = vec![
|
||||
"<tool_call>",
|
||||
r#"<function=invalid_function>"#,
|
||||
r#"<parameter=param>value</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut found_invalid = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
// Invalid function should be skipped
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
if name == "invalid_function" {
|
||||
found_invalid = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(!found_invalid, "Invalid function should not be parsed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_type_conversion() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=count>42</parameter>
|
||||
<parameter=rate>1.5</parameter>
|
||||
<parameter=enabled>true</parameter>
|
||||
<parameter=data>null</parameter>
|
||||
<parameter=text>string value</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// JSON values should be parsed
|
||||
assert_eq!(args["count"], 42);
|
||||
assert_eq!(args["rate"], 1.5);
|
||||
assert_eq!(args["enabled"], true);
|
||||
assert_eq!(args["data"], serde_json::Value::Null);
|
||||
assert_eq!(args["text"], "string value");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_special_characters_in_values() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=text>Special chars: @#$%^&*()</parameter>
|
||||
<parameter=emoji>🦀 Rust 🚀</parameter>
|
||||
<parameter=quotes>"double" and 'single' quotes</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Special chars: @#$%^&*()");
|
||||
assert_eq!(args["emoji"], "🦀 Rust 🚀");
|
||||
assert_eq!(args["quotes"], "\"double\" and 'single' quotes");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_whitespace_handling() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Test with various whitespace scenarios
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=trimmed> spaces around </parameter>
|
||||
<parameter=newlines>
|
||||
Line 1
|
||||
Line 2
|
||||
</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// Values should preserve internal whitespace but trim edges
|
||||
assert_eq!(args["trimmed"], "spaces around");
|
||||
assert!(args["newlines"].as_str().unwrap().contains("Line 1"));
|
||||
assert!(args["newlines"].as_str().unwrap().contains("Line 2"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_no_tools() {
|
||||
// Test input with no tool calls at all
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"This is just a normal response without any tool calls.
|
||||
I can provide information directly without using any tools.
|
||||
Even if I mention function names like get_weather or search,
|
||||
they are not actual tool calls unless properly formatted."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// No tools should be extracted
|
||||
assert_eq!(
|
||||
tools.len(),
|
||||
0,
|
||||
"Should not extract any tools from plain text"
|
||||
);
|
||||
|
||||
// All content should be returned as normal text
|
||||
assert_eq!(
|
||||
normal_text, input,
|
||||
"All content should be returned as normal text when no tools present"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_streaming_state_reset() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First tool
|
||||
let chunks1 = vec![
|
||||
r#"<tool_call><function=get_weather>"#,
|
||||
r#"<parameter=city>London</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
for chunk in chunks1 {
|
||||
parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
}
|
||||
|
||||
// Second tool - state should be reset
|
||||
let chunks2 = vec![
|
||||
r#"<tool_call><function=search>"#,
|
||||
r#"<parameter=query>rust</parameter>"#,
|
||||
"</function></tool_call>",
|
||||
];
|
||||
|
||||
let mut second_tool_name = None;
|
||||
for chunk in chunks2 {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
second_tool_name = Some(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(second_tool_name, Some("search".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_realistic_chunks() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Tokyo</parameter>
|
||||
<parameter=units>celsius</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 20, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_xml_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = QwenCoderParser::new();
|
||||
|
||||
let chunks = vec![
|
||||
"<to", "ol_", "cal", "l>", "<fun", "cti", "on=", "get", "_we", "ath", "er>", "<par", "ame",
|
||||
"ter=", "cit", "y>", "Tok", "yo", "</", "par", "ame", "ter>", "</", "func", "tion>", "</",
|
||||
"too", "l_c", "all>",
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_content_before_and_after_tool_calls() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"I'll analyze the weather for you now.
|
||||
<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Boston</parameter>
|
||||
<parameter=state>MA</parameter>
|
||||
</function>
|
||||
</tool_call>
|
||||
Based on the analysis, here's what I found."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Verify tool extraction
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
// Verify content preservation (only text before tool call is returned)
|
||||
assert!(normal_text.contains("I'll analyze the weather for you now."));
|
||||
// Text after tool call is not included in parse_complete
|
||||
assert!(!normal_text.contains("Based on the analysis, here's what I found."));
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["city"], "Boston");
|
||||
assert_eq!(args["state"], "MA");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_incomplete_tool_call() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Incomplete tool call - missing closing tag
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Chicago</parameter>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Should not extract incomplete tool calls
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return as normal text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_malformed_function_tag() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Malformed function tag - missing name attribute
|
||||
let input = r#"<tool_call>
|
||||
<function>
|
||||
<parameter=city>Miami</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Should not extract tool calls with malformed function tags
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(normal_text, input);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_many_parameters() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let mut params_xml = String::new();
|
||||
for i in 1..=20 {
|
||||
params_xml.push_str(&format!(
|
||||
r#"<parameter=param{}>value{}</parameter>
|
||||
"#,
|
||||
i, i
|
||||
));
|
||||
}
|
||||
|
||||
let input = format!(
|
||||
r#"<tool_call>
|
||||
<function=complex_func>
|
||||
{}
|
||||
</function>
|
||||
</tool_call>"#,
|
||||
params_xml
|
||||
);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "complex_func");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
|
||||
// Verify all 20 parameters are parsed
|
||||
for i in 1..=20 {
|
||||
let key = format!("param{}", i);
|
||||
let expected_value = format!("value{}", i);
|
||||
assert_eq!(args[key], expected_value);
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Edge Case Tests
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_malformed_xml_missing_parameter_close() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Missing </parameter> closing tag - parser regex won't match incomplete parameter
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Beijing
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// The parser extracts the tool call but with empty arguments since
|
||||
// the parameter block is malformed (no </parameter>)
|
||||
// This is acceptable behavior - we extract what we can
|
||||
if tools.is_empty() {
|
||||
// If no tools extracted, input returned as normal text
|
||||
assert_eq!(normal_text, input);
|
||||
} else {
|
||||
// If tool extracted, it should have the function name
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_malformed_xml_unclosed_function() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Missing </function> closing tag
|
||||
let input = r#"<tool_call>
|
||||
<function=get_weather>
|
||||
<parameter=city>Beijing</parameter>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Parser should still extract the tool since it has complete tool_call tags
|
||||
// and the function name + parameters are present
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_malformed_xml_nested_tool_calls() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Nested tool_call tags (invalid)
|
||||
let input = r#"<tool_call>
|
||||
<function=outer>
|
||||
<tool_call>
|
||||
<function=inner>
|
||||
</function>
|
||||
</tool_call>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
// Should handle gracefully - may parse first complete tool_call
|
||||
// The exact behavior depends on regex matching
|
||||
assert!(tools.len() <= 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_unicode_parameter_names() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Unicode characters in parameter names (Chinese, Japanese, emoji)
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=城市>北京</parameter>
|
||||
<parameter=天気>晴れ</parameter>
|
||||
<parameter=emoji_key>🌍🌎🌏</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["城市"], "北京");
|
||||
assert_eq!(args["天気"], "晴れ");
|
||||
assert_eq!(args["emoji_key"], "🌍🌎🌏");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_unicode_function_name() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Unicode function name
|
||||
let input = r#"<tool_call>
|
||||
<function=获取天气>
|
||||
<parameter=location>上海</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "获取天气");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["location"], "上海");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_very_large_parameter_value() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Generate a large parameter value (100KB)
|
||||
let large_value: String = "x".repeat(100_000);
|
||||
|
||||
let input = format!(
|
||||
r#"<tool_call>
|
||||
<function=process_large>
|
||||
<parameter=data>{}</parameter>
|
||||
</function>
|
||||
</tool_call>"#,
|
||||
large_value
|
||||
);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process_large");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["data"].as_str().unwrap().len(), 100_000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_very_large_nested_json_parameter() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Generate moderately nested JSON structure (10 levels to avoid stack overflow)
|
||||
let mut nested_json = String::from(r#"{"level": 0}"#);
|
||||
for i in 1..=10 {
|
||||
nested_json = format!(r#"{{"level": {}, "child": {}}}"#, i, nested_json);
|
||||
}
|
||||
|
||||
let input = format!(
|
||||
r#"<tool_call>
|
||||
<function=process_nested>
|
||||
<parameter=config>{}</parameter>
|
||||
</function>
|
||||
</tool_call>"#,
|
||||
nested_json
|
||||
);
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(&input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "process_nested");
|
||||
|
||||
// Verify the nested JSON was parsed correctly
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["config"].is_object());
|
||||
assert_eq!(args["config"]["level"], 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_streaming_malformed_recovery() {
|
||||
let mut parser = QwenCoderParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
// First: malformed tool call (invalid function name)
|
||||
// Second: valid tool call
|
||||
let chunks = vec![
|
||||
r#"<tool_call><function=invalid_func><parameter=x>1</parameter></function></tool_call>"#,
|
||||
r#"<tool_call><function=get_weather><parameter=city>Tokyo</parameter></function></tool_call>"#,
|
||||
];
|
||||
|
||||
let mut valid_tool_found = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
if name == "get_weather" {
|
||||
valid_tool_found = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
valid_tool_found,
|
||||
"Should recover and parse valid tool after invalid one"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_parameter_with_xml_like_content() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Parameter value contains XML-like content that shouldn't be parsed as tags
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=html_content><div class="test"><span>Hello</span></div></parameter>
|
||||
<parameter=xml_snippet><root><child attr="value"/></root></parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["html_content"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("<div class=\"test\">"));
|
||||
assert!(args["xml_snippet"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("<root><child"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_empty_parameter_value() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=empty></parameter>
|
||||
<parameter=whitespace> </parameter>
|
||||
<parameter=normal>value</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["empty"], "");
|
||||
assert_eq!(args["whitespace"], ""); // Trimmed
|
||||
assert_eq!(args["normal"], "value");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HTML Entity and Python Literal Tests
|
||||
// ============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_html_entity_decoding() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Test HTML entities in parameter values
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=ampersand>Tom & Jerry</parameter>
|
||||
<parameter=comparison>5 < 10 && 10 > 5</parameter>
|
||||
<parameter=quotes>"Hello" & 'World'</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["ampersand"], "Tom & Jerry");
|
||||
assert_eq!(args["comparison"], "5 < 10 && 10 > 5");
|
||||
assert_eq!(args["quotes"], "\"Hello\" & 'World'");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_html_numeric_entities() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Test numeric HTML entities
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=decimal><tag></parameter>
|
||||
<parameter=hex><tag></parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["decimal"], "<tag>");
|
||||
assert_eq!(args["hex"], "<tag>");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_python_literals() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Test Python-style literals (True, False, None)
|
||||
let input = r#"<tool_call>
|
||||
<function=process>
|
||||
<parameter=py_true>True</parameter>
|
||||
<parameter=py_false>False</parameter>
|
||||
<parameter=py_none>None</parameter>
|
||||
<parameter=json_true>true</parameter>
|
||||
<parameter=json_false>false</parameter>
|
||||
<parameter=json_null>null</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
// Python literals should be converted
|
||||
assert_eq!(args["py_true"], true);
|
||||
assert_eq!(args["py_false"], false);
|
||||
assert_eq!(args["py_none"], serde_json::Value::Null);
|
||||
// JSON literals should also work
|
||||
assert_eq!(args["json_true"], true);
|
||||
assert_eq!(args["json_false"], false);
|
||||
assert_eq!(args["json_null"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_coder_mixed_html_and_json() {
|
||||
let parser = QwenCoderParser::new();
|
||||
|
||||
// Test HTML entities within JSON structures
|
||||
let input = r#"<tool_call>
|
||||
<function=search>
|
||||
<parameter=query>price < 100 && rating > 4</parameter>
|
||||
<parameter=config>{"operator": "&&", "escape": true}</parameter>
|
||||
</function>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["query"], "price < 100 && rating > 4");
|
||||
// JSON with HTML entity inside - the entity gets decoded first, then JSON parsed
|
||||
assert!(args["config"].is_object());
|
||||
assert_eq!(args["config"]["operator"], "&&");
|
||||
}
|
||||
@@ -1,237 +0,0 @@
|
||||
//! Step3 Parser Integration Tests
|
||||
|
||||
use smg::tool_parser::{Step3Parser, ToolParser};
|
||||
|
||||
use crate::common::create_test_tools;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_complete_parsing() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"Let me help you.
|
||||
<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="search">
|
||||
<steptml:parameter name="query">rust programming</steptml:parameter>
|
||||
<steptml:parameter name="limit">10</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>
|
||||
Here are the results..."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Let me help you.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["query"], "rust programming");
|
||||
assert_eq!(args["limit"], 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_multiple_tools() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="get_weather">
|
||||
<steptml:parameter name="location">Tokyo</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="get_news">
|
||||
<steptml:parameter name="category">tech</steptml:parameter>
|
||||
<steptml:parameter name="limit">5</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert_eq!(tools[1].function.name, "get_news");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_type_conversion() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="process">
|
||||
<steptml:parameter name="count">100</steptml:parameter>
|
||||
<steptml:parameter name="rate">2.5</steptml:parameter>
|
||||
<steptml:parameter name="active">true</steptml:parameter>
|
||||
<steptml:parameter name="optional">null</steptml:parameter>
|
||||
<steptml:parameter name="text">hello world</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["count"], 100);
|
||||
assert_eq!(args["rate"], 2.5);
|
||||
assert_eq!(args["active"], true);
|
||||
assert_eq!(args["optional"], serde_json::Value::Null);
|
||||
assert_eq!(args["text"], "hello world");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_streaming() {
|
||||
let mut parser = Step3Parser::new();
|
||||
|
||||
let tools = create_test_tools();
|
||||
|
||||
// Simulate streaming chunks
|
||||
let chunks = vec![
|
||||
"<｜tool_calls_begin｜>\n",
|
||||
"<｜tool_call_begin｜>function",
|
||||
"<｜tool_sep｜><steptml:invoke name=\"calc\">",
|
||||
"\n<steptml:parameter name=\"x\">10</steptml:parameter>",
|
||||
"\n<steptml:parameter name=\"y\">20</steptml:parameter>",
|
||||
"\n</steptml:invoke><｜tool_call_end｜>",
|
||||
"\n<｜tool_calls_end｜>",
|
||||
];
|
||||
|
||||
let mut found_complete = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
if !result.calls.is_empty() {
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "calc");
|
||||
found_complete = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(found_complete);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_step3_format_detection() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
// Should detect Step3 format
|
||||
assert!(parser.has_tool_markers("<｜tool_calls_begin｜>"));
|
||||
assert!(parser.has_tool_markers("text with <｜tool_calls_begin｜> marker"));
|
||||
|
||||
// Should not detect other formats
|
||||
assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
|
||||
assert!(!parser.has_tool_markers("<tool_call>"));
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_nested_steptml() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="config">
|
||||
<steptml:parameter name="settings">{"nested": {"key": "value"}}</steptml:parameter>
|
||||
<steptml:parameter name="array">[1, 2, 3]</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "config");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["settings"].is_object());
|
||||
assert!(args["array"].is_array());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_python_literals() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="test">
|
||||
<steptml:parameter name="bool_true">True</steptml:parameter>
|
||||
<steptml:parameter name="bool_false">False</steptml:parameter>
|
||||
<steptml:parameter name="none_value">None</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["bool_true"], true);
|
||||
assert_eq!(args["bool_false"], false);
|
||||
assert_eq!(args["none_value"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_steptml_format() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"Text before.
|
||||
<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="search">
|
||||
<steptml:parameter name="query">rust lang</steptml:parameter>
|
||||
<steptml:parameter name="limit">10</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>Text after."#;
|
||||
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "Text before.\n");
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["query"], "rust lang");
|
||||
assert_eq!(args["limit"], 10);
|
||||
// TODO: Verify normal text extraction
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_parameter_values() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="config">
|
||||
<steptml:parameter name="settings">{"nested": {"value": true}}</steptml:parameter>
|
||||
<steptml:parameter name="items">[1, 2, 3]</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert!(args["settings"].is_object());
|
||||
assert!(args["items"].is_array());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_parameter_with_angle_brackets() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="compare">
|
||||
<steptml:parameter name="expression">a < b && b > c</steptml:parameter>
|
||||
<steptml:parameter name="context">comparison test</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].function.name, "compare");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["expression"], "a < b && b > c");
|
||||
assert_eq!(args["context"], "comparison test");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_step3_empty_function_name() {
|
||||
let parser = Step3Parser::new();
|
||||
|
||||
let input = r#"<｜tool_calls_begin｜>
|
||||
<｜tool_call_begin｜>function<｜tool_sep｜><steptml:invoke name="">
|
||||
<steptml:parameter name="param">value</steptml:parameter>
|
||||
</steptml:invoke><｜tool_call_end｜>
|
||||
<｜tool_calls_end｜>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0); // Should reject empty function name
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
//! Tool parser integration tests
|
||||
|
||||
#[path = "common/mod.rs"]
|
||||
pub mod common;
|
||||
|
||||
mod tool_parser;
|
||||
|
||||
pub use tool_parser::*;
|
||||
Reference in New Issue
Block a user