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

@@ -933,6 +933,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,14 @@
-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;
 
-/// 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()))
-}
-
 /// 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;
 use crate::pdl::requirements::PDL_INTERPRETER;
 
 #[cfg(desktop)]
@@ -19,11 +19,11 @@ 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(&source_file_path);
     let bin_path_future = pip_install_if_needed(&PDL_INTERPRETER);
 
     // wait for any model pulls to finish
-    block_on(pull_future)?;
+    //block_on(pull_future)?;
 
     // wait for any pip installs to finish
     let bin_path = block_on(bin_path_future)?;