swiftide_agents/tools/
arg_preprocessor.rs

1use std::borrow::Cow;
2
3use serde_json::{Map, Value};
4use swiftide_core::chat_completion::ToolCall;
5
6/// Preprocesses arguments for tool calls and tries to fix common errors
7/// This must be infallible and the result is always forwarded to the tool
8pub struct ArgPreprocessor;
9
10impl ArgPreprocessor {
11    pub fn preprocess_tool_calls(tool_calls: &mut [ToolCall]) {
12        for tool_call in tool_calls.iter_mut() {
13            let args = Self::preprocess(tool_call.args());
14
15            if args.as_ref().is_some_and(|a| match a {
16                Cow::Borrowed(_) => false,
17                Cow::Owned(_) => true,
18            }) {
19                tool_call.with_args(args.map(|a| a.to_string()));
20            }
21        }
22    }
23    pub fn preprocess(value: Option<&str>) -> Option<Cow<'_, str>> {
24        Some(take_first_occurrence_in_object(value?))
25    }
26}
27
28/// Strips duplicate keys from JSON objects
29fn take_first_occurrence_in_object(value: &str) -> Cow<'_, str> {
30    let Ok(parsed) = &serde_json::from_str(value) else {
31        return Cow::Borrowed(value);
32    };
33    if let Value::Object(obj) = parsed {
34        let mut new_map = Map::with_capacity(obj.len());
35        for (k, v) in obj {
36            // Only insert if we haven't seen this key yet.
37            new_map.entry(k).or_insert(v.clone());
38        }
39        Cow::Owned(Value::Object(new_map).to_string())
40    } else {
41        // If the top-level isn't even an object, just pass it as is,
42        // or decide how you want to handle that situation.
43        Cow::Borrowed(value)
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use serde_json::json;
51
52    #[test]
53    fn test_preprocess_regular_json() {
54        let input = json!({
55            "key1": "value1",
56            "key2": "value2"
57        })
58        .to_string();
59        let expected = json!({
60            "key1": "value1",
61            "key2": "value2"
62        });
63        let result = ArgPreprocessor::preprocess(Some(&input));
64        assert_eq!(result.as_deref(), Some(expected.to_string().as_str()));
65    }
66
67    #[test]
68    fn test_preprocess_json_with_duplicate_keys() {
69        let input = json!({
70            "key1": "value1",
71            "key1": "value2"
72        })
73        .to_string();
74        let expected = json!({
75            "key1": "value2"
76        });
77        let result = ArgPreprocessor::preprocess(Some(&input));
78        assert_eq!(result.as_deref(), Some(expected.to_string().as_str()));
79    }
80
81    #[test]
82    fn test_no_preprocess_invalid_json() {
83        let input = "invalid json";
84        let result = ArgPreprocessor::preprocess(Some(input));
85        assert_eq!(result.as_deref(), Some(input));
86    }
87
88    #[test]
89    fn test_no_input() {
90        let result = ArgPreprocessor::preprocess(None);
91        assert_eq!(result, None);
92    }
93}
swiftide_agents/tools/arg_preprocessor.rs

swiftide_agents/tools/
arg_preprocessor.rs
