shell_gpt/
lib.rs

1pub mod config;
2pub mod encryption;
3pub mod errors;
4pub mod openai;
5
6use crate::config::{CliArgs, Config};
7use rand::Rng;
8use spinoff::{spinners, Color, Spinner};
9use std::{io, process};
10
11const LOADING_MESSAGES: [&'static str; 7] = [
12    "Asking ChatGPT...",
13    "ChatGPT is thinking...",
14    "ChatGPT is writing...",
15    "Assistant is thinking...",
16    "Assistant is writing...",
17    "Invoking the AI gods...",
18    "Calling an AI friend...",
19];
20
21pub fn run(cli_args: &CliArgs, config: &Config) -> anyhow::Result<String> {
22    if !cli_args.raw {
23        let spinner = Spinner::new(spinners::Dots, get_loading_message(), Color::Blue);
24        let response = openai::ask_chatgpt(
25            cli_args.input.as_str(),
26            cli_args.pre_prompt,
27            config.openai_api_key.as_str(),
28        );
29        spinner.stop();
30        let response = response?;
31        println!("\n{response}");
32
33        if cli_args.pre_prompt == openai::PrePrompt::ShellScript {
34            ask_then_run_script(response.as_str());
35        }
36        Ok(response)
37    } else {
38        let response = openai::ask_chatgpt(
39            cli_args.input.as_str(),
40            cli_args.pre_prompt,
41            config.openai_api_key.as_str(),
42        );
43        let mut response = response?;
44        if cli_args.pre_prompt == openai::PrePrompt::ShellScript {
45            response = openai::extract_code_block_if_needed(response.as_str())
46        }
47        println!("\n{response}");
48        Ok(response)
49    }
50}
51
52fn get_loading_message() -> String {
53    let i = rand::thread_rng().gen_range(0..LOADING_MESSAGES.len());
54    LOADING_MESSAGES[i].to_string()
55}
56
57fn ask_then_run_script(response: &str) {
58    println!("\nDo you want to run this script? (y/N)");
59
60    let mut yes_no = String::new();
61    io::stdin()
62        .read_line(&mut yes_no)
63        .expect("Failed to read your input");
64    yes_no = yes_no.trim().to_lowercase();
65    if yes_no == "y" || yes_no == "yes" || yes_no == "o" || yes_no == "oui" {
66        let script = openai::extract_code_block_if_needed(response);
67        // let script = format!("set -x;{script}");
68        println!("\nExecuting script...\n\n--------------\n\n");
69        let output = process::Command::new("bash")
70            .arg("-c")
71            .arg(script)
72            .output()
73            .expect("Failed to execute script");
74        if !output.stderr.is_empty() {
75            eprintln!(
76                "{}",
77                String::from_utf8(output.stderr)
78                    .expect("Could not parse script stderr output as UTF-8")
79            );
80        } else {
81            println!(
82                "{}",
83                String::from_utf8(output.stdout)
84                    .expect("Could not parse script stdout output as UTF-8")
85            );
86        }
87        if !output.status.success() {
88            process::exit(1);
89        }
90    }
91}
shell_gpt/lib.rs

shell_gpt/
lib.rs
