Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
201 changes: 153 additions & 48 deletions crates/tool_parser/src/parsers/glm4_moe.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::collections::HashMap;

use async_trait::async_trait;
use openai_protocol::common::Tool;
use regex::Regex;
Expand Down Expand Up @@ -89,37 +91,23 @@ impl Glm4MoeParser {
Self::new(r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>")
}

/// Parse arguments from key-value pairs
fn parse_arguments(&self, args_text: &str) -> serde_json::Map<String, Value> {
/// Parse arguments from key-value pairs, coerced by the function's declared
/// schema. A `string` parameter keeps a numeric/bool/array-looking value as a
/// string (matches vLLM); unknown types fall back to [`infer_value`].
fn parse_arguments(
&self,
args_text: &str,
param_types: &HashMap<String, String>,
) -> 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())
}
};
let value =
helpers::coerce_by_schema_type(value_str, param_types.get(key).map(String::as_str))
.unwrap_or_else(|| infer_value(value_str));
Comment thread
key4ng marked this conversation as resolved.

arguments.insert(key.to_string(), value);
}
Expand All @@ -128,16 +116,17 @@ impl Glm4MoeParser {
}

/// Parse a single tool call block
fn parse_tool_call(&self, block: &str) -> ParserResult<Option<ToolCall>> {
fn parse_tool_call(&self, block: &str, tools: &[Tool]) -> 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);
// Parse arguments, coerced by this function's declared schema.
let param_types = helpers::param_types_for_function(tools, func_name);
let arguments = self.parse_arguments(args_text, &param_types);

let arguments_str = serde_json::to_string(&arguments)
.map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
Expand All @@ -154,12 +143,12 @@ impl Glm4MoeParser {
}

/// Parse all tool calls from text (shared logic for complete and incremental parsing)
fn parse_tool_calls_from_text(&self, text: &str) -> Vec<ToolCall> {
let mut tools = Vec::new();
fn parse_tool_calls_from_text(&self, text: &str, tools: &[Tool]) -> Vec<ToolCall> {
let mut parsed = 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),
match self.parse_tool_call(mat.as_str(), tools) {
Ok(Some(tool)) => parsed.push(tool),
Ok(None) => continue,
Err(e) => {
tracing::debug!("Failed to parse tool call: {}", e);
Expand All @@ -168,20 +157,17 @@ impl Glm4MoeParser {
}
}

tools
}
}

impl Default for Glm4MoeParser {
fn default() -> Self {
Self::glm45()
parsed
}
}

