[router][tool call] Improve normal content extraction and error handling (non-stream) (#11050)
This commit is contained in:
@@ -615,155 +615,3 @@ impl ToolParser for JsonParser {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_single_tool_call() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
|
||||
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 1);
|
||||
assert_eq!(tool_calls[0].function.name, "get_weather");
|
||||
assert_eq!(normal_text, ""); // Pure JSON should have no normal text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extract_json_with_normal_text() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Test extraction of JSON from mixed text
|
||||
let input =
|
||||
r#"Here is some text before {"name": "test", "arguments": {}} and some text after."#;
|
||||
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,
|
||||
"Here is some text before and some text after."
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extract_json_array_with_normal_text() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
// Test extraction of JSON array from mixed text
|
||||
let input = r#"Prefix text [{"name": "func1", "arguments": {}}, {"name": "func2", "arguments": {}}] suffix text"#;
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
|
||||
assert_eq!(tool_calls.len(), 2);
|
||||
assert_eq!(tool_calls[0].function.name, "func1");
|
||||
assert_eq!(tool_calls[1].function.name, "func2");
|
||||
assert_eq!(normal_text, "Prefix text suffix text");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_multiple_tool_calls() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"[
|
||||
{"name": "get_weather", "arguments": {"location": "SF"}},
|
||||
{"name": "search", "arguments": {"query": "news"}}
|
||||
]"#;
|
||||
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 2);
|
||||
assert_eq!(tool_calls[0].function.name, "get_weather");
|
||||
assert_eq!(tool_calls[1].function.name, "search");
|
||||
assert_eq!(normal_text, ""); // Pure JSON should have no normal text
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_with_parameters_key() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"{"name": "calculate", "parameters": {"x": 10, "y": 20}}"#;
|
||||
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 1);
|
||||
assert_eq!(tool_calls[0].function.name, "calculate");
|
||||
assert!(tool_calls[0].function.arguments.contains("10"));
|
||||
assert_eq!(normal_text, ""); // Pure JSON should have no normal text
|
||||
}
|
||||
|
||||
#[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
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_parse_with_normal_text() {
|
||||
let parser = JsonParser::new();
|
||||
let input = r#"Here is the weather data: {"name": "get_weather", "arguments": {"location": "SF"}} Let me know if you need more info."#;
|
||||
|
||||
let (normal_text, tool_calls) = parser.parse_complete(input).await.unwrap();
|
||||
assert_eq!(tool_calls.len(), 1);
|
||||
assert_eq!(tool_calls[0].function.name, "get_weather");
|
||||
assert_eq!(
|
||||
normal_text,
|
||||
"Here is the weather data: Let me know if you need more info."
|
||||
); // Normal text is now extracted when JSON is found in mixed content
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_format() {
|
||||
let parser = JsonParser::new();
|
||||
|
||||
assert!(parser.detect_format(r#"{"name": "test", "arguments": {}}"#));
|
||||
assert!(parser.detect_format(r#"[{"name": "test"}]"#));
|
||||
assert!(!parser.detect_format("plain text"));
|
||||
assert!(!parser.detect_format(r#"{"key": "value"}"#));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_streaming_parse() {
|
||||
// Just verify that streaming eventually produces a complete tool call
|
||||
let parser = JsonParser::new();
|
||||
let mut state = ParseState::new();
|
||||
|
||||
// Send complete JSON in one go
|
||||
// TODO simplified version, address more complex version
|
||||
let full_json = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#;
|
||||
|
||||
let result = parser
|
||||
.parse_incremental(full_json, &mut state)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should get a complete tool immediately with complete JSON
|
||||
match result {
|
||||
StreamResult::ToolComplete(tool) => {
|
||||
assert_eq!(tool.function.name, "get_weather");
|
||||
assert!(tool.function.arguments.contains("SF"));
|
||||
}
|
||||
_ => panic!("Expected ToolComplete for complete JSON input"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user