feat: port rust model pull logic to use rust AST

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

Author: starpit

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

@@ -1,16 +1,14 @@
-use yaml_rust2::Yaml;
+use crate::pdl::ast::PdlBlock;
 
 /// Extract models referenced by the programs
-pub fn extract_models(programs: Vec<Yaml>) -> Vec<String> {
-    extract_values(programs, "model")
+pub fn extract_models(program: &PdlBlock) -> Vec<String> {
+    extract_values(program, "model")
 }
 
 /// Take a list of Yaml fragments and produce a vector of the string-valued entries of the given field
-pub fn extract_values(programs: Vec<Yaml>, field: &str) -> Vec<String> {
-    let mut values = programs
-        .into_iter()
-        .flat_map(|p| extract_one_values(p, field))
-        .collect::<Vec<String>>();
+pub fn extract_values(program: &PdlBlock, field: &str) -> Vec<String> {
+    let mut values = vec![];
+    extract_values_iter(program, field, &mut values);
 
     // A single program may specify the same model more than once. Dedup!
     values.sort();
@@ -20,38 +18,37 @@ pub fn extract_values(programs: Vec<Yaml>, field: &str) -> Vec<String> {
 }
 
 /// Take one Yaml fragment and produce a vector of the string-valued entries of the given field
-fn extract_one_values(program: Yaml, field: &str) -> Vec<String> {
-    let mut values: Vec<String> = Vec::new();
-
+fn extract_values_iter(program: &PdlBlock, field: &str, values: &mut Vec<String>) {
     match program {
-        Yaml::Hash(h) => {
-            for (key, val) in h {
-                match key {
-                    Yaml::String(f) if f == field => match &val {
-                        Yaml::String(m) => {
-                            values.push(m.to_string());
-                        }
-                        _ => {}
-                    },
-                    _ => {}
-                }
-
-                for m in extract_one_values(val, field) {
-                    values.push(m)
-                }
-            }
+        PdlBlock::Model(b) => values.push(b.model.clone()),
+        PdlBlock::Repeat(b) => {
+            extract_values_iter(&b.repeat, field, values);
         }
-
-        Yaml::Array(a) => {
-            for val in a {
-                for m in extract_one_values(val, field) {
-                    values.push(m)
-                }
+        PdlBlock::Message(b) => {
+            extract_values_iter(&b.content, field, values);
+        }
+        PdlBlock::Array(b) => b
+            .array
+            .iter()
+            .for_each(|p| extract_values_iter(p, field, values)),
+        PdlBlock::Text(b) => b
+            .text
+            .iter()
+            .for_each(|p| extract_values_iter(p, field, values)),
+        PdlBlock::LastOf(b) => b
+            .last_of
+            .iter()
+            .for_each(|p| extract_values_iter(p, field, values)),
+        PdlBlock::If(b) => {
+            extract_values_iter(&b.then, field, values);
+            if let Some(else_) = &b.else_ {
+                extract_values_iter(else_, field, values);
             }
         }
-
+        PdlBlock::Object(b) => b
+            .object
+            .values()
+            .for_each(|p| extract_values_iter(p, field, values)),
         _ => {}
     }
-
-    values
 }

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

@@ -466,23 +466,27 @@ impl<'a> Interpreter<'a> {
             let scope = vm.new_scope_with_builtins();
 
             // TODO vm.new_syntax_error(&err, Some(block.code.as_str()))
-            let code_obj = match vm
-                .compile(
-                    block.code.as_str(),
-                    vm::compiler::Mode::Exec,
-                    "<embedded>".to_owned(),
-                ) {
-                    Ok(x) => Ok(x),
-                    Err(exc) => Err(Box::<dyn Error + Send + Sync>::from(format!("Syntax error in Python code {:?}", exc))),
-                }?;
+            let code_obj = match vm.compile(
+                block.code.as_str(),
+                vm::compiler::Mode::Exec,
+                "<embedded>".to_owned(),
+            ) {
+                Ok(x) => Ok(x),
+                Err(exc) => Err(Box::<dyn Error + Send + Sync>::from(format!(
+                    "Syntax error in Python code {:?}",
+                    exc
+                ))),
+            }?;
 
             // TODO vm.print_exception(exc);
             match vm.run_code_obj(code_obj, scope.clone()) {
                 Ok(_) => Ok(()),
                 Err(exc) => {
                     vm.print_exception(exc);
-                    Err(Box::<dyn Error + Send + Sync>::from("Error executing Python code"))
-                },
+                    Err(Box::<dyn Error + Send + Sync>::from(
+                        "Error executing Python code",
+                    ))
+                }
             }?;
 
             match scope.globals.get_item("result", vm) {
@@ -491,8 +495,10 @@ impl<'a> Interpreter<'a> {
                         Ok(x) => Ok(x),
                         Err(exc) => {
                             vm.print_exception(exc);
-                            Err(Box::<dyn Error + Send + Sync>::from("Unable to stringify Python 'result' value"))
-                        },
+                            Err(Box::<dyn Error + Send + Sync>::from(
+                                "Unable to stringify Python 'result' value",
+                            ))
+                        }
                     }?;
                     let messages = vec![ChatMessage::user(result_string.as_str().to_string())];
                     let trace = PdlBlock::PythonCode(block.clone());
@@ -927,7 +933,7 @@ pub fn run_sync(program: &PdlBlock, cwd: Option<PathBuf>, debug: bool) -> Interp
 }
 
 /// Read in a file from disk and parse it as a PDL program
-fn parse_file(path: &PathBuf) -> Result<PdlBlock, PdlError> {
+pub fn parse_file(path: &PathBuf) -> Result<PdlBlock, PdlError> {
     from_reader(File::open(path)?)
         .map_err(|err| Box::<dyn Error + Send + Sync>::from(err.to_string()))
 }
@@ -937,6 +943,7 @@ pub async fn run_file(source_file_path: &str, debug: bool) -> Interpretation {
     let cwd = path.parent().and_then(|cwd| Some(cwd.to_path_buf()));
     let program = parse_file(&path)?;
 
+    crate::pdl::pull::pull_if_needed(&program).await?;
     run(&program, cwd, debug).await
 }
 

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

@@ -1,20 +1,24 @@
-use ::std::io::{Error, ErrorKind};
+use ::std::io::Error;
 
 use duct::cmd;
 use rayon::prelude::*;
-use yaml_rust2::{Yaml, YamlLoader};
 
+use crate::pdl::ast::PdlBlock;
 use crate::pdl::extract;
+use crate::pdl::interpreter::parse_file;
 
-/// Read the given filesystem path and produce a potentially multi-document Yaml
-fn from_path(path: &str) -> Result<Vec<Yaml>, Error> {
-    let content = std::fs::read_to_string(path)?;
-    YamlLoader::load_from_str(&content).map_err(|e| Error::new(ErrorKind::Other, e.to_string()))
+pub async fn pull_if_needed_from_path(
+    source_file_path: &str,
+) -> Result<(), Box<dyn ::std::error::Error + Send + Sync>> {
+    let program = parse_file(&::std::path::PathBuf::from(source_file_path))?;
+    pull_if_needed(&program)
+        .await
+        .map_err(|e| Box::from(e.to_string()))
 }
 
 /// Pull models (in parallel) from the PDL program in the given filepath.
-pub async fn pull_if_needed(path: &str) -> Result<(), Error> {
-    extract::extract_models(from_path(path)?)
+pub async fn pull_if_needed(program: &PdlBlock) -> Result<(), Error> {
+    extract::extract_models(program)
         .into_par_iter()
         .try_for_each(|model| match model {
             m if model.starts_with("ollama/") => ollama_pull_if_needed(&m[7..]),

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

@@ -3,7 +3,7 @@ use duct::cmd;
 use futures::executor::block_on;
 
 use crate::pdl::pip::pip_install_if_needed;
-use crate::pdl::pull::pull_if_needed;
+use crate::pdl::pull::pull_if_needed_from_path;
 use crate::pdl::requirements::PDL_INTERPRETER;
 
 #[cfg(desktop)]
@@ -19,11 +19,13 @@ pub fn run_pdl_program(
     );
 
     // async the model pull and pip installs
-    let pull_future = pull_if_needed(&source_file_path);
+    let pull_future = pull_if_needed_from_path(&source_file_path);
     let bin_path_future = pip_install_if_needed(&PDL_INTERPRETER);
 
     // wait for any model pulls to finish
-    block_on(pull_future)?;
+    if let Err(e) = block_on(pull_future) {
+        return Err(e);
+    }
 
     // wait for any pip installs to finish
     let bin_path = block_on(bin_path_future)?;