[model-gateway] change sgl-router to sgl-model-gateway (#14312)

This commit is contained in:
Simo Lin
2025-12-05 12:04:48 -08:00
committed by GitHub
parent 1ea6b740a7
commit 49dfa1d891
431 changed files with 86 additions and 93 deletions
@@ -0,0 +1,32 @@
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),
}
@@ -0,0 +1,384 @@
// 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, 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("pythonic", || Box::new(PythonicParser::new()));
registry.register_parser("llama", || Box::new(LlamaParser::new()));
registry.register_parser("deepseek", || Box::new(DeepSeekParser::new()));
registry.register_parser("glm4_moe", || Box::new(Glm4MoeParser::new()));
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(&registry);
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
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*", "glm4_moe");
registry.map_model("glm-4.6*", "glm4_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()
}
}
+27
View File
@@ -0,0 +1,27 @@
/// 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 state;
pub mod traits;
pub mod types;
// Parser implementations
pub mod parsers;
#[cfg(test)]
mod tests;
// Re-export commonly used types
pub use errors::{ParserError, ParserResult};
pub use factory::{ParserFactory, ParserRegistry, PooledParser};
// Re-export parsers for convenience
pub use parsers::{
DeepSeekParser, Glm4MoeParser, JsonParser, KimiK2Parser, LlamaParser, MinimaxM2Parser,
MistralParser, PythonicParser, QwenParser, Step3Parser,
};
pub use traits::{PartialJsonParser, ToolParser};
pub use types::{FunctionCall, PartialToolCall, StreamingParseResult, ToolCall};
@@ -0,0 +1,328 @@
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::warn!("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::warn!("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();
}
}
@@ -0,0 +1,326 @@
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 the GLM-4 MoE specific format:
/// `<tool_call>{name}\n<arg_key>{key}</arg_key>\n<arg_value>{value}</arg_value>\n</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 GLM-4 MoE parser
pub fn new() -> 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_pattern = r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>";
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>",
}
}
/// 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 and return StreamingParseResult (mirrors Python's detect_and_parse)
/// 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::warn!("Failed to parse tool call: {}", e);
continue;
}
}
}
Ok(tools)
}
}
impl Default for Glm4MoeParser {
fn default() -> Self {
Self::new()
}
}
#[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(&current_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::warn!("Invalid tool name '{}' - skipping", tool_call.function.name);
helpers::reset_current_tool_state(
&mut self.buffer,
&mut false, // glm4_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();
}
}
@@ -0,0 +1,488 @@
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 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 = &current_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::warn!("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);
}
}
@@ -0,0 +1,305 @@
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::warn!("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;
}
}
@@ -0,0 +1,348 @@
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::warn!(
"Failed to parse JSON arguments for {}: {}",
func_name,
e
);
continue;
}
}
} else {
tracing::warn!("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::warn!("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();
}
}
@@ -0,0 +1,245 @@
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::warn!("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::warn!("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,
);
}
}
@@ -0,0 +1,549 @@
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("&lt;", "<")
.replace("&gt;", ">")
.replace("&amp;", "&")
.replace("&quot;", "\"")
.replace("&apos;", "'")
}
/// 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(&parameters)
.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::warn!("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::warn!("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;
}
}
@@ -0,0 +1,287 @@
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::warn!("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;
}
}
@@ -0,0 +1,32 @@
/// 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 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 use passthrough::PassthroughParser;
pub use pythonic::PythonicParser;
pub use qwen::QwenParser;
pub use step3::Step3Parser;
@@ -0,0 +1,55 @@
//! 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 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
}
}
@@ -0,0 +1,414 @@
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 warning and return entire text as fallback
tracing::warn!("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::warn!(
"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::warn!("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())
}
}
@@ -0,0 +1,258 @@
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::warn!("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,
);
}
}
@@ -0,0 +1,576 @@
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::warn!("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(&parameters)
.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::warn!("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;
}
}
@@ -0,0 +1,553 @@
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()))
}
}
}
/// Utility function to check if a string contains complete JSON
pub fn is_complete_json(input: &str) -> bool {
serde_json::from_str::<Value>(input).is_ok()
}
/// Utility function to find common prefix between two strings
pub fn find_common_prefix(s1: &str, s2: &str) -> usize {
s1.chars()
.zip(s2.chars())
.take_while(|(a, b)| a == b)
.count()
}
/// Utility function to compute diff between old and new strings
pub fn compute_diff(old: &str, new: &str) -> String {
let common_len = find_common_prefix(old, new);
// Convert character count to byte offset
new.chars().skip(common_len).collect()
}
@@ -0,0 +1,16 @@
/// Placeholder for Harmony streaming metadata captured during token-aware parsing.
#[derive(Debug, Clone, Default)]
pub struct HarmonyStreamState {
/// All tokens observed so far for the current assistant response.
pub tokens: Vec<u32>,
/// Number of tokens that have already been processed by the Harmony parser.
pub processed_tokens: usize,
/// Number of tool calls emitted downstream.
pub emitted_calls: usize,
/// Pending analysis-channel content awaiting flush into normal text output.
pub analysis_buffer: String,
/// Whether the tool name has been surfaced for the current call.
pub emitted_name: bool,
/// Whether arguments have been surfaced for the current call.
pub emitted_args: bool,
}
+599
View File
@@ -0,0 +1,599 @@
use super::*;
use crate::tool_parser::{
parsers::JsonParser,
partial_json::{compute_diff, find_common_prefix, is_complete_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());
}
#[test]
fn test_is_complete_json() {
assert!(is_complete_json(r#"{"name": "test"}"#));
assert!(is_complete_json(r#"[1, 2, 3]"#));
assert!(is_complete_json(r#""string""#));
assert!(is_complete_json("42"));
assert!(is_complete_json("true"));
assert!(is_complete_json("null"));
assert!(!is_complete_json(r#"{"name": "#));
assert!(!is_complete_json(r#"[1, 2, "#));
assert!(!is_complete_json(r#""unclosed"#));
}
#[test]
fn test_find_common_prefix() {
assert_eq!(find_common_prefix("hello", "hello"), 5);
assert_eq!(find_common_prefix("hello", "help"), 3);
assert_eq!(find_common_prefix("hello", "world"), 0);
assert_eq!(find_common_prefix("", "hello"), 0);
assert_eq!(find_common_prefix("hello", ""), 0);
}
#[test]
fn test_compute_diff() {
assert_eq!(compute_diff("hello", "hello world"), " world");
assert_eq!(compute_diff("", "hello"), "hello");
assert_eq!(compute_diff("hello", "hello"), "");
assert_eq!(compute_diff("test", "hello"), "hello");
}
// 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();
}
}
}
@@ -0,0 +1,76 @@
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>;
}
@@ -0,0 +1,88 @@
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,
}
/// Streaming parse result
#[derive(Debug, Clone)]
pub enum StreamResult {
/// Need more data to continue parsing
Incomplete,
/// Found a tool name (for streaming)
ToolName { index: usize, name: String },
/// Found incremental arguments (for streaming)
ToolArguments { index: usize, arguments: String },
/// Completed parsing a tool
ToolComplete(ToolCall),
/// Normal text (not part of tool call)
NormalText(String),
}
/// Token configuration for parsing
#[derive(Debug, Clone)]
pub struct TokenConfig {
/// Start tokens for tool calls
pub start_tokens: Vec<String>,
/// End tokens for tool calls
pub end_tokens: Vec<String>,
/// Separator between multiple tool calls
pub separator: String,
}
impl TokenConfig {
/// Iterate over start/end token pairs
pub fn iter_pairs(&self) -> impl Iterator<Item = (&str, &str)> {
self.start_tokens
.iter()
.zip(self.end_tokens.iter())
.map(|(s, e)| (s.as_str(), e.as_str()))
}
}
/// 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,
}