swiftide_agents/tools/
arg_preprocessor.rs

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

swiftide_agents/tools/
arg_preprocessor.rs
