refactor: extract a common extract_values() in the rust model pulling logic

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

@@ -0,0 +1,57 @@
+use yaml_rust2::Yaml;
+
+/// Extract models referenced by the programs
+pub fn extract_models(programs: Vec<Yaml>) -> Vec<String> {
+    extract_values(programs, "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>>();
+
+    // A single program may specify the same model more than once. Dedup!
+    values.sort();
+    values.dedup();
+
+    values
+}
+
+/// 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();
+
+    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)
+                }
+            }
+        }
+
+        Yaml::Array(a) => {
+            for val in a {
+                for m in extract_one_values(val, field) {
+                    values.push(m)
+                }
+            }
+        }
+
+        _ => {}
+    }
+
+    values
+}

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

@@ -1,3 +1,4 @@
+pub mod extract;
 pub mod pip;
 pub mod pull;
 pub mod shasum;

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

@@ -3,62 +3,17 @@ use rayon::prelude::*;
 use yaml_rust2::yaml::LoadError;
 use yaml_rust2::{ScanError, Yaml, YamlLoader};
 
+use crate::interpreter::extract;
+
 /// Read the given filesystem path and produce a potentially multi-document Yaml
 fn from_path(path: &String) -> Result<Vec<Yaml>, ScanError> {
     let content = std::fs::read_to_string(path).unwrap();
     YamlLoader::load_from_str(&content)
 }
 
-/// Take one Yaml fragment and produce the a vector of the models that are used
-fn extract_models(program: Yaml) -> Vec<String> {
-    let mut models: Vec<String> = Vec::new();
-
-    match program {
-        Yaml::Hash(h) => {
-            for (key, val) in h {
-                match key.as_str() {
-                    Some("model") => match &val {
-                        Yaml::String(m) => {
-                            models.push(m.to_string());
-                        }
-                        _ => {}
-                    },
-                    _ => {}
-                }
-
-                for m in extract_models(val) {
-                    models.push(m)
-                }
-            }
-        }
-
-        Yaml::Array(a) => {
-            for val in a {
-                for m in extract_models(val) {
-                    models.push(m)
-                }
-            }
-        }
-
-        _ => {}
-    }
-
-    models
-}
-
 /// Pull models (in parallel) from the PDL program in the given filepath.
 pub async fn pull_if_needed(path: &String) -> Result<(), LoadError> {
-    let mut models = from_path(path)
-        .unwrap()
-        .into_iter()
-        .flat_map(extract_models)
-        .collect::<Vec<String>>();
-
-    // A single program may specify the same model more than once. Dedup!
-    models.sort();
-    models.dedup();
-
-    models
+    extract::extract_models(from_path(path).unwrap())
         .into_par_iter()
         .try_for_each(|model| match model {
             m if model.starts_with("ollama/") => ollama_pull_if_needed(&m[7..]),