[router][grpc] Fix streaming bugs: empty tool names, state pollution, and panics (#11373)
This commit is contained in:
Chang Su
@@ -4,6 +4,7 @@
|
||||
pub mod mock_mcp_server;
|
||||
pub mod mock_openai_server;
|
||||
pub mod mock_worker;
|
||||
pub mod streaming_helpers;
|
||||
pub mod test_app;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
134
sgl-router/tests/common/streaming_helpers.rs
Normal file
134
sgl-router/tests/common/streaming_helpers.rs
Normal file
@@ -0,0 +1,134 @@
|
||||
//! Streaming Test Helpers
|
||||
//!
|
||||
//! Utilities for creating realistic streaming chunks that simulate
|
||||
//! how LLM tokens actually arrive (1-5 characters at a time).
|
||||
|
||||
/// Split input into realistic char-level chunks (2-3 chars each for determinism)
|
||||
pub fn create_realistic_chunks(input: &str) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
let mut i = 0;
|
||||
|
||||
while i < chars.len() {
|
||||
// Take 2-3 characters at a time (deterministic for testing)
|
||||
let chunk_size = if i + 3 <= chars.len() && chars[i].is_ascii_alphanumeric() {
|
||||
3 // Longer chunks for alphanumeric sequences
|
||||
} else {
|
||||
2 // Shorter chunks for special characters
|
||||
};
|
||||
|
||||
let end = (i + chunk_size).min(chars.len());
|
||||
let chunk: String = chars[i..end].iter().collect();
|
||||
chunks.push(chunk);
|
||||
i = end;
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Split input at strategic positions to test edge cases
|
||||
/// This creates chunks that break at critical positions like after quotes, colons, etc.
|
||||
pub fn create_strategic_chunks(input: &str) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut current = String::new();
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
|
||||
for (i, &ch) in chars.iter().enumerate() {
|
||||
current.push(ch);
|
||||
|
||||
// Break after strategic characters
|
||||
let should_break = matches!(ch, '"' | ':' | ',' | '{' | '}' | '[' | ']')
|
||||
|| (i > 0 && chars[i-1] == '"' && ch == ' ') // Space after quote
|
||||
|| current.len() >= 5; // Max 5 chars per chunk
|
||||
|
||||
if should_break && !current.is_empty() {
|
||||
chunks.push(current.clone());
|
||||
current.clear();
|
||||
}
|
||||
}
|
||||
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Create the bug scenario chunks: `{"name": "` arrives in parts
|
||||
pub fn create_bug_scenario_chunks() -> Vec<&'static str> {
|
||||
vec![
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"name"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#, // Bug occurs here: parser has {"name": "
|
||||
r#"search"#, // Use valid tool name
|
||||
r#"""#,
|
||||
r#","#,
|
||||
r#" "#,
|
||||
r#"""#,
|
||||
r#"arguments"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"query"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#,
|
||||
r#"test query"#,
|
||||
r#"""#,
|
||||
r#"}"#,
|
||||
r#"}"#,
|
||||
]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[allow(unused_imports)]
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_realistic_chunks() {
|
||||
let input = r#"{"name": "test"}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
// Should have multiple chunks
|
||||
assert!(chunks.len() > 3);
|
||||
|
||||
// Reconstructed should equal original
|
||||
let reconstructed: String = chunks.join("");
|
||||
assert_eq!(reconstructed, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strategic_chunks_breaks_after_quotes() {
|
||||
let input = r#"{"name": "value"}"#;
|
||||
let chunks = create_strategic_chunks(input);
|
||||
|
||||
// Should break after quotes and colons
|
||||
assert!(chunks.iter().any(|c| c.ends_with('"')));
|
||||
assert!(chunks.iter().any(|c| c.ends_with(':')));
|
||||
|
||||
// Reconstructed should equal original
|
||||
let reconstructed: String = chunks.join("");
|
||||
assert_eq!(reconstructed, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug_scenario_chunks() {
|
||||
let chunks = create_bug_scenario_chunks();
|
||||
let reconstructed: String = chunks.join("");
|
||||
|
||||
// Should reconstruct to valid JSON
|
||||
assert!(reconstructed.contains(r#"{"name": "search""#));
|
||||
|
||||
// The critical chunk sequence should be present (space after colon, then quote in next chunk)
|
||||
let joined = chunks.join("|");
|
||||
assert!(joined.contains(r#" |"#)); // The bug happens at {"name": " and then "
|
||||
}
|
||||
}
|
||||
@@ -126,28 +126,6 @@ fn test_glm4_format_detection() {
|
||||
assert!(!parser.has_tool_markers("plain text"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glm4_python_literal_values() {
|
||||
let parser = Glm4MoeParser::new();
|
||||
|
||||
let input = r#"<tool_call>config
|
||||
<arg_key>debug</arg_key>
|
||||
<arg_value>True</arg_value>
|
||||
<arg_key>verbose</arg_key>
|
||||
<arg_value>False</arg_value>
|
||||
<arg_key>optional</arg_key>
|
||||
<arg_value>None</arg_value>
|
||||
</tool_call>"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["debug"], true);
|
||||
assert_eq!(args["verbose"], false);
|
||||
assert_eq!(args["optional"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_python_literals() {
|
||||
let parser = Glm4MoeParser::new();
|
||||
@@ -172,7 +150,7 @@ async fn test_python_literals() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_nested_values() {
|
||||
async fn test_glm4_nested_json_in_arg_values() {
|
||||
let parser = Glm4MoeParser::new();
|
||||
|
||||
let input = r#"<tool_call>process
|
||||
|
||||
156
sgl-router/tests/tool_parser_partial_json.rs
Normal file
156
sgl-router/tests/tool_parser_partial_json.rs
Normal file
@@ -0,0 +1,156 @@
|
||||
//! Partial JSON Parser Tests
|
||||
//!
|
||||
//! Tests for the partial JSON parser with allow_partial_strings flag behavior
|
||||
|
||||
use sglang_router_rs::tool_parser::partial_json::PartialJson;
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_disallows_incomplete_strings() {
|
||||
// Test case from the bug report: {"name": "
|
||||
// With allow_partial_strings=false, should return {} (stop before incomplete string)
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should parse just the opening brace and stop at the incomplete string
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
|
||||
// Should have empty object (stopped before parsing incomplete "name" key)
|
||||
assert!(
|
||||
obj_map.is_empty() || !obj_map.contains_key("name"),
|
||||
"Should not parse incomplete string key, got: {:?}",
|
||||
obj_map
|
||||
);
|
||||
|
||||
// Should consume characters up to the incomplete string
|
||||
assert!(consumed <= input.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_allows_incomplete_strings() {
|
||||
// Test case: {"name": "
|
||||
// With allow_partial_strings=true, should parse the incomplete string
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should parse the object with incomplete string value
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
|
||||
// With allow_partial_strings=true, should parse "name" key with empty string value
|
||||
assert!(
|
||||
obj_map.contains_key("name"),
|
||||
"Should parse incomplete string with allow_partial_strings=true"
|
||||
);
|
||||
|
||||
assert_eq!(consumed, input.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_flag_complete_json() {
|
||||
// Test case: {"name": "test"}
|
||||
// Both flags should parse complete JSON the same way
|
||||
let input = r#"{"name": "test"}"#;
|
||||
|
||||
let parser = PartialJson::new(32, true);
|
||||
let result1 = parser.parse_value(input, false);
|
||||
assert!(result1.is_ok());
|
||||
let (obj1, consumed1) = result1.unwrap();
|
||||
|
||||
let result2 = parser.parse_value(input, true);
|
||||
assert!(result2.is_ok());
|
||||
let (obj2, consumed2) = result2.unwrap();
|
||||
|
||||
// Both should parse the same complete JSON
|
||||
assert_eq!(obj1, obj2);
|
||||
assert_eq!(consumed1, consumed2);
|
||||
assert_eq!(consumed1, input.len());
|
||||
|
||||
// Check the parsed value
|
||||
assert!(obj1.is_object());
|
||||
let obj_map = obj1.as_object().unwrap();
|
||||
assert_eq!(obj_map.get("name").and_then(|v| v.as_str()), Some("test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backward_compatibility_default() {
|
||||
// Test that default PartialJson still allows partial strings (backward compatible)
|
||||
let parser = PartialJson::default();
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, true);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, _) = result.unwrap();
|
||||
assert!(obj.is_object());
|
||||
|
||||
// Default behavior should allow partial strings
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
assert!(
|
||||
obj_map.contains_key("name"),
|
||||
"Default should allow partial strings for backward compatibility"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_string_in_nested_object() {
|
||||
// Test case: {"tool": {"name": "
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"tool": {"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, _) = result.unwrap();
|
||||
assert!(obj.is_object());
|
||||
|
||||
// With allow_partial_strings=false, should stop before incomplete nested string
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
if let Some(tool) = obj_map.get("tool") {
|
||||
if let Some(tool_map) = tool.as_object() {
|
||||
assert!(
|
||||
!tool_map.contains_key("name")
|
||||
|| tool_map.get("name").and_then(|v| v.as_str()).is_none(),
|
||||
"Should not parse incomplete nested string"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug_fix_exact_scenario() {
|
||||
// This test verifies the exact bug scenario from the issue:
|
||||
// buffer = "{\"name\": \""
|
||||
// flags = Allow.ALL & ~Allow.STR
|
||||
// Python returns: Parsed object: {}, consumed length: 10
|
||||
|
||||
let parser = PartialJson::new(32, true);
|
||||
let input = r#"{"name": ""#;
|
||||
|
||||
let result = parser.parse_value(input, false);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (obj, consumed) = result.unwrap();
|
||||
|
||||
// Should return empty object (not {"name": null} or {"name": ""})
|
||||
assert!(obj.is_object());
|
||||
let obj_map = obj.as_object().unwrap();
|
||||
assert!(
|
||||
obj_map.is_empty(),
|
||||
"Expected empty object, got: {:?}. This matches Python behavior with Allow.ALL & ~Allow.STR",
|
||||
obj_map
|
||||
);
|
||||
|
||||
// Should consume all characters (10 bytes)
|
||||
assert_eq!(consumed, 10, "Should consume all 10 characters");
|
||||
}
|
||||
@@ -1,73 +1,199 @@
|
||||
//! Streaming Parser Tests
|
||||
//! Realistic Streaming Parser Tests
|
||||
//!
|
||||
//! Tests for incremental/streaming parsing capabilities across all parsers
|
||||
//! Tests incremental parsing with realistic char-level chunks (2-5 chars)
|
||||
//! that simulate how LLM tokens actually arrive.
|
||||
//!
|
||||
//! These tests are designed to catch bugs like `{"name": "` being parsed
|
||||
//! as an empty tool name.
|
||||
|
||||
use sglang_router_rs::tool_parser::{
|
||||
JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
|
||||
};
|
||||
use sglang_router_rs::tool_parser::{JsonParser, LlamaParser, QwenParser, ToolParser};
|
||||
|
||||
mod common;
|
||||
use common::create_test_tools;
|
||||
use common::{create_test_tools, streaming_helpers::*};
|
||||
|
||||
// =============================================================================
|
||||
// THE BUG SCENARIO - Most Critical Test
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_streaming_simple() {
|
||||
async fn test_json_bug_incomplete_tool_name_string() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let full_json = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
|
||||
|
||||
let result = parser.parse_incremental(full_json, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("get_weather".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_streaming_array() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
// This exact sequence triggered the bug:
|
||||
// Parser receives {"name": " and must NOT parse it as empty name
|
||||
let chunks = vec![
|
||||
r#"["#,
|
||||
r#"{"name": "tool1", "#,
|
||||
r#""arguments": {}}, "#,
|
||||
r#"{"name": "tool2", "#,
|
||||
r#""arguments": {"x": 1"#,
|
||||
r#"}}]"#,
|
||||
r#"{"#,
|
||||
r#"""#,
|
||||
r#"name"#,
|
||||
r#"""#,
|
||||
r#":"#,
|
||||
r#" "#,
|
||||
r#"""#, // ← Critical moment: parser has {"name": "
|
||||
// At this point, partial_json should NOT allow incomplete strings
|
||||
// when current_tool_name_sent=false
|
||||
r#"search"#, // Use valid tool name from create_test_tools()
|
||||
r#"""#,
|
||||
r#", "#,
|
||||
r#"""#,
|
||||
r#"arguments"#,
|
||||
r#"""#,
|
||||
r#": {"#,
|
||||
r#"""#,
|
||||
r#"query"#,
|
||||
r#"""#,
|
||||
r#": "#,
|
||||
r#"""#,
|
||||
r#"rust programming"#,
|
||||
r#"""#,
|
||||
r#"}}"#,
|
||||
];
|
||||
|
||||
let mut tool_count = 0;
|
||||
let mut got_tool_name = false;
|
||||
let mut saw_empty_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
for chunk in chunks.iter() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
tool_count += 1;
|
||||
if let Some(name) = &call.name {
|
||||
if name.is_empty() {
|
||||
saw_empty_name = true;
|
||||
}
|
||||
if name == "search" {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Current implementation may handle this differently
|
||||
assert!(tool_count <= 2, "Should parse at most 2 tools");
|
||||
assert!(
|
||||
!saw_empty_name,
|
||||
"Parser should NEVER return empty tool name"
|
||||
);
|
||||
assert!(got_tool_name, "Should have parsed tool name correctly");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// JSON PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_realistic_chunks_simple_tool() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "get_weather", "arguments": {"city": "Paris"}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 10, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mistral_streaming() {
|
||||
async fn test_json_strategic_chunks_with_quotes() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let mut parser = MistralParser::new();
|
||||
let input = r#"{"name": "search", "arguments": {"query": "rust programming"}}"#;
|
||||
let chunks = create_strategic_chunks(input);
|
||||
|
||||
// Strategic chunks break after quotes and colons
|
||||
assert!(chunks.iter().any(|c| c.ends_with('"')));
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_json_incremental_arguments_streaming() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let input = r#"{"name": "search", "arguments": {"query": "test", "limit": 10}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
let mut tool_name_sent = false;
|
||||
let mut got_arguments = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
tool_name_sent = true;
|
||||
}
|
||||
if tool_name_sent && !call.parameters.is_empty() {
|
||||
got_arguments = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(tool_name_sent, "Should have sent tool name");
|
||||
assert!(got_arguments, "Should have sent arguments");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// LLAMA PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_realistic_chunks_with_python_tag() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
let input = r#"<|python_tag|>{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 15, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "calculate");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_python_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = LlamaParser::new();
|
||||
|
||||
// Python tag itself arrives in small chunks
|
||||
let chunks = vec![
|
||||
r#"Here is the result: "#,
|
||||
r#"[TOOL_CALLS] ["#,
|
||||
r#"{"name": "#,
|
||||
r#""search", "#,
|
||||
r#""arguments": "#,
|
||||
r#"{"query": "#,
|
||||
r#""rust lang""#,
|
||||
r#"}}]"#,
|
||||
"<|p", "yth", "on_", "tag", "|>{", r#"""#, "na", r#"me""#, ": ", r#"""#, "sea", "rch",
|
||||
r#"""#, ", ", r#"""#, "par", "ame", "ter", "s", r#"""#, ": {", r#"""#, "q", r#"""#, ": ",
|
||||
r#"""#, "tes", "t", r#"""#, "}}",
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
@@ -82,40 +208,47 @@ async fn test_mistral_streaming() {
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have found tool name");
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// QWEN PARSER REALISTIC STREAMING
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_realistic_chunks_with_xml_tags() {
|
||||
let tools = create_test_tools();
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let input = "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>";
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(chunks.len() > 20, "Should have many small chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "get_weather");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pythonic_streaming() {
|
||||
async fn test_qwen_xml_tag_arrives_in_parts() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = PythonicParser::new();
|
||||
|
||||
let full_input = r#"[get_weather(city="London", units="celsius")]"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("get_weather".to_string()));
|
||||
let args: serde_json::Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
|
||||
assert_eq!(args["city"], "London");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_llama_streaming_with_python_tag() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = LlamaParser::new();
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let chunks = vec![
|
||||
r#"Let me help. "#,
|
||||
r#"<|python"#,
|
||||
r#"_tag|>"#,
|
||||
r#"{"name": "#,
|
||||
r#""calculate", "#,
|
||||
r#""arguments": "#,
|
||||
r#"{"x": 10}"#,
|
||||
r#"}"#,
|
||||
"<to", "ol_", "cal", "l>\n", "{", r#"""#, "na", "me", r#"""#, ": ", r#"""#, "tra", "nsl",
|
||||
"ate", r#"""#, ", ", r#"""#, "arg", "ume", "nts", r#"""#, ": {", r#"""#, "tex", "t",
|
||||
r#"""#, ": ", r#"""#, "hel", "lo", r#"""#, "}}\n", "</t", "ool", "_ca", "ll>",
|
||||
];
|
||||
|
||||
let mut got_tool_name = false;
|
||||
@@ -124,191 +257,66 @@ async fn test_llama_streaming_with_python_tag() {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if let Some(name) = call.name {
|
||||
assert_eq!(name, "calculate");
|
||||
assert_eq!(name, "translate");
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have found tool name");
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_qwen_streaming() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
// Note: Parser expects newline after both tags
|
||||
let full_input = "<tool_call>\n{\"name\": \"translate\", \"arguments\": {\"text\": \"hello\", \"to\": \"zh\"}}\n</tool_call>";
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("translate".to_string()));
|
||||
}
|
||||
// =============================================================================
|
||||
// EDGE CASES WITH REALISTIC CHUNKS
|
||||
// =============================================================================
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_incomplete_stays_incomplete() {
|
||||
async fn test_json_very_long_url_in_arguments() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let chunks = vec![r#"{"na"#, r#"me": "#];
|
||||
// Simulate long URL arriving in many chunks
|
||||
let long_url = "https://example.com/very/long/path/".to_string() + &"segment/".repeat(50);
|
||||
let input = format!(
|
||||
r#"{{"name": "search", "arguments": {{"query": "{}"}}}}"#,
|
||||
long_url
|
||||
);
|
||||
let chunks = create_realistic_chunks(&input);
|
||||
|
||||
assert!(chunks.len() > 100, "Long URL should create many chunks");
|
||||
|
||||
let mut got_tool_name = false;
|
||||
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(chunk, &tools).await.unwrap();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should return empty calls for partial JSON, got: {:?}",
|
||||
result
|
||||
);
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(got_tool_name, "Should have parsed tool name");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_buffer_accumulation() {
|
||||
async fn test_json_unicode_arrives_byte_by_byte() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let result1 = parser.parse_incremental(r#"{"na"#, &tools).await.unwrap();
|
||||
let input = r#"{"name": "search", "arguments": {"query": "Hello 世界 🌍"}}"#;
|
||||
let chunks = create_realistic_chunks(input);
|
||||
|
||||
assert!(result1.calls.is_empty(), "Should not parse incomplete JSON");
|
||||
let mut got_tool_name = false;
|
||||
|
||||
let result2 = parser
|
||||
.parse_incremental(r#"me": "test", "arguments": {}}"#, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
!result2.calls.is_empty(),
|
||||
"Should parse complete JSON after buffering"
|
||||
);
|
||||
assert_eq!(result2.calls[0].name, Some("test".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_multiple_tools_sequential() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = QwenParser::new();
|
||||
|
||||
let full_input = r#"<tool_call>
|
||||
{"name": "tool1", "arguments": {}}
|
||||
</tool_call>"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
assert_eq!(result.calls[0].name, Some("tool1".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_reset_after_error() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser1 = JsonParser::new();
|
||||
|
||||
let _ = parser1
|
||||
.parse_incremental(r#"{"name": invalid}"#, &tools)
|
||||
.await;
|
||||
|
||||
// Use a new parser instance for clean state
|
||||
let mut parser2 = JsonParser::new();
|
||||
let result = parser2
|
||||
.parse_incremental(r#"{"name": "test", "arguments": {}}"#, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should parse valid JSON");
|
||||
assert_eq!(result.calls[0].name, Some("test".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_with_unicode_chunks() {
|
||||
let tools = create_test_tools();
|
||||
|
||||
let mut parser = JsonParser::new();
|
||||
|
||||
let full_input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍"}}"#;
|
||||
|
||||
let result = parser.parse_incremental(full_input, &tools).await.unwrap();
|
||||
|
||||
assert!(!result.calls.is_empty(), "Should have parsed a tool call");
|
||||
|
||||
// Check if we got the tool name
|
||||
if let Some(name) = &result.calls[0].name {
|
||||
assert_eq!(name, "translate");
|
||||
for chunk in chunks {
|
||||
let result = parser.parse_incremental(&chunk, &tools).await.unwrap();
|
||||
for call in result.calls {
|
||||
if call.name.is_some() {
|
||||
got_tool_name = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// In streaming mode, need to make another call to get parameters
|
||||
let result2 = parser.parse_incremental("", &tools).await.unwrap();
|
||||
|
||||
// Parameters should be in either result.calls[1] or result2.calls[0]
|
||||
let params = if result.calls.len() > 1 {
|
||||
&result.calls[1].parameters
|
||||
} else if !result2.calls.is_empty() {
|
||||
&result2.calls[0].parameters
|
||||
} else {
|
||||
&result.calls[0].parameters
|
||||
};
|
||||
|
||||
if !params.is_empty() {
|
||||
let args: serde_json::Value = serde_json::from_str(params).unwrap();
|
||||
assert!(args["text"].as_str().unwrap().contains("世界"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_with_partial_chunks() {
|
||||
let mut parser = JsonParser::new();
|
||||
let tools = create_test_tools();
|
||||
|
||||
let partial = r#"{"#;
|
||||
let result = parser.parse_incremental(partial, &tools).await.unwrap();
|
||||
assert!(
|
||||
result.calls.is_empty(),
|
||||
"Should return empty calls for just opening brace"
|
||||
);
|
||||
|
||||
let mut parser2 = JsonParser::new();
|
||||
let complete = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#;
|
||||
let result = parser2.parse_incremental(complete, &tools).await.unwrap();
|
||||
|
||||
assert!(
|
||||
!result.calls.is_empty(),
|
||||
"Expected tool call for complete JSON"
|
||||
);
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "get_weather");
|
||||
|
||||
// In streaming mode, need to make another call to get parameters
|
||||
let result2 = parser2.parse_incremental("", &tools).await.unwrap();
|
||||
|
||||
// Parameters should be in either result.calls[1] or result2.calls[0]
|
||||
let params = if result.calls.len() > 1 {
|
||||
&result.calls[1].parameters
|
||||
} else if !result2.calls.is_empty() {
|
||||
&result2.calls[0].parameters
|
||||
} else {
|
||||
&result.calls[0].parameters
|
||||
};
|
||||
|
||||
if !params.is_empty() {
|
||||
let args: serde_json::Value = serde_json::from_str(params).unwrap();
|
||||
assert_eq!(args["location"], "SF");
|
||||
}
|
||||
|
||||
// The PartialJson parser can complete partial JSON by filling in missing values
|
||||
let mut parser3 = JsonParser::new();
|
||||
let partial_with_name = r#"{"name": "test", "argum"#;
|
||||
let result = parser3
|
||||
.parse_incremental(partial_with_name, &tools)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Parser behavior may vary - either complete with partial data or wait for more
|
||||
if !result.calls.is_empty() {
|
||||
assert_eq!(result.calls[0].name.as_ref().unwrap(), "test");
|
||||
}
|
||||
assert!(got_tool_name, "Should have parsed with unicode");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user