diff --git a/pdl-live-react/src-tauri/src/interpreter/extract.rs b/pdl-live-react/src-tauri/src/interpreter/extract.rs new file mode 100644 index 000000000..640915e8e --- /dev/null +++ b/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) -> Vec { + 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, field: &str) -> Vec { + let mut values = programs + .into_iter() + .flat_map(|p| extract_one_values(p, field)) + .collect::>(); + + // 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 { + let mut values: Vec = 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 +} diff --git a/pdl-live-react/src-tauri/src/interpreter/mod.rs b/pdl-live-react/src-tauri/src/interpreter/mod.rs index db2595326..800d54f66 100644 --- a/pdl-live-react/src-tauri/src/interpreter/mod.rs +++ b/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; diff --git a/pdl-live-react/src-tauri/src/interpreter/pull.rs b/pdl-live-react/src-tauri/src/interpreter/pull.rs index dae8fb1c0..accdbe2ae 100644 --- a/pdl-live-react/src-tauri/src/interpreter/pull.rs +++ b/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, 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 { - let mut models: Vec = 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::>(); - - // 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..]),