fix: slipping <tool_calls_end> tool call parsing for DeepSeek v3.1 (#3995)

Signed-off-by: Guan Luo <[email protected]>

Author: GuanLuo

lib/llm/tests/test_jail.rs

@@ -1792,6 +1792,123 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_deepseek_v3_1() {
+        // DeepSeek v3.1 format with two tool calls encoded in special tags
+        let text = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_current_weather<｜tool▁sep｜>{"location": "Berlin", "units": "metric"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_weather_forecast<｜tool▁sep｜>{"location": "Berlin", "days": 7, "units": "imperial"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_air_quality<｜tool▁sep｜>{"location": "Berlin", "radius": 50}<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜>"#;
+
+        let chunks = vec![create_mock_response_chunk(text.to_string(), 0)];
+
+        let input_stream = stream::iter(chunks);
+
+        let jail = JailedStream::builder()
+            .tool_call_parser("deepseek_v3_1")
+            .build();
+        let jailed_stream = jail.apply(input_stream);
+        let results: Vec<_> = jailed_stream.collect().await;
+
+        // Should have at least one output containing both analysis text and parsed tool call
+        assert!(!results.is_empty());
+
+        // Verify a tool call was parsed with expected name and args
+        let tool_call_idx = results
+            .iter()
+            .position(test_utils::has_tool_call)
+            .expect("Should have a tool call result");
+        test_utils::assert_tool_call(
+            &results[tool_call_idx],
+            "get_current_weather",
+            json!({"location": "Berlin", "units": "metric"}),
+        );
+        for result in results {
+            let Some(data) = result.data else {
+                continue;
+            };
+            for choice in data.choices {
+                if let Some(content) = choice.delta.content {
+                    assert!(
+                        !content.contains("<｜tool▁calls▁end｜>"),
+                        "Should not contain deepseek special tokens in content"
+                    );
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_deepseek_v3_1_chunk() {
+        // DeepSeek v3.1 format with two tool calls encoded in special tags
+        let text = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_current_weather<｜tool▁sep｜>{"location": "Berlin", "units": "metric"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_weather_forecast<｜tool▁sep｜>{"location": "Berlin", "days": 7, "units": "imperial"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_air_quality<｜tool▁sep｜>{"location": "Berlin", "radius": 50}<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜>"#;
+
+        // Split text into words, treating angle-bracketed tokens as one word
+        let mut words = Vec::new();
+        let mut i = 0;
+        let chars: Vec<char> = text.chars().collect();
+        while i < chars.len() {
+            if chars[i] == '<' {
+                // Find the next '>'
+                if let Some(end) = chars[i..].iter().position(|&c| c == '>') {
+                    let word: String = chars[i..=i + end].iter().collect();
+                    words.push(word);
+                    i += end + 1;
+                } else {
+                    // Malformed, just push the rest
+                    words.push(chars[i..].iter().collect());
+                    break;
+                }
+            } else if chars[i].is_whitespace() {
+                i += 1;
+            } else {
+                // Collect until next whitespace or '<'
+                let start = i;
+                while i < chars.len() && !chars[i].is_whitespace() && chars[i] != '<' {
+                    i += 1;
+                }
+                words.push(chars[start..i].iter().collect());
+            }
+        }
+
+        let chunks = words
+            .into_iter()
+            .map(|word| create_mock_response_chunk(word, 0))
+            .collect::<Vec<_>>();
+
+        let input_stream = stream::iter(chunks);
+
+        let jail = JailedStream::builder()
+            .tool_call_parser("deepseek_v3_1")
+            .build();
+        let jailed_stream = jail.apply(input_stream);
+        let results: Vec<_> = jailed_stream.collect().await;
+
+        // Should have at least one output containing both analysis text and parsed tool call
+        assert!(!results.is_empty());
+
+        // Verify a tool call was parsed with expected name and args
+        let tool_call_idx = results
+            .iter()
+            .position(test_utils::has_tool_call)
+            .expect("Should have a tool call result");
+        test_utils::assert_tool_call(
+            &results[tool_call_idx],
+            "get_current_weather",
+            json!({"location": "Berlin", "units": "metric"}),
+        );
+        for result in results {
+            let Some(data) = result.data else {
+                continue;
+            };
+            for choice in data.choices {
+                if let Some(content) = choice.delta.content {
+                    assert!(
+                        !content.contains("<｜tool▁calls▁end｜>"),
+                        "Should not contain deepseek special tokens in content"
+                    );
+                }
+            }
+        }
+    }
+
     #[tokio::test]
     async fn test_jailed_stream_mistral_false_positive_curly() {
         // Curly brace in normal text should not trigger tool call detection for mistral

lib/parsers/src/tool_calling/config.rs

@@ -153,16 +153,22 @@ impl ToolCallConfig {
     }
 
     pub fn deepseek_v3_1() -> Self {
+        // The whole tool calls block is wrapped between
+        // <｜tool▁calls▁begin｜> ... <｜tool▁calls▁end｜>
+        // regardless of number of tool calls. For external use of this
+        // config, we want them to only be operating on the whole block,
+        // so the tool parser can properly consume all tool call tokens.
+        // https://huggingface.co/deepseek-ai/DeepSeek-V3.1#toolcall
         Self {
             format: ToolCallParserType::Json,
             json: JsonParserConfig {
                 tool_call_start_tokens: vec![
                     "<｜tool▁calls▁begin｜>".to_string(),
-                    "<｜tool▁call▁begin｜>".to_string(),
+                    // "<｜tool▁call▁begin｜>".to_string(),
                 ],
                 tool_call_end_tokens: vec![
                     "<｜tool▁calls▁end｜>".to_string(),
-                    "<｜tool▁call▁end｜>".to_string(),
+                    // "<｜tool▁call▁end｜>".to_string(),
                 ],
                 tool_call_separator_tokens: vec!["<｜tool▁sep｜>".to_string()],
                 parser_type: JsonParserType::DeepseekV31,

lib/parsers/src/tool_calling/json/deepseek_parser.rs

@@ -126,8 +126,28 @@ pub fn parse_tool_calls_deepseek_v3_1(
         return Ok((vec![], Some(String::new())));
     }
 
-    let tool_call_start_tokens = &config.tool_call_start_tokens;
-    let tool_call_end_tokens = &config.tool_call_end_tokens;
+    // For DeepSeek_v3_1, we consider the tool call block to be
+    // <｜tool▁calls▁begin｜>...<｜tool▁calls▁end｜> and only start parsing
+    // if seeing <｜tool▁calls▁begin｜>, even though the individual calls are
+    // parsed by <｜tool▁call▁begin｜>...<｜tool▁call▁end｜>.
+    // This is because if we start parsing by considering all call(s) tokens,
+    // we are not properly grouping the tool calls and results in groups:
+    // 1. <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>...<｜tool▁call▁end｜>
+    // 2. <｜tool▁calls▁end｜>
+    // where 2. will not be recognized as part of the tool call block due
+    // to missing start token and will not be consumed.
+    let has_end_token = config
+        .tool_call_end_tokens
+        .iter()
+        .any(|token| !token.is_empty() && trimmed.contains(token));
+    if !has_end_token {
+        return Ok((vec![], Some(trimmed.to_string())));
+    }
+
+    let mut tool_call_start_tokens = config.tool_call_start_tokens.clone();
+    tool_call_start_tokens.extend(vec!["<｜tool▁call▁begin｜>".to_string()]);
+    let mut tool_call_end_tokens = config.tool_call_end_tokens.clone();
+    tool_call_end_tokens.extend(vec!["<｜tool▁call▁end｜>".to_string()]);
     let separator_tokens = &config.tool_call_separator_tokens;
 
     // Early exit if no tokens configured
@@ -166,7 +186,7 @@ pub fn parse_tool_calls_deepseek_v3_1(
     };
 
     // Extract individual tool call blocks
-    let blocks = extract_tool_call_blocks(trimmed, tool_call_start_tokens, tool_call_end_tokens);
+    let blocks = extract_tool_call_blocks(trimmed, &tool_call_start_tokens, &tool_call_end_tokens);
 
     if blocks.is_empty() {
         // Found start token but no valid blocks
@@ -398,7 +418,7 @@ mod detect_parser_tests {
         let text = r#"<｜tool▁call▁begin｜>get_current_weather宽带}"#;
         let config = ToolCallConfig::deepseek_v3_1().json;
         let result = detect_tool_call_start_deepseek_v3_1(text, &config);
-        assert!(result);
+        assert!(!result);
     }
 
     #[test]

lib/parsers/src/tool_calling/parsers.rs

@@ -2413,15 +2413,15 @@ mod detect_parser_tests {
     }
 
     #[test]
-    fn test_e2e_detect_tool_call_start_deepseek_v3_1() {
+    fn test_e2e_detect_incomplete_tool_call_start_deepseek_v3_1() {
         let text =
             r#"<｜tool▁call▁begin｜>get_current_weather{"location": "Tokyo"}<｜tool▁call▁end｜>"#;
         let result = detect_tool_call_start(text, Some("deepseek_v3_1")).unwrap();
-        assert!(result);
+        assert!(!result);
     }
 
     #[test]
-    fn test_e2e_detect_tool_call_multiple_start_deepseek_v3_1() {
+    fn test_e2e_detect_tool_call_start_deepseek_v3_1() {
         let text = r#"<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_current_weather{"location": "Tokyo"}<｜tool▁call▁end｜>"#;
         let result = detect_tool_call_start(text, Some("deepseek_v3_1")).unwrap();
         assert!(result);