feat: pull models in rust

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,7 @@ urlencoding = "2.1.3"
 tempfile = "3.16.0"
 file_diff = "1.0.0"
 duct = "0.13.7"
+rayon = "1.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()
     );
+
+    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,27 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+enum PdlBlock {
+    LiteLlmModelBLock,
+    TextBlock,
+}
+
+#[derive(Serialize, Deserialize)]
+struct TextBlock {
+  description: Option<str>,
+  defs?: Defs10
+  def?: Def9
+  contribute?: Contribute9
+  parser?: Parser9
+  fallback?: Fallback9
+  role?: Role9
+  context?: Context9
+  pdl__id?: PdlId9
+  pdl__result?: unknown
+  pdl__location?: PdlLocationType | null
+  pdl__timing?: PdlTiming | null
+  pdl__is_leaf?: PdlIsLeaf9
+  kind?: Kind9
+  text: Text
+    
+}

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

@@ -0,0 +1 @@
+[email protected]

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

@@ -0,0 +1,71 @@
+use duct::cmd;
+use rayon::prelude::*;
+use serde_json::{Error, Map, Result, Value};
+
+fn from_path(path: &String) -> Result<Map<String, Value>> {
+    let content = std::fs::read_to_string(path).unwrap();
+    let map: Map<String, Value> = serde_json::from_str(&content).unwrap();
+
+    Ok(map)
+}
+
+fn extract_models(program: Map<String, Value>) -> Result<Vec<String>> {
+    let mut map: Vec<String> = Vec::new();
+    for (key, val) in program {
+        match key.as_str() {
+            "model" => match &val {
+                Value::String(m) => {
+                    map.push(m.to_string());
+                }
+                _ => {}
+            },
+            _ => {}
+        }
+
+        match val {
+            Value::Array(a) => {
+                for v in a {
+                    match v {
+                        Value::Object(h) => {
+                            for m in extract_models(h).unwrap() {
+                                map.push(m)
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            }
+            Value::Object(h) => {
+                for m in extract_models(h).unwrap() {
+                    map.push(m)
+                }
+            }
+            _ => {}
+        }
+    }
+
+    Ok(map)
+}
+
+pub fn pull_if_needed(path: &String) -> Result<()> {
+    extract_models(from_path(path)?)?
+        .into_par_iter()
+        .try_for_each(|model: String| -> Result<()> {
+            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(())
+}
+
+fn ollama_pull(model: &str) -> Result<()> {
+    cmd!("ollama", "pull", model).run().map_err(Error::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() {