Support Linux refactoring checks

src/dsml.rs

@@ -0,0 +1,210 @@
+use serde_json::{Map, Value};
+
+#[derive(Clone, Copy)]
+pub(crate) struct Syntax {
+    pub(crate) tool_start: &'static str,
+    pub(crate) tool_end: &'static str,
+    pub(crate) invoke_start: &'static str,
+    pub(crate) invoke_end: &'static str,
+    pub(crate) parameter_start: &'static str,
+    pub(crate) parameter_end: &'static str,
+}
+
+pub(crate) const SYNTAXES: [Syntax; 3] = [
+    Syntax {
+        tool_start: "<｜DSML｜tool_calls>",
+        tool_end: "</｜DSML｜tool_calls>",
+        invoke_start: "<｜DSML｜invoke",
+        invoke_end: "</｜DSML｜invoke>",
+        parameter_start: "<｜DSML｜parameter",
+        parameter_end: "</｜DSML｜parameter>",
+    },
+    Syntax {
+        tool_start: "<DSML｜tool_calls>",
+        tool_end: "</DSML｜tool_calls>",
+        invoke_start: "<DSML｜invoke",
+        invoke_end: "</DSML｜invoke>",
+        parameter_start: "<DSML｜parameter",
+        parameter_end: "</DSML｜parameter>",
+    },
+    Syntax {
+        tool_start: "<tool_calls>",
+        tool_end: "</tool_calls>",
+        invoke_start: "<invoke",
+        invoke_end: "</invoke>",
+        parameter_start: "<parameter",
+        parameter_end: "</parameter>",
+    },
+];
+
+pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value)>), String> {
+    let Some((start, syntax)) = SYNTAXES
+        .iter()
+        .filter_map(|syntax| text.find(syntax.tool_start).map(|start| (start, *syntax)))
+        .min_by_key(|(start, _)| *start)
+    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 content = text[..start].trim_end().to_owned();
+    let raw = &text[start..end];
+    let mut cursor = syntax.tool_start.len();
+    let mut calls = Vec::new();
+    loop {
+        skip_whitespace(raw, &mut cursor);
+        if raw[cursor..].starts_with(syntax.tool_end) {
+            break;
+        }
+        if !raw[cursor..].starts_with(syntax.invoke_start) {
+            return Err("invalid or incomplete DSML tool call".into());
+        }
+        let tag_end = raw[cursor..]
+            .find('>')
+            .map(|offset| cursor + offset + 1)
+            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+        let name = attribute(&raw[cursor..tag_end], "name")
+            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+        cursor = tag_end;
+        let mut arguments = Map::new();
+        loop {
+            skip_whitespace(raw, &mut cursor);
+            if raw[cursor..].starts_with(syntax.invoke_end) {
+                cursor += syntax.invoke_end.len();
+                break;
+            }
+            let (name, value) = parse_parameter(raw, &mut cursor, syntax)?
+                .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+            arguments.insert(name, value);
+        }
+        calls.push((name, Value::Object(arguments)));
+    }
+    if calls.is_empty() {
+        return Err("invalid or incomplete DSML tool call".into());
+    }
+    Ok((content, calls))
+}
+
+fn parse_parameter(
+    text: &str,
+    cursor: &mut usize,
+    syntax: Syntax,
+) -> Result<Option<(String, Value)>, String> {
+    if !text[*cursor..].starts_with(syntax.parameter_start) {
+        return Ok(None);
+    }
+    let Some(tag_end) = text[*cursor..].find('>').map(|end| *cursor + end + 1) else {
+        return Ok(None);
+    };
+    let tag = &text[*cursor..tag_end];
+    let Some(name) = attribute(tag, "name") else {
+        return Ok(None);
+    };
+    let is_string = attribute(tag, "string");
+    *cursor = tag_end;
+    let mut nested_start = *cursor;
+    skip_whitespace(text, &mut nested_start);
+    if is_string.is_none() && text[nested_start..].starts_with(syntax.parameter_start) {
+        *cursor = nested_start;
+        let mut nested = Map::new();
+        loop {
+            skip_whitespace(text, cursor);
+            if !text[*cursor..].starts_with(syntax.parameter_start) {
+                break;
+            }
+            let Some((name, value)) = parse_parameter(text, cursor, syntax)? else {
+                return Ok(None);
+            };
+            nested.insert(name, value);
+        }
+        skip_whitespace(text, cursor);
+        if !text[*cursor..].starts_with(syntax.parameter_end) {
+            return Ok(None);
+        }
+        *cursor += syntax.parameter_end.len();
+        return Ok(Some((name, Value::Object(nested))));
+    }
+    let Some(value_end) = text[*cursor..]
+        .find(syntax.parameter_end)
+        .map(|end| *cursor + end)
+    else {
+        return Ok(None);
+    };
+    let raw = &text[*cursor..value_end];
+    *cursor = value_end + syntax.parameter_end.len();
+    let value = if is_string.as_deref().unwrap_or("true") == "true" {
+        Value::String(unescape(raw))
+    } else {
+        serde_json::from_str(raw)
+            .map_err(|error| format!("invalid DSML tool arguments: {error}"))?
+    };
+    Ok(Some((name, value)))
+}
+
+fn attribute(tag: &str, name: &str) -> Option<String> {
+    let start = tag.find(&format!("{name}=\""))? + name.len() + 2;
+    let end = start + tag[start..].find('"')?;
+    Some(unescape(&tag[start..end]))
+}
+
+fn skip_whitespace(text: &str, cursor: &mut usize) {
+    while text[*cursor..]
+        .chars()
+        .next()
+        .is_some_and(char::is_whitespace)
+    {
+        *cursor += text[*cursor..].chars().next().unwrap().len_utf8();
+    }
+}
+
+fn unescape(text: &str) -> String {
+    text.replace("&quot;", "\"")
+        .replace("&gt;", ">")
+        .replace("&lt;", "<")
+        .replace("&amp;", "&")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parses_reference_and_legacy_dsml_syntaxes() {
+        for (start, end, invoke, invoke_end, parameter, parameter_end) in [
+            (
+                "<｜DSML｜tool_calls>",
+                "</｜DSML｜tool_calls>",
+                "<｜DSML｜invoke",
+                "</｜DSML｜invoke>",
+                "<｜DSML｜parameter",
+                "</｜DSML｜parameter>",
+            ),
+            (
+                "<DSML｜tool_calls>",
+                "</DSML｜tool_calls>",
+                "<DSML｜invoke",
+                "</DSML｜invoke>",
+                "<DSML｜parameter",
+                "</DSML｜parameter>",
+            ),
+            (
+                "<tool_calls>",
+                "</tool_calls>",
+                "<invoke",
+                "</invoke>",
+                "<parameter",
+                "</parameter>",
+            ),
+        ] {
+            let raw = format!(
+                "done{start}{invoke} name=\"read\">{parameter} name=\"path\" string=\"true\">src/main.rs{parameter_end}{invoke_end}{end}"
+            );
+            let (content, calls) = parse_tool_calls(&raw).unwrap();
+            assert_eq!(content, "done");
+            assert_eq!(calls[0].0, "read");
+            assert_eq!(calls[0].1["path"], "src/main.rs");
+        }
+    }
+}