tool_parser_step3.rs 8.4 KB
Newer Older
1
2
//! Step3 Parser Integration Tests

3
4
5
6
use sglang_router_rs::tool_parser::{Step3Parser, ToolParser};

mod common;
use common::create_test_tools;
7
8
9
10
11
12
13
14
15
16
17
18
19
20

#[tokio::test]
async fn test_step3_complete_parsing() {
    let parser = Step3Parser::new();

    let input = r#"Let me help you.
<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="search">
<steptml:parameter name="query">rust programming</steptml:parameter>
<steptml:parameter name="limit">10</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>
Here are the results..."#;

21
    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
22
    assert_eq!(tools.len(), 1);
23
    assert_eq!(normal_text, "Let me help you.\n");
24
    assert_eq!(tools[0].function.name, "search");
25

26
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
    assert_eq!(args["query"], "rust programming");
    assert_eq!(args["limit"], 10);
}

#[tokio::test]
async fn test_step3_multiple_tools() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="get_weather">
<steptml:parameter name="location">Tokyo</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="get_news">
<steptml:parameter name="category">tech</steptml:parameter>
<steptml:parameter name="limit">5</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

45
46
47
48
    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");
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
}

#[tokio::test]
async fn test_step3_type_conversion() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="process">
<steptml:parameter name="count">100</steptml:parameter>
<steptml:parameter name="rate">2.5</steptml:parameter>
<steptml:parameter name="active">true</steptml:parameter>
<steptml:parameter name="optional">null</steptml:parameter>
<steptml:parameter name="text">hello world</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

65
66
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
67

68
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
69
70
71
72
73
74
75
76
77
    assert_eq!(args["count"], 100);
    assert_eq!(args["rate"], 2.5);
    assert_eq!(args["active"], true);
    assert_eq!(args["optional"], serde_json::Value::Null);
    assert_eq!(args["text"], "hello world");
}

#[tokio::test]
async fn test_step3_streaming() {
78
79
80
    let mut parser = Step3Parser::new();

    let tools = create_test_tools();
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

    // Simulate streaming chunks
    let chunks = vec![
        "<|tool_calls_begin|>\n",
        "<|tool_call_begin|>function",
        "<|tool_sep|><steptml:invoke name=\"calc\">",
        "\n<steptml:parameter name=\"x\">10</steptml:parameter>",
        "\n<steptml:parameter name=\"y\">20</steptml:parameter>",
        "\n</steptml:invoke><|tool_call_end|>",
        "\n<|tool_calls_end|>",
    ];

    let mut found_complete = false;

    for chunk in chunks {
96
        let result = parser.parse_incremental(chunk, &tools).await.unwrap();
97

98
99
        if !result.calls.is_empty() {
            if let Some(name) = &result.calls[0].name {
100
101
102
103
104
105
                assert_eq!(name, "calc");
                found_complete = true;
            }
        }
    }

106
    assert!(found_complete);
107
108
109
110
111
112
113
}

#[test]
fn test_step3_format_detection() {
    let parser = Step3Parser::new();

    // Should detect Step3 format
114
115
    assert!(parser.has_tool_markers("<|tool_calls_begin|>"));
    assert!(parser.has_tool_markers("text with <|tool_calls_begin|> marker"));
116
117

    // Should not detect other formats
118
119
120
    assert!(!parser.has_tool_markers("[TOOL_CALLS]"));
    assert!(!parser.has_tool_markers("<tool_call>"));
    assert!(!parser.has_tool_markers("plain text"));
121
122
123
124
125
126
127
128
129
130
131
132
133
}

#[tokio::test]
async fn test_step3_nested_steptml() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="config">
<steptml:parameter name="settings">{"nested": {"key": "value"}}</steptml:parameter>
<steptml:parameter name="array">[1, 2, 3]</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

134
135
136
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "config");
137

138
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
    assert!(args["settings"].is_object());
    assert!(args["array"].is_array());
}

#[tokio::test]
async fn test_step3_python_literals() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="test">
<steptml:parameter name="bool_true">True</steptml:parameter>
<steptml:parameter name="bool_false">False</steptml:parameter>
<steptml:parameter name="none_value">None</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

155
156
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
157

158
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    assert_eq!(args["bool_true"], true);
    assert_eq!(args["bool_false"], false);
    assert_eq!(args["none_value"], serde_json::Value::Null);
}

#[tokio::test]
async fn test_steptml_format() {
    let parser = Step3Parser::new();

    let input = r#"Text before.
<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="search">
<steptml:parameter name="query">rust lang</steptml:parameter>
<steptml:parameter name="limit">10</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>Text after."#;

176
    let (normal_text, tools) = parser.parse_complete(input).await.unwrap();
177
    assert_eq!(tools.len(), 1);
178
    assert_eq!(normal_text, "Text before.\n");
179
    assert_eq!(tools[0].function.name, "search");
180

181
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
    assert_eq!(args["query"], "rust lang");
    assert_eq!(args["limit"], 10);
    // TODO: Verify normal text extraction
}

#[tokio::test]
async fn test_json_parameter_values() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="config">
<steptml:parameter name="settings">{"nested": {"value": true}}</steptml:parameter>
<steptml:parameter name="items">[1, 2, 3]</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

198
199
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
200

201
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
    assert!(args["settings"].is_object());
    assert!(args["items"].is_array());
}

#[tokio::test]
async fn test_step3_parameter_with_angle_brackets() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="compare">
<steptml:parameter name="expression">a < b && b > c</steptml:parameter>
<steptml:parameter name="context">comparison test</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

217
218
219
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "compare");
220

221
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
222
223
224
225
226
227
228
229
230
231
232
233
234
235
    assert_eq!(args["expression"], "a < b && b > c");
    assert_eq!(args["context"], "comparison test");
}

#[tokio::test]
async fn test_step3_empty_function_name() {
    let parser = Step3Parser::new();

    let input = r#"<|tool_calls_begin|>
<|tool_call_begin|>function<|tool_sep|><steptml:invoke name="">
<steptml:parameter name="param">value</steptml:parameter>
</steptml:invoke><|tool_call_end|>
<|tool_calls_end|>"#;

236
237
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 0); // Should reject empty function name
238
}