feat: pull models in rust (#743)

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

Author: starpit

.github/workflows/tauri-cli.yml

@@ -25,10 +25,11 @@ jobs:
         with:
           node-version: 22
       - name: Install dependencies
+        # sleep 2 to wait for ollama to be running... hack warning
         run: |
           npm ci & sudo apt update && \
             sudo apt install -y libgtk-3-dev libwebkit2gtk-4.1-dev librsvg2-dev patchelf at-spi2-core && \
-            (curl -fsSL https://ollama.com/install.sh | sudo -E sh && sleep 2 && ollama pull granite3.2:2b)
+            (curl -fsSL https://ollama.com/install.sh | sudo -E sh && sleep 2)
           wait
       - name: Test production build
         run: npm run tauri build -- --bundles deb # Skip testing appimage, is this dangerous? It's slow...

pdl-live-react/src-tauri/Cargo.lock

@@ -29,6 +29,8 @@ urlencoding = "2.1.3"
 tempfile = "3.16.0"
 file_diff = "1.0.0"
 duct = "0.13.7"
+rayon = "1.10.0"
+yaml-rust2 = "0.10.0"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 tauri-plugin-cli = "2"

pdl-live-react/src-tauri/Cargo.toml

@@ -6,6 +6,8 @@ use duct::cmd;
 use tauri::path::BaseDirectory;
 use tauri::Manager;
 
+use crate::interpreter::load;
+
 #[cfg(desktop)]
 fn pip_install_if_needed(app_handle: tauri::AppHandle) -> Result<PathBuf, tauri::Error> {
     let cache_path = app_handle.path().cache_dir()?.join("pdl");
@@ -79,6 +81,9 @@ pub fn run_pdl_program(
         "Running {:#?}",
         Path::new(&source_file_path).file_name().unwrap()
     );
+
+    let _ = load::pull_if_needed(&source_file_path);
+
     let bin_path = pip_install_if_needed(app_handle)?;
     let trace_arg = if let Some(arg) = trace_file {
         if let serde_json::Value::String(f) = &arg.value {

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

@@ -0,0 +1,74 @@
+use duct::cmd;
+use rayon::prelude::*;
+use yaml_rust2::yaml::LoadError;
+use yaml_rust2::{ScanError, Yaml, YamlLoader};
+
+/// 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 fn pull_if_needed(path: &String) -> Result<(), LoadError> {
+    from_path(path)
+        .unwrap()
+        .into_iter()
+        .flat_map(extract_models)
+        .collect::<Vec<String>>()
+        .into_par_iter()
+        .try_for_each(|model| match model {
+            m if model.starts_with("ollama/") => ollama_pull(&m[7..]),
+            m if model.starts_with("ollama_chat/") => ollama_pull(&m[12..]),
+            _ => {
+                eprintln!("Skipping model pull for {}", model);
+                Ok(())
+            }
+        })
+        .expect("successfully pulled models");
+
+    Ok(())
+}
+
+/// The Ollama implementation of a single model pull
+fn ollama_pull(model: &str) -> Result<(), LoadError> {
+    cmd!("ollama", "pull", model).run().map_err(LoadError::IO)?;
+    Ok(())
+}

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

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

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

@@ -4,6 +4,7 @@ use tauri_plugin_pty;
 mod cli;
 mod commands;
 mod gui;
+mod interpreter;
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {