tool_parser_deepseek.rs 4.86 KB
Newer Older
1
2
//! DeepSeek V3 Parser Integration Tests

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

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

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

    let input = r#"Let me help you with that.
<|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_weather
```json
{"location": "Tokyo", "units": "celsius"}
```<|tool▁call▁end|><|tool▁calls▁end|>
The weather in Tokyo is..."#;

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

24
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
    assert_eq!(args["location"], "Tokyo");
    assert_eq!(args["units"], "celsius");
}

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

    let input = r#"<|tool▁calls▁begin|>
<|tool▁call▁begin|>function<|tool▁sep|>search
```json
{"query": "rust programming"}
```<|tool▁call▁end|>
<|tool▁call▁begin|>function<|tool▁sep|>translate
```json
{"text": "Hello World", "to": "ja"}
```<|tool▁call▁end|>
<|tool▁calls▁end|>"#;

44
45
46
47
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 2);
    assert_eq!(tools[0].function.name, "search");
    assert_eq!(tools[1].function.name, "translate");
48
49
50
51
}

#[tokio::test]
async fn test_deepseek_streaming() {
52
53
54
    let tools = create_test_tools();

    let mut parser = DeepSeekParser::new();
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

    // Simulate streaming chunks
    let chunks = vec![
        "<|tool▁calls▁begin|><|tool▁call▁begin|>",
        "function<|tool▁sep|>get_weather\n",
        "```json\n",
        r#"{"location": "#,
        r#""Beijing", "#,
        r#""units": "metric"}"#,
        "\n```<|tool▁call▁end|><|tool▁calls▁end|>",
    ];

    let mut found_name = false;

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

72
73
        for call in result.calls {
            if let Some(name) = call.name {
74
75
76
77
78
79
                assert_eq!(name, "get_weather");
                found_name = true;
            }
        }
    }

80
    assert!(found_name, "Should have found tool name during streaming");
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
}

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

    let input = r#"<|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>process
```json
{
    "data": {
        "nested": {
            "deep": [1, 2, 3]
        }
    }
}
```<|tool▁call▁end|><|tool▁calls▁end|>"#;

98
99
100
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "process");
101

102
    let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap();
103
104
105
106
107
108
109
110
    assert!(args["data"]["nested"]["deep"].is_array());
}

#[test]
fn test_deepseek_format_detection() {
    let parser = DeepSeekParser::new();

    // Should detect DeepSeek format
111
112
    assert!(parser.has_tool_markers("<|tool▁calls▁begin|>"));
    assert!(parser.has_tool_markers("text with <|tool▁calls▁begin|> marker"));
113
114

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

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

    // Malformed JSON should be skipped
    let input = r#"<|tool▁calls▁begin|>
<|tool▁call▁begin|>function<|tool▁sep|>broken
```json
{invalid json}
```<|tool▁call▁end|>
<|tool▁call▁begin|>function<|tool▁sep|>valid
```json
{"key": "value"}
```<|tool▁call▁end|>
<|tool▁calls▁end|>"#;

136
    let (_normal_text, tools) = parser.parse_complete(input).await.unwrap();
137
    // Only the valid tool call should be parsed
138
139
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].function.name, "valid");
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
}

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

    let input = r#"<|tool▁calls▁begin|>
<|tool▁call▁begin|>function<|tool▁sep|>get_weather
```json
{"location": "Tokyo"}
```<|tool▁call▁end|>
<|tool▁call▁begin|>function<|tool▁sep|>get_weather
```json
{"location": "Paris"}
```<|tool▁call▁end|>
<|tool▁calls▁end|><|end▁of▁sentence|>"#;

157
158
159
160
    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_weather");
161
}