ravenrs_ai/
tools.rs

1//! Letting a model choose what to do: it names one tool, and the app does the rest.
2
3use std::collections::HashMap;
4
5use serde_json::Value;
6
7/// Something a model can ask the app to do.
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub struct Tool {
10    /// Its name.
11    pub name: &'static str,
12    /// When to choose it, for the model.
13    pub description: &'static str,
14}
15
16/// What a model chose.
17#[derive(Clone, Debug, Default, PartialEq, Eq)]
18pub struct ToolCall {
19    /// The tool.
20    pub name: String,
21    /// Its arguments, as text.
22    pub arguments: HashMap<String, String>,
23}
24
25impl ToolCall {
26    /// An argument, if it came and is not blank.
27    pub fn argument(&self, name: &str) -> Option<&str> {
28        self.arguments.get(name).map(|value| value.trim()).filter(|value| !value.is_empty())
29    }
30}
31
32/// Instructions that make a model answer with one tool. Pass them as the system
33/// text of [`LocalAssistant::write`](crate::LocalAssistant::write), with the person's
34/// words as the question.
35pub fn tool_rules(intro: &str, tools: &[Tool]) -> String {
36    let list: String = tools.iter().map(|tool| format!("- {}: {}\n", tool.name, tool.description)).collect();
37    format!("{intro} Answer only with JSON in the form {{\"tool\":\"name\",\"args\":{{}}}}.\nTools:\n{list}")
38}
39
40/// Reads what a model answered. Anything but a known tool, however it is written,
41/// is no choice at all, so a model cannot ask for what the app did not offer.
42pub fn read_tool_call(answer: &str, tools: &[Tool]) -> Option<ToolCall> {
43    let (start, end) = (answer.find('{')?, answer.rfind('}')?);
44    let value: Value = serde_json::from_str(answer.get(start..=end)?).ok()?;
45    let folded = |name: &str| name.chars().flat_map(ravenrs_text::fold).collect::<String>();
46    let written = folded(value.get("tool")?.as_str()?);
47    let name = tools.iter().find(|tool| folded(tool.name) == written)?.name;
48    let arguments = value
49        .get("args")
50        .and_then(Value::as_object)
51        .map(|args| {
52            args.iter()
53                .filter_map(|(key, value)| {
54                    let text = match value {
55                        Value::String(text) => text.clone(),
56                        Value::Number(number) => number.to_string(),
57                        _ => return None,
58                    };
59                    Some((key.clone(), text))
60                })
61                .collect()
62        })
63        .unwrap_or_default();
64    Some(ToolCall { name: name.to_owned(), arguments })
65}
66
67/// Whether `text` has a whole JSON object: the braces it opened are all closed.
68pub(crate) fn json_closed(text: &str) -> bool {
69    let (mut depth, mut in_text, mut escaped) = (0_usize, false, false);
70    for letter in text.chars() {
71        match letter {
72            _ if escaped => escaped = false,
73            '\\' if in_text => escaped = true,
74            '"' => in_text = !in_text,
75            '{' if !in_text => depth += 1,
76            '}' if !in_text && depth > 0 => {
77                depth -= 1;
78                if depth == 0 {
79                    return true;
80                }
81            }
82            _ => {}
83        }
84    }
85    false
86}
87
88#[cfg(test)]
89mod tests {
90    use super::{Tool, json_closed, read_tool_call, tool_rules};
91
92    const TOOLS: &[Tool] = &[
93        Tool { name: "balance", description: "the balance" },
94        Tool { name: "spend", description: "adds a spending" },
95    ];
96
97    #[test]
98    fn it_reads_a_choice_and_its_arguments() {
99        let call = read_tool_call(r#"{"tool":"spend","args":{"what":"bread","amount":12.5}}"#, TOOLS);
100        let call = call.expect("a choice");
101        assert_eq!(call.name, "spend");
102        assert_eq!(call.argument("what"), Some("bread"));
103        assert_eq!(call.argument("amount"), Some("12.5"));
104    }
105
106    #[test]
107    fn it_finds_the_json_among_other_words() {
108        let call = read_tool_call("Sure:\n{\"tool\":\"balance\",\"args\":{}}\nDone", TOOLS);
109        assert_eq!(call.expect("a choice").name, "balance");
110    }
111
112    #[test]
113    fn it_sees_when_the_json_is_whole() {
114        assert!(!json_closed("Sure: {\"tool\":\"a\",\"args\":{}"));
115        assert!(json_closed("Sure: {\"tool\":\"a\",\"args\":{}}"));
116        assert!(!json_closed("{\"tool\":\"}\""));
117        assert!(!json_closed("no json"));
118    }
119
120    #[test]
121    fn a_name_with_accents_is_the_same_tool() {
122        let call = read_tool_call(r#"{"tool":"Spênd","args":{}}"#, TOOLS);
123        assert_eq!(call.expect("a choice").name, "spend");
124    }
125
126    #[test]
127    fn anything_else_is_no_choice() {
128        assert_eq!(read_tool_call(r#"{"tool":"delete_everything","args":{}}"#, TOOLS), None);
129        assert_eq!(read_tool_call("not json", TOOLS), None);
130        assert_eq!(read_tool_call(r#"{"tool":3}"#, TOOLS), None);
131        assert_eq!(read_tool_call("", TOOLS), None);
132    }
133
134    #[test]
135    fn a_blank_argument_is_no_argument() {
136        let call = read_tool_call(r#"{"tool":"spend","args":{"what":"  "}}"#, TOOLS).expect("a choice");
137        assert_eq!(call.argument("what"), None);
138    }
139
140    #[test]
141    fn the_rules_list_the_tools() {
142        let rules = tool_rules("Choose.", TOOLS);
143        assert!(rules.contains("- balance: the balance"));
144        assert!(rules.contains("\"tool\""));
145    }
146}
ravenrs_ai/tools.rs

ravenrs_ai/
tools.rs
