"models/vscode:/vscode.git/clone" did not exist on "cbb19ee84ee495e819e4ef5723cc3412b237ff1d"
Unverified Commit dba751a8 authored by Chang Su's avatar Chang Su Committed by GitHub
Browse files

[router][tool call] Support normal content extraction before tool call (streaming) (#11038)

parent 2e763398
...@@ -154,8 +154,18 @@ impl ToolParser for DeepSeekParser { ...@@ -154,8 +154,18 @@ impl ToolParser for DeepSeekParser {
// Check for tool markers // Check for tool markers
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
// No markers found, return as incomplete // No tool markers detected - return all buffered content as normal text
return Ok(StreamResult::Incomplete); let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
if let Some(marker_pos) = state.buffer.find("<|tool▁calls▁begin|>") {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Look for start of tool calls // Look for start of tool calls
...@@ -220,7 +230,7 @@ impl ToolParser for DeepSeekParser { ...@@ -220,7 +230,7 @@ impl ToolParser for DeepSeekParser {
}); });
} }
Err(_) => { Err(_) => {
// Can't parse yet, keep buffering // Can't parse yet, continue waiting for more data
} }
} }
} }
......
...@@ -177,8 +177,18 @@ impl ToolParser for Glm4MoeParser { ...@@ -177,8 +177,18 @@ impl ToolParser for Glm4MoeParser {
// Check for tool markers // Check for tool markers
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
// No markers found, return as incomplete // No tool markers detected - return all buffered content as normal text
return Ok(StreamResult::Incomplete); let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
if let Some(marker_pos) = state.buffer.find("<tool_call>") {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Look for start of tool call // Look for start of tool call
......
...@@ -227,10 +227,34 @@ impl JsonParser { ...@@ -227,10 +227,34 @@ impl JsonParser {
} }
// Check for any start token // Check for any start token
self.token_config let has_start_token = self
.token_config
.start_tokens .start_tokens
.iter() .iter()
.any(|token| text.contains(token)) .any(|token| text.contains(token));
// Also check if we have what looks like JSON even without start token
// This handles cases where we've already processed the start token
// and are working on subsequent tools
has_start_token || (text.trim_start().starts_with('{') && text.contains(r#""name""#))
}
/// Check if text might contain a partial start token (for streaming)
fn has_partial_start_token(&self, text: &str) -> bool {
if self.token_config.start_tokens.is_empty() {
return false;
}
// Check if the end of the buffer could be the start of any start token
for start_token in &self.token_config.start_tokens {
for i in 1..start_token.len() {
let token_prefix = &start_token[..i];
if text.ends_with(token_prefix) {
return true;
}
}
}
false
} }
} }
...@@ -382,8 +406,42 @@ impl ToolParser for JsonParser { ...@@ -382,8 +406,42 @@ impl ToolParser for JsonParser {
// Check if we have potential tool calls // Check if we have potential tool calls
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
// No tool markers, return as incomplete if self.has_partial_start_token(&state.buffer) {
return Ok(StreamResult::Incomplete); // We might be in the middle of receiving a start token, wait for more data
return Ok(StreamResult::Incomplete);
}
// No tool markers and no partial tokens - return all buffered content as normal text
let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
if !self.token_config.start_tokens.is_empty() {
let start_token = &self.token_config.start_tokens[0];
if !start_token.is_empty() {
if let Some(marker_pos) = state.buffer.find(start_token) {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
}
}
} else {
// For JSON without start tokens, look for the start of JSON structure
// Find whichever comes first: '{' or '['
let brace_pos = state.buffer.find('{');
let bracket_pos = state.buffer.find('[');
let json_pos = brace_pos.iter().chain(bracket_pos.iter()).min().copied();
if let Some(pos) = json_pos {
if pos > 0 {
// We have text before JSON structure - extract it as normal text
let normal_text: String = state.buffer.drain(..pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
}
} }
// Extract JSON content first to check for separators // Extract JSON content first to check for separators
...@@ -407,9 +465,8 @@ impl ToolParser for JsonParser { ...@@ -407,9 +465,8 @@ impl ToolParser for JsonParser {
// We need to figure out how much to remove from the original buffer // We need to figure out how much to remove from the original buffer
// Find where the separator is in the original buffer and remove up to and including it // Find where the separator is in the original buffer and remove up to and including it
if let Some(sep_in_original) = state.buffer.find(separator.as_str()) { if let Some(sep_in_original) = state.buffer.find(separator.as_str()) {
let remaining = // Remove processed content up to and including separator
state.buffer[sep_in_original + separator.len()..].to_string(); state.buffer.drain(..=sep_in_original + separator.len() - 1);
state.buffer = remaining;
} }
// Return the first tool as complete // Return the first tool as complete
...@@ -518,7 +575,7 @@ impl ToolParser for JsonParser { ...@@ -518,7 +575,7 @@ impl ToolParser for JsonParser {
} }
Err(_) => { Err(_) => {
// Failed to parse even as partial JSON // Failed to parse even as partial JSON
// Keep buffering // Continue waiting for more data
} }
} }
......
...@@ -152,9 +152,22 @@ impl ToolParser for KimiK2Parser { ...@@ -152,9 +152,22 @@ impl ToolParser for KimiK2Parser {
self.has_tool_markers(&state.buffer) || state.buffer.contains("<|tool_call_begin|>"); self.has_tool_markers(&state.buffer) || state.buffer.contains("<|tool_call_begin|>");
if !has_tool_call { if !has_tool_call {
// No markers found, clear buffer and return // No tool markers detected - return all buffered content as normal text
state.buffer.clear(); let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::Incomplete); return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
let marker1_pos = state.buffer.find("<|tool_calls_section_begin|>");
let marker2_pos = state.buffer.find("<|tool_call_begin|>");
let marker_pos = marker1_pos.iter().chain(marker2_pos.iter()).min().copied();
if let Some(pos) = marker_pos {
if pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Try to match streaming pattern // Try to match streaming pattern
......
...@@ -75,16 +75,18 @@ impl ToolParser for LlamaParser { ...@@ -75,16 +75,18 @@ impl ToolParser for LlamaParser {
chunk: &str, chunk: &str,
state: &mut ParseState, state: &mut ParseState,
) -> ToolParserResult<StreamResult> { ) -> ToolParserResult<StreamResult> {
// Try with the python_tag parser first // First, try with the configured json_parser (which handles python_tag)
let result = self.json_parser.parse_incremental(chunk, state).await?; let result = self.json_parser.parse_incremental(chunk, state).await?;
// If we get Incomplete and buffer starts with '{', might be plain JSON // If we get Incomplete and no python_tag in buffer, might be plain JSON
if matches!(result, StreamResult::Incomplete) && state.buffer.trim_start().starts_with('{') if matches!(result, StreamResult::Incomplete) {
{ let trimmed = state.buffer.trim_start();
// Check if we have python_tag in the buffer if trimmed.starts_with('{') && !state.buffer.contains("<|python_tag|>") {
if !state.buffer.contains("<|python_tag|>") { // Likely plain JSON, try with a plain parser
// Likely plain JSON, create temporary parser // Note: We need to be careful not to double-add the chunk
let plain_parser = JsonParser::new(); let plain_parser = JsonParser::new();
// The chunk was already added to state.buffer by json_parser above
// So we call with empty string to just process what's in the buffer
return plain_parser.parse_incremental("", state).await; return plain_parser.parse_incremental("", state).await;
} }
} }
......
...@@ -195,7 +195,18 @@ impl ToolParser for MistralParser { ...@@ -195,7 +195,18 @@ impl ToolParser for MistralParser {
// Check if we have the start marker // Check if we have the start marker
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
return Ok(StreamResult::Incomplete); // No tool markers detected - return all buffered content as normal text
let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before [TOOL_CALLS] and extract it as normal text
if let Some(marker_pos) = state.buffer.find("[TOOL_CALLS]") {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Try to extract complete JSON array // Try to extract complete JSON array
......
...@@ -190,7 +190,18 @@ impl ToolParser for QwenParser { ...@@ -190,7 +190,18 @@ impl ToolParser for QwenParser {
// Check if we have the start marker // Check if we have the start marker
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
return Ok(StreamResult::Incomplete); // No tool markers detected - return all buffered content as normal text
let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
if let Some(marker_pos) = state.buffer.find("<tool_call>") {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Find start and end positions // Find start and end positions
...@@ -212,7 +223,12 @@ impl ToolParser for QwenParser { ...@@ -212,7 +223,12 @@ impl ToolParser for QwenParser {
} }
} }
Err(_) => { Err(_) => {
// JSON parsing failed, might be incomplete // JSON parsing failed, might be incomplete or malformed
// If we have what looks like a complete tool call block, treat as normal text
if state.buffer[start_pos..end_pos].contains("\n</tool_call>") {
let malformed_text: String = state.buffer.drain(..end_pos).collect();
return Ok(StreamResult::NormalText(malformed_text));
}
} }
} }
} else { } else {
......
...@@ -209,8 +209,18 @@ impl ToolParser for Step3Parser { ...@@ -209,8 +209,18 @@ impl ToolParser for Step3Parser {
// Check for tool markers // Check for tool markers
if !self.has_tool_markers(&state.buffer) { if !self.has_tool_markers(&state.buffer) {
// No markers found, return as incomplete // No tool markers detected - return all buffered content as normal text
return Ok(StreamResult::Incomplete); let normal_text = std::mem::take(&mut state.buffer);
return Ok(StreamResult::NormalText(normal_text));
}
// Check for text before tool markers and extract it as normal text
if let Some(marker_pos) = state.buffer.find("<|tool_calls_begin|>") {
if marker_pos > 0 {
// We have text before the tool marker - extract it as normal text
let normal_text: String = state.buffer.drain(..marker_pos).collect();
return Ok(StreamResult::NormalText(normal_text));
}
} }
// Look for start of tool calls // Look for start of tool calls
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment