tool_parser_gpt_oss.rs 6.58 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
//! GPT-OSS Parser Integration Tests

use sglang_router_rs::tool_parser::{GptOssParser, ParseState, StreamResult, ToolParser};

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

    let input = r#"Let me search for that information.
<|channel|>commentary to=functions.search<|constrain|>json<|message|>{"query": "rust programming", "limit": 10}<|call|>
Here are the results..."#;

13
14
15
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "search");
16

17
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
18
19
20
21
22
23
24
25
26
27
28
    assert_eq!(args["query"], "rust programming");
    assert_eq!(args["limit"], 10);
}

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

    let input = r#"<|channel|>commentary to=functions.get_weather<|constrain|>json<|message|>{"location": "Paris"}<|call|>commentary
<|channel|>commentary to=functions.search<|constrain|>json<|message|>{"query": "Paris tourism"}<|call|>"#;

29
30
31
32
    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, "search");
33
34
35
36
37
38
39
40
41
}

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

    let input = r#"<|channel|>commentary to=api.users.create<|constrain|>json<|message|>{"name": "John", "email": "john@example.com"}<|call|>
<|channel|>commentary to=tools.calculator.add<|constrain|>json<|message|>{"x": 10, "y": 20}<|call|>"#;

42
43
44
45
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 2);
    assert_eq!(tools[0].function.name, "create"); // Should extract last part
    assert_eq!(tools[1].function.name, "add");
46
47
48
49
50
51
52
53
}

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

    let input = r#"<|start|>assistant<|channel|>commentary to=functions.test<|constrain|>json<|message|>{"key": "value"}<|call|>"#;

54
55
56
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "test");
57
58
59
60
61
62
63
64
65
}

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

    let input =
        r#"<|channel|>commentary to=functions.get_time<|constrain|>json<|message|>{}<|call|>"#;

66
67
68
69
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "get_time");
    assert_eq!(tools[0].function.arguments, "{}");
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
}

#[tokio::test]
async fn test_gpt_oss_streaming() {
    let parser = GptOssParser::new();
    let mut state = ParseState::new();

    // Simulate streaming chunks
    let chunks = vec![
        "<|channel|>commentary to=",
        "functions.calculate",
        "<|constrain|>json<|message|>",
        r#"{"x": 10"#,
        r#", "y": 20}"#,
        "<|call|>",
    ];

    let mut found_name = false;
    let mut found_complete = false;

    for chunk in chunks {
        let result = parser.parse_incremental(chunk, &mut state).await.unwrap();

        match result {
            StreamResult::ToolName { name, .. } => {
                assert_eq!(name, "calculate");
                found_name = true;
            }
            StreamResult::ToolComplete(tool) => {
                assert_eq!(tool.function.name, "calculate");
                found_complete = true;
            }
            _ => {}
        }
    }

    assert!(found_name || found_complete);
}

#[test]
fn test_gpt_oss_format_detection() {
    let parser = GptOssParser::new();

    // Should detect GPT-OSS format
    assert!(parser.detect_format("<|channel|>commentary to="));
    assert!(parser.detect_format("<|channel|>commentary"));
    assert!(parser.detect_format("text with <|channel|>commentary to= marker"));

    // Should not detect other formats
    assert!(!parser.detect_format("[TOOL_CALLS]"));
    assert!(!parser.detect_format("<tool_call>"));
    assert!(!parser.detect_format("plain text"));
}

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

    let input = r#"<|channel|>commentary to=functions.test  <|constrain|>json<|message|>{"key": "value"}<|call|>"#;

130
131
132
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "test");
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
}

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

    let input = r#"<|channel|>commentary to=functions.process<|constrain|>json<|message|>{
    "nested": {
        "data": [1, 2, 3],
        "config": {
            "enabled": true
        }
    }
}<|call|>"#;

148
149
150
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "process");
151

152
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
153
154
155
156
157
158
159
160
161
162
163
    assert!(args["nested"]["data"].is_array());
    assert_eq!(args["nested"]["config"]["enabled"], true);
}

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

    // Python should extract commentary as normal text
    let input = r#"<|channel|>commentary<|message|>**Action plan**: 1. Do X 2. Do Y<|end|>"#;

164
165
166
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 0); // No tool calls
                                // TODO: Verify normal text = "**Action plan**: 1. Do X 2. Do Y"
167
168
169
170
171
172
173
174
175
}

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

    let input = r#"<|channel|>commentary to=functions.test<|constrain|>json<|message|>{"x": 1}<|call|>
<|channel|>final<|message|>The result is calculated.<|return|>"#;

176
177
178
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "test");
179
180
181
182
183
184
185
186
187
188
189
    // TODO: Verify normal text = "The result is calculated."
}

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

    let input = r#"<|channel|>commentary<|message|>Let me think<|end|>
<|channel|>commentary to=functions.calc<|constrain|>json<|message|>{"x": 5}<|call|>
<|channel|>commentary<|message|>Processing...<|end|>"#;

190
191
192
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "calc");
193
194
    // TODO: Verify normal text = "Let me think Processing..."
}