feat: compile bee to pdl

Signed-off-by: Nick Mitchell <[email protected]>

This also improves error handling in the rust code.

Signed-off-by: Nick Mitchell <[email protected]>

Author: starpit

.github/workflows/tauri-cli.yml

@@ -42,8 +42,14 @@ jobs:
         run: |
           PATH=./src-tauri/target/release/:$PATH
           
-          # 1. `run` subcommand works without any arguments
-          pdl run | grep Usage
+          # 1a. `run` subcommand errors due to missing required positional parameter
+          pdl run && (echo "This should have failed" && exit 1) || (echo "Great, expected failure received" && exit 0)
+
+          # 1b.`run` subcommand works without any arguments to print Usage
+          pdl run 2>&1 | grep Usage
+
+          # 1c.`run` subcommand works with -h to print Usage
+          pdl run -h 2>&1 | grep Usage
 
           # 2. `run` subcommand works with UI demos (yaml source)
           pdl run ./demos/demo1.pdl | grep 'write a hello'

pdl-live-react/src-tauri/src/cli/run.rs

@@ -13,7 +13,7 @@ pub fn run_pdl_program(
     trace_file: Option<&tauri_plugin_cli::ArgData>,
     data: Option<&tauri_plugin_cli::ArgData>,
     stream: Option<&tauri_plugin_cli::ArgData>,
-) -> Result<(), tauri::Error> {
+) -> Result<(), Box<dyn ::std::error::Error>> {
     println!(
         "Running {:#?}",
         Path::new(&source_file_path).file_name().unwrap()

pdl-live-react/src-tauri/src/cli/setup.rs

@@ -1,77 +1,86 @@
-use std::path::Path;
-use std::process::exit;
+use ::std::path::Path;
 
 use serde_json::Value;
 use urlencoding::encode;
 
 use tauri_plugin_cli::CliExt;
 
 use crate::cli::run;
+use crate::compile;
 use crate::gui::setup as gui_setup;
 
 #[cfg(desktop)]
-pub fn cli(app: &mut tauri::App) -> Result<(), tauri::Error> {
+pub fn cli(app: &mut tauri::App) -> Result<(), Box<dyn ::std::error::Error>> {
     app.handle().plugin(tauri_plugin_cli::init())?;
 
-    match app.cli().matches() {
-        // `matches` here is a Struct with { args, subcommand }.
-        // `args` is `HashMap<String, ArgData>` where `ArgData` is a struct with { value, occurrences }.
-        // `subcommand` is `Option<Box<SubcommandMatches>>` where `SubcommandMatches` is a struct with { name, matches }.
-        Ok(matches) => match matches.subcommand {
-            Some(subcommand_matches) => match subcommand_matches.name.as_str() {
-                "run" => {
-                    if let Some(source) = subcommand_matches.matches.args.get("source") {
-                        if let Value::String(source_file_path) = &source.value {
-                            match run::run_pdl_program(
-                                source_file_path.clone(),
-                                app.handle().clone(),
-                                subcommand_matches.matches.args.get("trace"),
-                                subcommand_matches.matches.args.get("data"),
-                                subcommand_matches.matches.args.get("stream"),
-                            ) {
-                                Ok(()) => exit(0),
-                                _ => exit(1),
-                            }
-                        }
-                    }
-                    println!("Usage: run <source.pdl>");
-                    exit(1)
-                }
-                "view" => match subcommand_matches.matches.args.get("trace") {
-                    Some(trace) => match &trace.value {
-                        Value::String(trace_file) => {
-                            gui_setup(
-                                app.handle().clone(),
-                                Path::new("/local")
-                                    .join(encode(trace_file).as_ref())
-                                    .display()
-                                    .to_string(),
-                            )?;
-                            Ok(())
-                        }
-                        _ => {
-                            println!("Usage: view <tracefile.json>");
-                            exit(1)
-                        }
-                    },
-                    _ => {
-                        println!("Usage: view <tracefile.json>");
-                        exit(1)
-                    }
-                },
-                _ => {
-                    println!("Invalid subcommand");
-                    exit(1)
+    // `matches` here is a Struct with { args, subcommand }.
+    // `args` is `HashMap<String, ArgData>` where `ArgData` is a struct with { value, occurrences }.
+    // `subcommand` is `Option<Box<SubcommandMatches>>` where `SubcommandMatches` is a struct with { name, matches }.
+    let Some(subcommand_matches) = app.cli().matches()?.subcommand else {
+        if let Some(help) = app.cli().matches()?.args.get("help") {
+            return Err(Box::from(help.value.as_str().or(Some("Internal Error")).unwrap()));
+        } else {
+            return Err(Box::from("Internal Error"));
+        }
+    };
+
+    match subcommand_matches.name.as_str() {
+        "compile" => {
+            let Some(compile_subcommand_matches) = subcommand_matches.matches.subcommand else {
+                return Err(Box::from("Missing compile subcommand"));
+            };
+
+            match compile_subcommand_matches.name.as_str() {
+                "beeai" => {
+                    let Some(source) = compile_subcommand_matches.matches.args.get("source") else {
+                        return Err(Box::from("Missing source file"));
+                    };
+                    let Value::String(source_file_path) = &source.value else {
+                        return Err(Box::from("Invalid source file argument"));
+                    };
+                    let Some(output) = compile_subcommand_matches.matches.args.get("output") else {
+                        return Err(Box::from("Missing output argument"));
+                    };
+                    let Value::String(output_file_path) = &output.value else {
+                        return Err(Box::from("Invalid output file argument"));
+                    };
+                    return compile::beeai::compile(source_file_path, output_file_path);
                 }
-            },
-            None => {
-                println!("Invalid command");
-                exit(1)
+                _ => {}
             }
-        },
-        Err(s) => {
-            println!("{:?}", s);
-            exit(1)
         }
+        "run" => {
+            let Some(source) = subcommand_matches.matches.args.get("source") else {
+                return Err(Box::from("Missing source file"));
+            };
+            let Value::String(source_file_path) = &source.value else {
+                return Err(Box::from("Invalid source file argument"));
+            };
+            return run::run_pdl_program(
+                source_file_path.clone(),
+                app.handle().clone(),
+                subcommand_matches.matches.args.get("trace"),
+                subcommand_matches.matches.args.get("data"),
+                subcommand_matches.matches.args.get("stream"),
+            );
+        }
+        "view" => {
+            let Some(trace) = subcommand_matches.matches.args.get("trace") else {
+                return Err(Box::from("Missing trace file"));
+            };
+            let Value::String(trace_file) = &trace.value else {
+                return Err(Box::from("Invalid trace file argument"));
+            };
+            gui_setup(
+                app.handle().clone(),
+                Path::new("/local")
+                    .join(encode(trace_file).as_ref())
+                    .display()
+                    .to_string(),
+            )?
+        }
+        _ => {}
     }
+
+    Ok(())
 }

pdl-live-react/src-tauri/src/compile/beeai.rs

@@ -0,0 +1,186 @@
+use ::std::collections::HashMap;
+use ::std::error::Error;
+use ::std::fs::File;
+use ::std::io::BufReader;
+
+use serde::{Deserialize, Serialize};
+use serde_json::{from_reader, to_string, Value};
+
+macro_rules! zip {
+    ($x: expr) => ($x);
+    ($x: expr, $($y: expr), +) => (
+        $x.into_iter().zip(
+            zip!($($y), +))
+    )
+}
+
+#[derive(Deserialize, Debug)]
+struct BeeAiInputStateDict {
+    prompt: Option<String>,
+    // expected_output: Option<String>,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiInputState {
+    #[serde(rename = "__dict__")]
+    dict: BeeAiInputStateDict,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiInput {
+    #[serde(rename = "py/state")]
+    state: BeeAiInputState,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiTool {
+    //#[serde(rename = "py/object")]
+    //tool: String,
+    //options: Option<String>, // TODO maybe more general than String?
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiLlmParametersState {
+    #[serde(rename = "__dict__")]
+    dict: HashMap<String, Value>,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiLlmParameters {
+    #[serde(rename = "py/state")]
+    state: BeeAiLlmParametersState,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiLlmSettings {
+    api_key: String,
+    // base_url: String,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiLlm {
+    // might be helpful to know it's Ollama?
+    //#[serde(rename = "py/object")]
+    //object: String,
+    parameters: BeeAiLlmParameters,
+
+    #[serde(rename = "_model_id")]
+    model_id: String,
+    //#[serde(rename = "_litellm_provider_id")]
+    //provider_id: String,
+    #[serde(rename = "_settings")]
+    settings: BeeAiLlmSettings,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflowStepStateMeta {
+    //name: String,
+    role: String,
+    llm: BeeAiLlm,
+    instructions: Option<String>,
+    //tools: Option<Vec<BeeAiTool>>,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflowStepStateDict {
+    meta: BeeAiWorkflowStepStateMeta,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflowStepState {
+    #[serde(rename = "__dict__")]
+    dict: BeeAiWorkflowStepStateDict,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflowStep {
+    #[serde(rename = "py/state")]
+    state: BeeAiWorkflowStepState,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflowInner {
+    #[serde(rename = "_name")]
+    name: String,
+    #[serde(rename = "_steps")]
+    steps: HashMap<String, BeeAiWorkflowStep>,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiWorkflow {
+    workflow: BeeAiWorkflowInner,
+}
+#[derive(Deserialize, Debug)]
+struct BeeAiProgram {
+    inputs: Vec<BeeAiInput>,
+    workflow: BeeAiWorkflow,
+}
+
+#[derive(Serialize, Debug)]
+#[serde(untagged)]
+enum PdlBlock {
+    String(String),
+    Text {
+        #[serde(skip_serializing_if = "Option::is_none")]
+        description: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        role: Option<String>,
+        text: Vec<PdlBlock>,
+    },
+    Model {
+        #[serde(skip_serializing_if = "Option::is_none")]
+        description: Option<String>,
+        model: String,
+        parameters: HashMap<String, Value>,
+    },
+}
+
+pub fn compile(source_file_path: &String, output_path: &String) -> Result<(), Box<dyn Error>> {
+    println!("Compiling beeai {} to {}", source_file_path, output_path);
+
+    // Open the file in read-only mode with buffer.
+    let file = File::open(source_file_path)?;
+    let reader = BufReader::new(file);
+
+    // Read the JSON contents of the file as an instance of `User`.
+    let bee: BeeAiProgram = from_reader(reader)?;
+
+    let inputs: Vec<PdlBlock> = bee
+        .inputs
+        .into_iter()
+        .map(|input| input.state.dict.prompt)
+        .flatten()
+        .map(|prompt| PdlBlock::String(format!("{}\n", prompt)))
+        .collect::<Vec<_>>();
+
+    let system_prompts = bee
+        .workflow
+        .workflow
+        .steps
+        .values()
+        .filter_map(|step| step.state.dict.meta.instructions.clone())
+        .map(|instructions| PdlBlock::Text {
+            role: Some(String::from("system")),
+            text: vec![PdlBlock::String(instructions)],
+            description: None,
+        })
+        .collect::<Vec<_>>();
+
+    let model_calls = bee
+        .workflow
+        .workflow
+        .steps
+        .into_values()
+        .map(|step| (step.state.dict.meta.role, step.state.dict.meta.llm))
+        .map(|(role, llm)| PdlBlock::Model {
+            description: Some(role),
+            model: format!("{}/{}", llm.settings.api_key, llm.model_id),
+            parameters: llm.parameters.state.dict,
+        })
+        .collect::<Vec<_>>();
+
+    let pdl: PdlBlock = PdlBlock::Text {
+        description: Some(bee.workflow.workflow.name),
+        role: None,
+        text: zip!(inputs, system_prompts, model_calls)
+            .map(|(a, (b, c))| [a, b, c])
+            .flatten()
+            .collect(),
+    };
+
+    match output_path.as_str() {
+        "-" => println!("{}", to_string(&pdl)?),
+        _ => {
+            ::std::fs::write(output_path, to_string(&pdl)?)?;
+        }
+    }
+
+    Ok(())
+}

pdl-live-react/src-tauri/src/compile/mod.rs

@@ -0,0 +1 @@
+pub mod beeai;

pdl-live-react/src-tauri/src/lib.rs

@@ -3,6 +3,7 @@ use tauri_plugin_pty;
 
 mod cli;
 mod commands;
+mod compile;
 mod gui;
 mod interpreter;
 
@@ -14,7 +15,13 @@ pub fn run() {
             if args_os().count() <= 1 {
                 gui::setup(app.handle().clone(), "".to_owned())?;
             } else {
-                cli::setup::cli(app)?;
+                match cli::setup::cli(app) {
+                    Ok(()) => ::std::process::exit(0),
+                    Err(s) => {
+                        eprintln!("{}", s);
+                        ::std::process::exit(1)
+                    }
+                }
             }
 
             Ok(())