feat: add -c/--code to generate only code (#327)

Author: sigoden

README.md

@@ -311,10 +311,11 @@ Options:
   -r, --role <ROLE>          Choose a role
   -s, --session [<SESSION>]  Create or reuse a session
   -e, --execute              Execute commands using natural language
+  -c, --code                 Generate only code
   -f, --file <FILE>...       Attach files to the message to be sent
   -H, --no-highlight         Disable syntax highlighting
   -S, --no-stream            No stream output
-  -w, --wrap <WRAP>          Specify the text-wrapping mode (no*, auto, <max-width>)
+  -w, --wrap <WRAP>          Specify the text-wrapping mode (no, auto, <max-width>)
       --light-theme          Use light theme
       --dry-run              Run in dry run mode
       --info                 Print related information
@@ -328,18 +329,16 @@ Options:
 Here are some practical examples:
 
 ```sh
-aichat -s                                    # Start REPL with a new temp session
+aichat -s                                    # Start REPL with a new session
 aichat -s temp                               # Reuse temp session
 aichat -r shell -s                           # Create a session with a role
 aichat -m openai:gpt-4-32k -s                # Create a session with a model
-aichat -s sh unzip a file                    # Run session in command mode
+aichat -s temp unzip a file                  # Run session in command mode
 
 aichat -r shell unzip a file                 # Use role in command mode
 aichat -s shell unzip a file                 # Use session in command mode
 
-cat config.json | aichat convert to yaml     # Read stdin
-cat config.json | aichat -r convert:yaml     # Read stdin with a role
-cat config.json | aichat -s i18n             # Read stdin with a session
+cat data.toml | aichat -c to json            # Read stdin
 
 aichat --file a.png b.png -- diff images     # Attach files
 aichat --file screenshot.png -r ocr          # Attach files with a role
@@ -352,7 +351,7 @@ aichat --info                                # system-wide information
 aichat -s temp --info                        # Show session details
 aichat -r shell --info                       # Show role info
 
-$(echo "$data" | aichat -S -H to json)       # Use aichat in a script
+$(echo "$data" | aichat -c to json)          # Use aichat in a script
 ```
 
 ### Execute commands using natural language
@@ -399,7 +398,7 @@ This is a **very handy feature**, which allows you to use `aichat` shell complet
 
 ![aichat-integration](https://github.com/sigoden/aichat/assets/4012553/873ebf23-226c-412e-a34f-c5aaa7017524)
 
-To install shell integration, go to [./scripts/shell-integration](./scripts/shell-integration/) to download the script and source the script in rc file. After that restart your shell. You can invoke the completion with `alt+e` hotkey.
+To install shell integration, go to [./scripts/shell-integration](https://github.com/sigoden/aichat/tree/main/scripts/shell-integration) to download the script and source the script in rc file. After that restart your shell. You can invoke the completion with `alt+e` hotkey.
 
 ## License
 

src/cli.rs

@@ -15,6 +15,9 @@ pub struct Cli {
     /// Execute commands using natural language
     #[clap(short = 'e', long)]
     pub execute: bool,
+    /// Generate only code
+    #[clap(short = 'c', long)]
+    pub code: bool,
     /// Attach files to the message to be sent.
     #[clap(short = 'f', long, num_args = 1.., value_name = "FILE")]
     pub file: Option<Vec<String>>,

src/config/mod.rs

@@ -288,6 +288,11 @@ impl Config {
         self.set_role_obj(role)
     }
 
+    pub fn set_code_role(&mut self) -> Result<()> {
+        let role = Role::for_code();
+        self.set_role_obj(role)
+    }
+
     pub fn set_role_obj(&mut self, role: Role) -> Result<()> {
         if let Some(session) = self.session.as_mut() {
             session.update_role(Some(role.clone()))?;

src/config/role.rs

@@ -29,7 +29,7 @@ impl Role {
             _ => "&&",
         };
         Self {
-            name: "__for_execute__".into(),
+            name: "__execute__".into(),
             prompt: format!(
                 r#"Provide only {shell} commands for {os} without any description.
 If there is a lack of details, provide most logical solution.
@@ -44,7 +44,7 @@ Do not provide markdown formatting such as ```"#
 
     pub fn for_describe() -> Self {
         Self {
-            name: "__for_describe__".into(),
+            name: "__describe__".into(),
             prompt: r#"Provide a terse, single sentence description of the given shell command.
 Describe each argument and option of the command.
 Provide short responses in about 80 words.
@@ -54,6 +54,20 @@ APPLY MARKDOWN formatting when possible."#
         }
     }
 
+    pub fn for_code() -> Self {
+        Self {
+            name: "__code__".into(),
+            prompt: r#"Provide only code as output without any description.
+Provide only code in plain text format without Markdown formatting.
+Do not include symbols such as ``` or ```python.
+If there is a lack of details, provide most logical solution.
+You are not allowed to ask for more details.
+For example if the prompt is "Hello world Python", you should return "print('Hello world')"."#
+                .into(),
+            temperature: None,
+        }
+    }
+
     pub fn info(&self) -> Result<String> {
         let output = serde_yaml::to_string(&self)
             .with_context(|| format!("Unable to show info about role {}", &self.name))?;

src/main.rs

@@ -11,7 +11,7 @@ mod utils;
 
 use crate::cli::Cli;
 use crate::config::{Config, GlobalConfig};
-use crate::utils::run_command;
+use crate::utils::{extract_block, run_command};
 
 use anyhow::{bail, Result};
 use clap::Parser;
@@ -65,6 +65,8 @@ fn main() -> Result<()> {
     } else {
         if let Some(name) = &cli.role {
             config.write().set_role(name)?;
+        } else if cli.code {
+            config.write().set_code_role()?;
         }
         if let Some(session) = &cli.session {
             config
@@ -95,7 +97,7 @@ fn main() -> Result<()> {
     }
     config.write().prelude()?;
     if let Err(err) = match text {
-        Some(text) => start_directive(&config, &text, cli.file, cli.no_stream),
+        Some(text) => start_directive(&config, &text, cli.file, cli.no_stream, cli.code),
         None => start_interactive(&config),
     } {
         let highlight = stderr().is_terminal() && config.read().highlight;
@@ -109,6 +111,7 @@ fn start_directive(
     text: &str,
     include: Option<Vec<String>>,
     no_stream: bool,
+    code_mode: bool,
 ) -> Result<()> {
     if let Some(session) = &config.read().session {
         session.guard_save()?;
@@ -117,14 +120,19 @@ fn start_directive(
     let mut client = init_client(config)?;
     ensure_model_capabilities(client.as_mut(), input.required_capabilities())?;
     config.read().maybe_print_send_tokens(&input);
-    let output = if no_stream {
+    let output = if !stdout().is_terminal() || no_stream {
         let output = client.send_message(input.clone())?;
-        if stdout().is_terminal() {
+        let to_print = if code_mode && output.trim_start().starts_with("```") {
+            extract_block(&output)
+        } else {
+            output.clone()
+        };
+        if no_stream {
             let render_options = config.read().get_render_options()?;
             let mut markdown_render = MarkdownRender::init(render_options)?;
-            println!("{}", markdown_render.render(&output).trim());
+            println!("{}", markdown_render.render(&to_print).trim());
         } else {
-            println!("{}", output);
+            println!("{}", to_print);
         }
         output
     } else {
@@ -144,7 +152,10 @@ fn execute(config: &GlobalConfig, text: &str) -> Result<()> {
     let input = Input::from_str(text);
     let client = init_client(config)?;
     config.read().maybe_print_send_tokens(&input);
-    let eval_str = client.send_message(input.clone())?;
+    let mut eval_str = client.send_message(input.clone())?;
+    if eval_str.contains("```") {
+        eval_str = extract_block(&eval_str);
+    }
     config.write().save_message(input, &eval_str)?;
     let render_options = config.read().get_render_options()?;
     let mut markdown_render = MarkdownRender::init(render_options)?;

src/utils/mod.rs

@@ -10,10 +10,16 @@ pub use self::prompt_input::*;
 pub use self::render_prompt::render_prompt;
 pub use self::tiktoken::cl100k_base_singleton;
 
+use fancy_regex::Regex;
+use lazy_static::lazy_static;
 use sha2::{Digest, Sha256};
 use std::env;
 use std::process::Command;
 
+lazy_static! {
+    static ref CODE_BLOCK_RE: Regex = Regex::new(r"(?ms)```\w*(.*?)```").unwrap();
+}
+
 pub fn now() -> String {
     let now = chrono::Local::now();
     now.to_rfc3339_opts(chrono::SecondsFormat::Secs, false)
@@ -150,6 +156,18 @@ pub fn run_command(eval_str: &str) -> anyhow::Result<i32> {
     Ok(status.code().unwrap_or_default())
 }
 
+pub fn extract_block(input: &str) -> String {
+    let output: String = CODE_BLOCK_RE
+        .captures_iter(input)
+        .filter_map(|m| {
+            m.ok()
+                .and_then(|cap| cap.get(1))
+                .map(|m| String::from(m.as_str()))
+        })
+        .collect();
+    output.trim().to_string()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;