Run code blocks

Cargo.toml

@@ -32,9 +32,11 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0.95"
 serde_yaml = "0.9"
 smartstring = { version = "1.0", features = ["serde"] }
+tempfile = "3.8"
 thiserror = "1.0"
 tokio = { version = "1", features = ["full"] }
 tokio-stream = "0.1.12"
+tokio-util = {version = "0.7", features = ["io"]}
 
 aio-cargo-info = { path = "./crates/aio-cargo-info", version = "0.1" }
 

src/arguments.rs

@@ -2,6 +2,35 @@ use std::fmt::Display;
 
 use clap::{Parser, ValueEnum};
 
+/// Program to communicate with large language models and AI API 
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub struct Args {
+    /// Configuration file
+    #[arg(long, default_value_t = String::from("~/.config/aio/config.yaml"))]
+    pub config_path: String,
+    /// Credentials file
+    #[arg(long, default_value_t = String::from("~/.config/aio/creds.yaml"))]
+    pub creds_path: String,
+    /// Engine name
+    /// 
+    /// The name can be followed by custom prompt name from the configuration file
+    /// (ex: openai:command)
+    #[arg(long, short)]
+    pub engine: String,
+    /// Formatter
+    /// 
+    /// Possible values: markdown, raw
+    #[arg(long, short, default_value_t = Default::default())]
+    pub formatter: FormatterChoice,
+    /// Run code block if the language is supported
+    #[arg(long, short, default_value_t = Default::default())]
+    pub run: RunChoice,
+    /// Force to run code 
+    /// User text prompt
+    pub input: Option<String>,
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
 pub enum FormatterChoice {
     Markdown,
@@ -28,36 +57,37 @@ impl Display for FormatterChoice {
     }
 }
 
-/// Program to communicate with large language models and AI API 
-#[derive(Parser, Debug)]
-#[command(author, version, about, long_about = None)]
-pub struct Args {
-    /// Configuration file
-    #[arg(long, default_value_t = String::from("~/.config/aio/config.yaml"))]
-    pub config_path: String,
-    /// Credentials file
-    #[arg(long, default_value_t = String::from("~/.config/aio/creds.yaml"))]
-    pub creds_path: String,
-    /// Engine name
-    /// 
-    /// The name can be followed by custom prompt name from the configuration file
-    /// (ex: openai:command)
-    #[arg(long, short)]
-    pub engine: String,
-    /// Formatter
-    /// 
-    /// Possible values: markdown, raw
-    #[arg(long, short, default_value_t = Default::default())]
-    pub formatter: FormatterChoice,
-    /// User text prompt
-    pub input: Option<String>,
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
+pub enum RunChoice {
+    /// Doesn't run anything
+    No,
+    /// Ask to run code
+    Ask,
+    /// Run code without asking
+    Force
+}
+
+impl Default for RunChoice {
+    fn default() -> Self {
+        RunChoice::No
+    }
 }
 
+impl Display for RunChoice {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            RunChoice::No => write!(f, "no"),
+            RunChoice::Ask => write!(f, "ask"),
+            RunChoice::Force => write!(f, "force"),
+        }
+    }
+}
 pub struct ProcessedArgs {
     pub config_path: String,
     pub creds_path: String,
     pub engine: String,
     pub formatter: FormatterChoice,
+    pub run: RunChoice,
     pub input: String,
 }
 
@@ -68,6 +98,7 @@ impl From<Args> for ProcessedArgs {
             creds_path: args.creds_path,
             engine: args.engine,
             formatter: args.formatter,
+            run: args.run,
             input: args.input.unwrap_or_default(),
         }
     }

src/generators/debug.rs

@@ -0,0 +1,14 @@
+use crate::args;
+use super::{ResultRun, ResultStream, Error};
+
+pub async fn run(config: crate::config::Config, args: args::ProcessedArgs) -> ResultRun {
+    use tokio_stream::StreamExt;
+    let input = args.input;
+    let file = tokio::fs::File::open(&input).await.map_err(|e| Error::Custom(std::borrow::Cow::Owned(e.to_string())))?;
+
+    let stream = tokio_util::io::ReaderStream::new(file).map(|mut r| -> ResultStream {
+        let bytes = r.map_err(|e| Error::Custom(std::borrow::Cow::Owned(e.to_string())))?;
+        Ok(String::from_utf8(bytes.as_ref().to_vec()).map_err(|e| Error::Custom(std::borrow::Cow::Owned(e.to_string())))?)
+    });
+    Ok(Box::pin(stream))
+}

src/generators/mod.rs

@@ -1,4 +1,5 @@
 pub mod openai;
+pub mod debug;
 
 use tokio_stream::Stream;
 use thiserror::Error;

src/main.rs

@@ -1,4 +1,5 @@
 pub mod arguments;
+mod runner;
 mod generators;
 mod formatters;
 mod config;
@@ -21,13 +22,21 @@ macro_rules! raise_str {
     };
 }
 
+fn home_dir() -> &'static str {
+    #[cfg(unix)]
+    lazy_static::lazy_static! {
+        static ref HOME: String = std::env::var("HOME").expect("Failed to resolve home path");
+    }
+    #[cfg(windows)]
+    lazy_static::lazy_static! {
+        static ref HOME: String = std::env::var("USERPROFILE").expect("Failed to resolve user profile path");
+    }
+    &*HOME
+}
+
 fn resolve_path(path: &str) -> Cow<str> {
     if path.starts_with("~/") {
-        #[cfg(unix)]
-        let home = std::env::var("HOME").expect("Failed to resolve home path");
-        #[cfg(windows)]
-        let home = std::env::var("USERPROFILE").expect("Failed to resolve user profile path");
-        Cow::Owned(format!("{}{}{}", home, std::path::MAIN_SEPARATOR, &path[2..]))
+        Cow::Owned(format!("{}{}{}", home_dir(), std::path::MAIN_SEPARATOR, &path[2..]))
     } else {
         Cow::Borrowed(path)
     }
@@ -67,23 +76,30 @@ async fn main() -> Result<(), String> {
         args::FormatterChoice::Markdown => Box::new(formatters::new_markdown_formatter()),
         args::FormatterChoice::Raw => Box::new(formatters::new_raw_formatter()),
     };
+    let mut runner = runner::Runner::new(args.run);
 
-    let engine = args.engine
+    let (engine, _prompt) = args.engine
         .find(':')
-        .map(|i| &args.engine[..i])
-        .unwrap_or(args.engine.as_str());
+        .map(|i| (&args.engine[..i], Some(&args.engine[i+1..])))
+        .unwrap_or((args.engine.as_str(), None));
+
     let mut stream = match engine {
         "openai" => generators::openai::run(creds.openai, config, args).await,
+        "from-file" => generators::debug::run(config, args).await,
         _ => panic!("Unknown engine: {}", engine),
     }.map_err(|e| format!("Failed to request OpenAI API: {}", e))?;
 
     loop {
         match stream.next().await {
-            Some(Ok(token)) => raise_str!(formatter.push(&token), "Failed to parse markdown: {}"),
+            Some(Ok(token)) => {
+                raise_str!(formatter.push(&token), "Failed to parse markdown: {}");
+                raise_str!(runner.push(&token), "Failed push text in the runner system: {}");
+            },
             Some(Err(e)) => Err(e.to_string())?,
             None => break,
         }
     }
     raise_str!(formatter.end_of_document(), "Failed to end markdown: {}");
+    raise_str!(runner.end_of_document(), "Failed to run code: {}");
     Ok(())
 }