[router][tool call] Improve normal content extraction and error handling (non-stream) (#11050)
This commit is contained in:
@@ -3,15 +3,16 @@
|
||||
//! Tests for the JSON parser which handles OpenAI, Claude, and generic JSON formats
|
||||
|
||||
use serde_json::json;
|
||||
use sglang_router_rs::tool_parser::{JsonParser, ToolParser};
|
||||
use sglang_router_rs::tool_parser::{JsonParser, TokenConfig, ToolParser};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_simple_json_tool_call() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
@@ -21,13 +22,14 @@ async fn test_simple_json_tool_call() {
|
||||
#[tokio::test]
|
||||
async fn test_json_array_of_tools() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"[
|
||||
let input = r#"Hello, here are the results: [
|
||||
{"name": "get_weather", "arguments": {"location": "SF"}},
|
||||
{"name": "search", "arguments": {"query": "news"}}
|
||||
]"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 2);
|
||||
assert_eq!(normal_text, "Hello, here are the results: ");
|
||||
assert_eq!(tools[0].function.name, "get_weather");
|
||||
assert_eq!(tools[1].function.name, "search");
|
||||
}
|
||||
@@ -37,8 +39,9 @@ async fn test_json_with_parameters_key() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "calculate");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
@@ -51,8 +54,12 @@ async fn test_json_extraction_from_text() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"I'll help you with that. {"name": "search", "arguments": {"query": "rust"}} Let me search for that."#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"I'll help you with that. Let me search for that."
|
||||
);
|
||||
assert_eq!(tools[0].function.name, "search");
|
||||
}
|
||||
|
||||
@@ -73,8 +80,9 @@ async fn test_json_with_nested_objects() {
|
||||
}
|
||||
}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "update_config");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
@@ -87,8 +95,9 @@ async fn test_json_with_special_characters() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "echo", "arguments": {"text": "Line 1\nLine 2\tTabbed", "path": "C:\\Users\\test"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Line 1\nLine 2\tTabbed");
|
||||
@@ -100,8 +109,9 @@ async fn test_json_with_unicode() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "translate", "arguments": {"text": "Hello 世界 🌍", "emoji": "😊"}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
assert_eq!(args["text"], "Hello 世界 🌍");
|
||||
@@ -113,8 +123,9 @@ async fn test_json_empty_arguments() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "ping", "arguments": {}}"#;
|
||||
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(normal_text, "");
|
||||
assert_eq!(tools[0].function.name, "ping");
|
||||
|
||||
let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
|
||||
@@ -127,8 +138,12 @@ async fn test_json_invalid_format() {
|
||||
|
||||
// Missing closing brace
|
||||
let input = r#"{"name": "test", "arguments": {"key": "value""#;
|
||||
let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tools.len(), 0);
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"{\"name\": \"test\", \"arguments\": {\"key\": \"value\""
|
||||
);
|
||||
|
||||
// Not JSON at all
|
||||
let input = "This is just plain text";
|
||||
@@ -145,3 +160,32 @@ async fn test_json_format_detection() {
|
||||
assert!(!parser.detect_format("plain text"));
|
||||
assert!(!parser.detect_format(r#"{"key": "value"}"#)); // No name field
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_with_wrapper_tokens() {
|
||||
let parser = JsonParser::with_config(TokenConfig {
|
||||
start_tokens: vec!["<tool>".to_string()],
|
||||
end_tokens: vec!["</tool>".to_string()],
|
||||
separator: ", ".to_string(),
|
||||
});
|
||||
|
||||
let input = r#"<tool>{"name": "test", "arguments": {}}</tool>"#;
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 1);
|
||||
assert_eq!(tool_calls[0].function.name, "test");
|
||||
assert_eq!(normal_text, ""); // Wrapper tokens with no extra text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_with_start_token_invalid_json() {
|
||||
let parser = JsonParser::with_config(TokenConfig {
|
||||
start_tokens: vec!["<|python_tag|>".to_string()],
|
||||
end_tokens: vec!["".to_string()],
|
||||
separator: ";".to_string(),
|
||||
});
|
||||
|
||||
let input = r#"Hello world <|python_tag|>this is not valid json at all"#;
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 0);
|
||||
assert_eq!(normal_text, input); // Should return entire original text when JSON parsing fails
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user