#[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
impl Glm4MoeParser {
/// Shared non-streaming parse, schema-aware when `tools` are provided.
fn parse_complete_inner(
&self,
text: &str,
tools: &[Tool],
) -> ParserResult<(String, Vec<ToolCall>)> {
if !self.has_tool_markers(text) {
return Ok((text.to_string(), vec![]));
}
Expand All @@ -193,15 +179,58 @@ impl ToolParser for Glm4MoeParser {
.ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
let normal_text = text[..idx].to_string();

// Parse all tool calls using shared helper
let tools = self.parse_tool_calls_from_text(text);
let parsed = self.parse_tool_calls_from_text(text, tools);

// If no tools were successfully parsed despite having markers, return entire text as fallback
if tools.is_empty() {
if parsed.is_empty() {
return Ok((text.to_string(), vec![]));
}

Ok((normal_text, tools))
Ok((normal_text, parsed))
}
}

/// Infer a JSON value from raw text when the schema type is unknown: JSON
/// (numbers/bools/null/objects/arrays), then Python-style literals, then string.
fn infer_value(value_str: &str) -> Value {
if let Ok(json_val) = serde_json::from_str::<Value>(value_str) {
return json_val;
}
match value_str {
"true" | "True" => Value::Bool(true),
"false" | "False" => Value::Bool(false),
"null" | "None" => Value::Null,
_ => {
if let Ok(num) = value_str.parse::<i64>() {
Value::Number(num.into())
} else if let Ok(num) = value_str.parse::<f64>() {
serde_json::Number::from_f64(num)
.map_or_else(|| Value::String(value_str.to_string()), Value::Number)
} else {
Value::String(value_str.to_string())
}
}
}
}

impl Default for Glm4MoeParser {
fn default() -> Self {
Self::glm45()
}
}

#[async_trait]
impl ToolParser for Glm4MoeParser {
async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
self.parse_complete_inner(text, &[])
}

async fn parse_complete_with_tools(
&self,
text: &str,
tools: &[Tool],
) -> ParserResult<(String, Vec<ToolCall>)> {
self.parse_complete_inner(text, tools)
}

async fn parse_incremental(
Expand Down Expand Up @@ -250,7 +279,7 @@ impl ToolParser for Glm4MoeParser {

// 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]);
let parsed_tools = self.parse_tool_calls_from_text(&current_text[..block_end], tools);
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// Extract normal text before tool calls
let idx = current_text.find(self.bot_token);
Expand Down Expand Up @@ -346,3 +375,79 @@ impl ToolParser for Glm4MoeParser {
self.streamed_args_for_tool.clear();
}
}

#[cfg(test)]
mod tests {
use openai_protocol::common::Function;

use super::*;

fn tool_with_props(props: Value) -> Vec<Tool> {
vec![Tool {
tool_type: "function".to_string(),
function: Function {
name: "f".to_string(),
description: None,
parameters: serde_json::json!({"type": "object", "properties": props}),
strict: None,
},
}]
}

// String-typed params keep numeric/bool/array-looking values as strings (the
// JavaScript-category bug); non-string params still coerce.
#[tokio::test]
async fn test_schema_aware_coercion_keeps_strings() {
let tools = tool_with_props(serde_json::json!({
"limit": {"type": "string"},
"flag": {"type": "string"},
"coords": {"type": "string"},
"count": {"type": "integer"},
}));
let text = "<tool_call>f\n\
<arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
<arg_key>flag</arg_key>\n<arg_value>true</arg_value>\n\
<arg_key>coords</arg_key>\n<arg_value>[60,30]</arg_value>\n\
<arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
</tool_call>";
let (_, calls) = Glm4MoeParser::glm45()
.parse_complete_with_tools(text, &tools)
.await
.unwrap();
assert_eq!(calls.len(), 1);
let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
assert_eq!(args["limit"], Value::String("4".to_string()));
assert_eq!(args["flag"], Value::String("true".to_string()));
assert_eq!(args["coords"], Value::String("[60,30]".to_string()));
assert_eq!(args["count"], Value::Number(5.into()));
}

// Without a schema (no tools), behavior is unchanged: blind inference.
#[tokio::test]
async fn test_no_schema_infers_type() {
let text = "<tool_call>f\n<arg_key>count</arg_key>\n<arg_value>5</arg_value>\n</tool_call>";
let (_, calls) = Glm4MoeParser::glm45().parse_complete(text).await.unwrap();
let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
assert_eq!(args["count"], Value::Number(5.into()));
}

// The incremental path forwards tools, so it coerces by schema too.
#[tokio::test]
async fn test_streaming_schema_aware_coercion() {
let tools = tool_with_props(serde_json::json!({
"limit": {"type": "string"},
"count": {"type": "integer"},
}));
let text = "<tool_call>f\n\
<arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
<arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
</tool_call>";
let result = Glm4MoeParser::glm45()
.parse_incremental(text, &tools)
.await
.unwrap();
let args: Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
assert_eq!(args["limit"], Value::String("4".to_string()));
assert_eq!(args["count"], Value::Number(5.into()));
}
}
47 changes: 43 additions & 4 deletions crates/tool_parser/src/parsers/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,16 @@ pub fn param_types_for_function(tools: &[Tool], func_name: &str) -> HashMap<Stri
types
}

/// Coerce a raw value by its declared JSON-schema type. `string` is kept verbatim;
/// numeric/boolean/structured types are parsed. `None` (unknown type or parse
/// failure) means the caller should fall back to its own inference.
/// Coerce a raw value by its declared JSON-schema type. For `string`, a JSON
/// string literal (`"4"`) is unwrapped to its content while bare text (`4`,
/// `true`, `[60,30]`) is kept as the string itself; numeric/boolean/structured
/// types are parsed. `None` (unknown type or parse failure) means the caller
/// should fall back to its own inference.
pub fn coerce_by_schema_type(text: &str, declared_type: Option<&str>) -> Option<Value> {
match declared_type? {
"string" => Some(Value::String(text.to_string())),
"string" => Some(Value::String(
serde_json::from_str::<String>(text).unwrap_or_else(|_| text.to_string()),
)),
"integer" => text
.trim()
.parse::<i64>()
Expand Down Expand Up @@ -622,4 +626,39 @@ mod tests {
&serde_json::json!({"key": "value"})
);
}

#[test]
fn test_coerce_by_schema_type() {
use serde_json::json;
// string: bare numeric/bool/array-looking text stays a string...
assert_eq!(coerce_by_schema_type("4", Some("string")), Some(json!("4")));
assert_eq!(
coerce_by_schema_type("true", Some("string")),
Some(json!("true"))
);
assert_eq!(
coerce_by_schema_type("[60,30]", Some("string")),
Some(json!("[60,30]"))
);
// ...but a JSON string literal is unwrapped to its content.
assert_eq!(
coerce_by_schema_type("\"4\"", Some("string")),
Some(json!("4"))
);

// typed values are parsed
assert_eq!(coerce_by_schema_type("4", Some("integer")), Some(json!(4)));
assert_eq!(
coerce_by_schema_type("true", Some("boolean")),
Some(json!(true))
);
assert_eq!(
coerce_by_schema_type("[60,30]", Some("array")),
Some(json!([60, 30]))
);

// unknown type / value that fails to parse -> None (caller infers)
assert_eq!(coerce_by_schema_type("4", None), None);
assert_eq!(coerce_by_schema_type("abc", Some("integer")), None);
}
}
Loading
Loading