Restore DSML tool-call recovery
This commit is contained in:
26
src/dsml.rs
26
src/dsml.rs
@@ -45,10 +45,11 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
|
||||
else {
|
||||
return Ok((text.to_owned(), Vec::new()));
|
||||
};
|
||||
let Some(relative_end) = text[start..].find(syntax.tool_end) else {
|
||||
return Err("invalid or incomplete DSML tool call".into());
|
||||
};
|
||||
let end = start + relative_end + syntax.tool_end.len();
|
||||
let relative_end = text[start..].find(syntax.tool_end);
|
||||
let end = relative_end
|
||||
.map(|relative_end| start + relative_end + syntax.tool_end.len())
|
||||
.unwrap_or(text.len());
|
||||
let outer_call_complete = relative_end.is_some();
|
||||
let content = text[..start].trim_end().to_owned();
|
||||
let raw = &text[start..end];
|
||||
let mut cursor = syntax.tool_start.len();
|
||||
@@ -58,6 +59,11 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
|
||||
if raw[cursor..].starts_with(syntax.tool_end) {
|
||||
break;
|
||||
}
|
||||
// DS4 recovers calls whose invokes are complete even if generation
|
||||
// stopped before emitting the outer tool-calls closing tag.
|
||||
if !outer_call_complete && cursor == raw.len() && !calls.is_empty() {
|
||||
break;
|
||||
}
|
||||
if !raw[cursor..].starts_with(syntax.invoke_start) {
|
||||
return Err("invalid or incomplete DSML tool call".into());
|
||||
}
|
||||
@@ -207,4 +213,16 @@ mod tests {
|
||||
assert_eq!(calls[0].1["path"], "src/main.rs");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recovers_complete_invokes_without_the_outer_closing_tag() {
|
||||
let raw = "before<|DSML|tool_calls><|DSML|invoke name=\"read\"><|DSML|parameter name=\"path\" string=\"true\">src/main.rs</|DSML|parameter></|DSML|invoke>";
|
||||
|
||||
let (content, calls) = parse_tool_calls(raw).unwrap();
|
||||
|
||||
assert_eq!(content, "before");
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].0, "read");
|
||||
assert_eq!(calls[0].1["path"], "src/main.rs");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user