feat: update rust to pip install code block requirements

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

Author: starpit

examples/talk/6-code-json.pdl

@@ -26,6 +26,8 @@ text:
 - def: EVAL
   contribute: []
   lang: python
+  requirements:
+  - textdistance
   code:
     |
     import textdistance

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

@@ -2,10 +2,19 @@ use ::std::path::Path;
 use duct::cmd;
 use futures::executor::block_on;
 use yaml_rust2::yaml::LoadError;
+use yaml_rust2::{ScanError, Yaml, YamlLoader};
 
-use crate::interpreter::pip::pip_install_internal_if_needed;
+use crate::interpreter::pip::{
+    pip_install_code_blocks_if_needed, pip_install_internal_if_needed,
+};
 use crate::interpreter::pull::pull_if_needed;
 
+/// 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)
+}
+
 #[cfg(desktop)]
 pub fn run_pdl_program(
     source_file_path: String,
@@ -20,9 +29,10 @@ pub fn run_pdl_program(
     );
 
     // async the model pull and pip installs
-    let pull_future = pull_if_needed(&source_file_path);
-    let bin_path_future =
-        pip_install_internal_if_needed(app_handle, &"interpreter/requirements.txt");
+    let program = &from_path(&source_file_path).unwrap()[0];
+    let pull_future = pull_if_needed(&program);
+    let reqs_future = pip_install_code_blocks_if_needed(&app_handle, &program);
+    let bin_path_future = pip_install_internal_if_needed(&app_handle, &"interpreter/requirements.txt");
 
     // wait for any model pulls to finish
     block_on(pull_future).map_err(|e| match e {
@@ -33,6 +43,7 @@ pub fn run_pdl_program(
 
     // wait for any pip installs to finish
     let bin_path = block_on(bin_path_future)?;
+    block_on(reqs_future)?;
 
     let mut args = vec![
         source_file_path,

pdl-live-react/src-tauri/src/compile/beeai.rs

@@ -291,7 +291,7 @@ fn python_source_to_json(
         eprintln!("Compiling from Python source");
     }
     let bin_path = block_on(pip_install_internal_if_needed(
-        app_handle,
+        &app_handle,
         &"interpreter/beeai-requirements.txt",
     ))?;
 

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

@@ -1,16 +1,18 @@
 use yaml_rust2::Yaml;
 
-/// Extract models referenced by the programs
-pub fn extract_models(programs: Vec<Yaml>) -> Vec<String> {
-    extract_values(programs, "model")
+/// Extract models referenced by the program
+pub fn extract_models(program: &Yaml) -> Vec<String> {
+    extract_values(program, "model")
+}
+
+/// Extract requirements.txt referenced by the program
+pub fn extract_requirements(program: &Yaml) -> Vec<String> {
+    extract_values(program, "requirements")
 }
 
 /// 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: &Yaml, field: &str) -> Vec<String> {
+    let mut values = extract_one_values(program, field);
 
     // A single program may specify the same model more than once. Dedup!
     values.sort();
@@ -20,7 +22,7 @@ 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> {
+fn extract_one_values(program: &Yaml, field: &str) -> Vec<String> {
     let mut values: Vec<String> = Vec::new();
 
     match program {
@@ -31,6 +33,10 @@ fn extract_one_values(program: Yaml, field: &str) -> Vec<String> {
                         Yaml::String(m) => {
                             values.push(m.to_string());
                         }
+                        Yaml::Array(a) => a.into_iter().for_each(|v| match v {
+                            Yaml::String(m) => values.push(m.to_string()),
+                            _ => {}
+                        }),
                         _ => {}
                     },
                     _ => {}

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

@@ -1,23 +1,32 @@
-use ::std::fs::{copy, create_dir_all};
+use ::std::fs::{copy, create_dir_all, write};
 use ::std::path::{Path, PathBuf};
 
 use duct::cmd;
+use rayon::prelude::*;
 use tauri::path::BaseDirectory;
 use tauri::Manager;
+use tempfile::Builder;
+use yaml_rust2::Yaml;
 
+use crate::interpreter::extract;
 use crate::interpreter::shasum;
 
 #[cfg(desktop)]
-pub async fn pip_install_if_needed(
+fn pip_install_if_needed_with_hash(
     cache_path: &Path,
     requirements_path: &Path,
+    hash: String,
+    force: bool,
 ) -> Result<PathBuf, tauri::Error> {
     create_dir_all(&cache_path)?;
 
-    let hash = shasum::sha256sum(&requirements_path)?;
     let venv_path = cache_path.join("venvs").join(hash);
     let bin_path = venv_path.join(if cfg!(windows) { "Scripts" } else { "bin" });
 
+    // re: force, this is part of the short-term hack to install all
+    // code block dependencies in the main interpreter venv. Once we
+    // figure out how to support a separate venv for each code block
+    // (that needs it), we can undo this hack.
     if !venv_path.exists() {
         println!("Creating virtual environment...");
         let python = if cfg!(target_os = "macos") {
@@ -27,18 +36,65 @@ pub async fn pip_install_if_needed(
         };
         cmd!(python, "-mvenv", &venv_path).run()?;
 
-        cmd!(bin_path.join("pip"), "install", "-r", &requirements_path,).run()?;
+        if !force {
+            cmd!(bin_path.join("pip"), "install", "-r", &requirements_path).run()?;
+
+            let cached_requirements_path = venv_path.join("requirements.txt");
+            copy(requirements_path, cached_requirements_path)?;
+        }
+    }
 
-        let cached_requirements_path = venv_path.join("requirements.txt");
-        copy(requirements_path, cached_requirements_path)?;
+    if force {
+        cmd!(bin_path.join("pip"), "install", "-r", &requirements_path).run()?;
     }
 
     Ok(bin_path.to_path_buf())
 }
 
+#[cfg(desktop)]
+fn pip_install_if_needed(
+    cache_path: &Path,
+    requirements_path: &Path,
+) -> Result<PathBuf, tauri::Error> {
+    let hash = shasum::sha256sum(&requirements_path)?;
+    pip_install_if_needed_with_hash(cache_path, requirements_path, hash, false)
+}
+
+#[cfg(desktop)]
+pub async fn pip_install_code_blocks_if_needed(
+    app_handle: &tauri::AppHandle,
+    program: &Yaml,
+) -> Result<(), tauri::Error> {
+    let cache_path = app_handle.path().cache_dir()?.join("pdl");
+
+    // for now, install the requirements in the main interpreter venv
+    let requirements_path = app_handle
+        .path()
+        .resolve("interpreter/requirements.txt", BaseDirectory::Resource)?;
+
+    extract::extract_requirements(program)
+        .into_par_iter()
+        .try_for_each(|req| -> Result<(), tauri::Error> {
+            let req_path = Builder::new()
+                .prefix("pdl-requirements-")
+                .suffix(".txt")
+                .tempfile()?;
+            // This is part of the "force" hack described above, where
+            // we force the code block dependencies to be installed in
+            // the main interpreter venv.
+            let hash = shasum::sha256sum(&requirements_path)?;
+            write(&req_path, req)?;
+            pip_install_if_needed_with_hash(&cache_path, &req_path.path(), hash, true)?;
+            Ok(())
+        })
+        .expect("code block requirements installed");
+
+    Ok(())
+}
+
 #[cfg(desktop)]
 pub async fn pip_install_internal_if_needed(
-    app_handle: tauri::AppHandle,
+    app_handle: &tauri::AppHandle,
     requirements: &str,
 ) -> Result<PathBuf, tauri::Error> {
     // the interpreter requirements.txt
@@ -48,5 +104,5 @@ pub async fn pip_install_internal_if_needed(
 
     let cache_path = app_handle.path().cache_dir()?.join("pdl");
 
-    pip_install_if_needed(&cache_path, &requirements_path).await
+    pip_install_if_needed(&cache_path, &requirements_path)
 }

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

@@ -1,19 +1,13 @@
 use duct::cmd;
 use rayon::prelude::*;
 use yaml_rust2::yaml::LoadError;
-use yaml_rust2::{ScanError, Yaml, YamlLoader};
+use yaml_rust2::Yaml;
 
 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)
-}
-
 /// Pull models (in parallel) from the PDL program in the given filepath.
-pub async fn pull_if_needed(path: &String) -> Result<(), LoadError> {
-    extract::extract_models(from_path(path).unwrap())
+pub async fn pull_if_needed(program: &Yaml) -> Result<(), LoadError> {
+    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/helpers.ts

@@ -5,9 +5,10 @@ import type {
   PdlBlock,
   TextBlock,
   ArgsBlock,
-  CodeBlock,
   PdlModelInput,
   LocalizedExpression,
+  PythonCodeBlock,
+  NonPythonCodeBlock,
 } from "./pdl_ast"
 
 /** Re-export for convenience */
@@ -301,11 +302,13 @@ export function extractStructuredModelResponse({
   return { resultForDisplay, lang, meta }
 }
 
-export function isArgs(block: ArgsBlock | CodeBlock): block is ArgsBlock {
+export function isArgs(
+  block: ArgsBlock | PythonCodeBlock | NonPythonCodeBlock,
+): block is ArgsBlock {
   return Array.isArray((block as ArgsBlock).args)
 }
 
-export function extractCode({ code }: CodeBlock): string {
+export function extractCode({ code }: PythonCodeBlock | NonPythonCodeBlock): string {
   if (
     isNonScalarPdlBlock(code) &&
     hasResult(code) &&