From d201294b48cf1e8981615c1a33ba53191ea931ac Mon Sep 17 00:00:00 2001 From: ykdy3951 Date: Wed, 28 Feb 2024 18:50:40 +0900 Subject: [PATCH 1/3] Add set key & search commands to the CLI - Add set key command to the CLI - Add search command to the CLI - Add clipboard copy feature to the CLI - Delete "Assistant: " prefix from the output of the CLI - Update help message of the CLI - Add spinner display to the CLI when searching --- Cargo.toml | 10 +++-- src/commands/search/mod.rs | 86 ++++++++++++++++++++++++++++++++++++- src/commands/set/key/mod.rs | 28 +++++++++++- src/main.rs | 24 ++++++++++- 4 files changed, 140 insertions(+), 8 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 02b30d0..cf9ef45 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "cllm" -version = "0.1.0" +version = "0.1.1" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -8,5 +8,9 @@ edition = "2021" [dependencies] clap = { version = "4.5.1", features = ["derive"] } tokio = { version = "1.36.0", features = ["full"] } -llm-chain = "0.12.0" -llm-chain-openai = "0.12.0" \ No newline at end of file +llm-chain = "0.13.0" +llm-chain-openai = "0.13.0" +serde_json = "1.0" +dirs = "5.0" +spinners = "4.1.1" +cli-clipboard = "0.4.0" diff --git a/src/commands/search/mod.rs b/src/commands/search/mod.rs index c0c6c95..20e7979 100644 --- a/src/commands/search/mod.rs +++ b/src/commands/search/mod.rs @@ -1,5 +1,13 @@ use clap::Parser; - +use std::env; +use spinners::{Spinner, Spinners}; +use llm_chain::{ + chains::conversation::Chain, executor, parameters, prompt, step::Step, + options::{ModelRef, OptionsBuilder}, + prompt::{Conversation, ChatMessageCollection}, +}; +use llm_chain_openai::chatgpt::Model; +use cli_clipboard::{ClipboardContext, ClipboardProvider}; #[derive(Debug, Parser)] #[clap( name = "search", @@ -11,7 +19,81 @@ pub struct Search { qeury: String, } +pub fn few_shot_template(list: Vec<(String, String)>) -> ChatMessageCollection { + + let mut ret_prompt = Conversation::new(); + + for (user, assistant) in &list { + ret_prompt = ret_prompt.with_user(user.to_string()).with_assistant(assistant.to_string()); + } + + ret_prompt +} + pub async fn handle_search(search: Search) -> Result<(), Box> { - println!("Searching for: {}", search.qeury); + + if !env::var("OPENAI_API_KEY").is_ok() { + println!("Please set your OpenAI API key using the `set key` command."); + return Ok(()); + } + + let mut spinner = Spinner::new(Spinners::Dots9, "Searching for the command...".into()); + + let model = ModelRef::from_model_name(Model::Gpt35Turbo.to_string()); + + let mut option_builder = OptionsBuilder::new(); + option_builder.add_option(llm_chain::options::Opt::Model(model)); + let options = option_builder.build(); + + let exec = executor!( + chatgpt, + options + )?; + + let few_shot_examples: Vec<(String, String)> = vec![ + ("Show all pods in k8s".to_string(), "kubectl get pods".to_string()), + ("Find all files recursively within the current directory that contain 'a' in their filenames.".to_string(), "find . -type f -name '*a*' -print".to_string()), + ("Provide the command to build and push a Docker image from the current directory.".to_string(), "docker build -t myapp:latest --path".to_string()), + ]; + + + let mut conversation = Conversation::new() + .with_system_template( + "I want you to act as generating a command for request tasks on {{os_name}}. Also please don't explain the commands, just generate the command.", + ¶meters!{"os_name" => env::consts::OS} + ).unwrap(); + + // i want to append the few shot examples to the conversation + let few_shot_prompt = few_shot_template(few_shot_examples); + + conversation.append(few_shot_prompt); + + let conversation = conversation + .with_system( + " + Only generate the command, don't explain it. + ".to_string() + ); + + let mut chain = Chain::new_with_message_collection( + &conversation + ); + + let step = Step::for_prompt_template( + prompt!( + user: "task : {{query}}" + ) + ); + let parameters = parameters!().with("query", search.qeury); + let res = chain.send_message(step, ¶meters, &exec).await?; + let res = res.to_immediate().await?.as_content().to_chat().to_string(); + let res = res.split("Assistant: ").collect::>()[1].to_string().trim().to_string(); + + let mut ctx: ClipboardContext = ClipboardProvider::new().unwrap(); + ctx.set_contents(res.clone().to_string()).unwrap(); + + spinner.stop_and_persist("✔", "Finished searching for the command and copied to your clipboard :)".into()); + + println!("{}", res); Ok(()) } \ No newline at end of file diff --git a/src/commands/set/key/mod.rs b/src/commands/set/key/mod.rs index 4652586..9debb76 100644 --- a/src/commands/set/key/mod.rs +++ b/src/commands/set/key/mod.rs @@ -1,4 +1,8 @@ use clap::Parser; +use dirs; +use std::env; +use std::fs::File; +use std::io::prelude::*; #[derive(Debug, Parser)] #[clap( @@ -12,6 +16,28 @@ pub struct Key { } pub async fn handle_key(key: Key) -> Result<(), Box> { - println!("Setting API Key: {}", key.api_key); + + let home_dir = dirs::home_dir().unwrap(); + let save_dir = home_dir.join(".cllm"); + let config_path = save_dir.join("credentials.json"); + + if !save_dir.exists() { + std::fs::create_dir_all(&save_dir)?; + } + + let mut config = if config_path.exists() { + let config = std::fs::read_to_string(config_path.clone())?; + serde_json::from_str(&config)? + } + else { + serde_json::json!({}) + }; + + config["OPEN_AI"] = key.api_key.clone().into(); + let config = serde_json::to_string_pretty(&config)?; + File::create(config_path)?.write_all(config.as_bytes())?; + env::set_var("OPENAI_API_KEY", key.api_key); + + println!("API key set successfully."); Ok(()) } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 670fafd..60c4bea 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,21 +1,41 @@ pub mod commands; +use std::env; use clap::Parser; use commands::{Commands, handle_command}; +use dirs; #[derive(Debug, Parser)] #[clap( version, - author, + about="Empower your CLI experience with a command search tool driven by LLM magic!\n\ + Github: https://github.com/dev-backpack/cllm\n\ + If you have any questions or suggestions, feel free to open an issue on the github repo." )] struct Cli { #[clap(subcommand)] pub commands: Commands, } - #[tokio::main] async fn main() { + + // Set the OPENAI_API_KEY environment variable + let home_dir = dirs::home_dir().unwrap(); + let save_dir = home_dir.join(".cllm"); + let config_path = save_dir.join("credentials.json"); + + if config_path.exists() { + let config = std::fs::read_to_string(config_path).unwrap(); + let config: serde_json::Value = serde_json::from_str(&config).unwrap(); + + if config["OPEN_AI"].is_string() { + let api_key = config["OPEN_AI"].as_str().unwrap(); + env::set_var("OPENAI_API_KEY", api_key); + } + } + + // Parse the command line arguments let cli = Cli::parse(); if let Err(_error) = handle_command(cli.commands).await { From 49b7073439029cc098c583e326dbfe116df311f3 Mon Sep 17 00:00:00 2001 From: ykdy3951 Date: Wed, 28 Feb 2024 19:03:53 +0900 Subject: [PATCH 2/3] Update lib.rs --- src/lib.rs | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/lib.rs b/src/lib.rs index bc8a79c..412ef34 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,12 +1,16 @@ pub(crate) mod commands; +use std::env; use clap::Parser; use commands::{Commands, handle_command}; +use dirs; #[derive(Debug, Parser)] #[clap( version, - author, + about="Empower your CLI experience with a command search tool driven by LLM magic!\n\ + Github: https://github.com/dev-backpack/cllm\n\ + If you have any questions or suggestions, feel free to open an issue on the github repo." )] struct Cli { #[clap(subcommand)] @@ -14,6 +18,22 @@ struct Cli { } pub async fn run() -> Result<(), Box> { + + // Set the OPENAI_API_KEY environment variable + let home_dir = dirs::home_dir().unwrap(); + let save_dir = home_dir.join(".cllm"); + let config_path = save_dir.join("credentials.json"); + + if config_path.exists() { + let config = std::fs::read_to_string(config_path).unwrap(); + let config: serde_json::Value = serde_json::from_str(&config).unwrap(); + + if config["OPEN_AI"].is_string() { + let api_key = config["OPEN_AI"].as_str().unwrap(); + env::set_var("OPENAI_API_KEY", api_key); + } + } + let cli: Cli = Cli::parse(); if let Err(_error) = handle_command(cli.commands).await { From 06dcd98338dda890ab9167f4b7cd318f6fd5a717 Mon Sep 17 00:00:00 2001 From: ykdy3951 Date: Wed, 28 Feb 2024 21:35:36 +0900 Subject: [PATCH 3/3] Update search command --- src/commands/search/mod.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/commands/search/mod.rs b/src/commands/search/mod.rs index 20e7979..b41c362 100644 --- a/src/commands/search/mod.rs +++ b/src/commands/search/mod.rs @@ -63,16 +63,13 @@ pub async fn handle_search(search: Search) -> Result<(), Box env::consts::OS} ).unwrap(); - // i want to append the few shot examples to the conversation let few_shot_prompt = few_shot_template(few_shot_examples); conversation.append(few_shot_prompt); let conversation = conversation .with_system( - " - Only generate the command, don't explain it. - ".to_string() + "Only generate the command, don't explain it".to_string() ); let mut chain = Chain::new_with_message_collection(