1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
use std::error::Error;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

use crate::tools::Tool;

pub struct CommandExecutor {
    platform: String,
}

impl CommandExecutor {
    /// Create a new CommandExecutor instance
    /// # Example
    /// ```rust,ignore
    /// let tool = CommandExecutor::new("linux");
    /// ```
    pub fn new<S: Into<String>>(platform: S) -> Self {
        Self {
            platform: platform.into(),
        }
    }
}

impl Default for CommandExecutor {
    fn default() -> Self {
        Self::new("linux")
    }
}

#[derive(Deserialize, Serialize, Debug)]
struct CommandInput {
    cmd: String,
    #[serde(default)]
    args: Vec<String>,
}
#[derive(Serialize, Deserialize, Debug)]
struct CommandsWrapper {
    commands: Vec<CommandInput>,
}

#[async_trait]
impl Tool for CommandExecutor {
    fn name(&self) -> String {
        String::from("Command_Executor")
    }
    fn description(&self) -> String {
        String::from(format!(
            r#""This tool let you run command on the terminal"
            "The input should be an array with comands for the following platform: {}"
            "examle of input: [{{ "cmd": "ls", "args": [] }},{{"cmd":"mkdir","args":["test"]}}]"
            "Should be a comma separeted comands"
            "#,
            self.platform
        ))
    }

    fn parameters(&self) -> Value {
        let prompt = format!(
            "This tool let you run command on the terminal.
        The input should be an array with comands for the following platform: {}",
            self.platform
        );
        json!(

        {
          "description": prompt,
          "type": "object",
          "properties": {
            "commands": {
              "description": "An array of command objects to be executed",
              "type": "array",
              "items": {
                "type": "object",
                "properties": {
                  "cmd": {
                    "type": "string",
                    "description": "The command to execute"
                  },
                  "args": {
                    "type": "array",
                    "items": {
                      "type": "string"
                    },
                    "default": [],
                    "description": "List of arguments for the command"
                  }
                },
                "required": ["cmd"],
                "additionalProperties": false,
                "description": "Object representing a command and its optional arguments"
              }
            }
          },
          "required": ["commands"],
          "additionalProperties": false
        }
                )
    }

    async fn parse_input(&self, input: &str) -> Value {
        log::info!("Parsing input: {}", input);

        // Attempt to parse input string into CommandsWrapper struct first
        let wrapper_result = serde_json::from_str::<CommandsWrapper>(input);

        if let Ok(wrapper) = wrapper_result {
            // If successful, serialize the `commands` back into a serde_json::Value
            // this is for llm like open ai tools
            serde_json::to_value(wrapper.commands).unwrap_or_else(|err| {
                log::error!("Serialization error: {}", err);
                Value::Null
            })
        } else {
            // If the first attempt fails, try parsing it as Vec<CommandInput> directly
            // This works on any llm
            let commands_result = serde_json::from_str::<Vec<CommandInput>>(input);

            commands_result.map_or_else(
                |err| {
                    log::error!("Failed to parse input: {}", err);
                    Value::Null
                },
                |commands| serde_json::to_value(commands).unwrap_or(Value::Null),
            )
        }
    }

    async fn run(&self, input: Value) -> Result<String, Box<dyn Error>> {
        let commands: Vec<CommandInput> = serde_json::from_value(input)?;
        let mut result = String::new();

        for command in commands {
            let mut command_to_execute = std::process::Command::new(&command.cmd);
            command_to_execute.args(&command.args);

            let output = command_to_execute.output()?;

            result.push_str(&format!(
                "Command: {}\nOutput: {}",
                command.cmd,
                String::from_utf8_lossy(&output.stdout),
            ));

            if !output.status.success() {
                return Err(Box::new(std::io::Error::new(
                    std::io::ErrorKind::Other,
                    format!(
                        "Command {} failed with status: {}",
                        command.cmd, output.status
                    ),
                )));
            }
        }

        Ok(result)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use serde_json::json;
    #[tokio::test]
    async fn test_with_string_executor() {
        let tool = CommandExecutor::new("linux");
        let input = json!({
            "commands": [
                {
                    "cmd": "ls",
                    "args": []
                }
            ]
        });
        println!("{}", &input.to_string());
        let result = tool.call(&input.to_string()).await.unwrap();
        println!("Res: {}", result);
    }
}