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_interpreter_if_needed;
+use crate::interpreter::pip::{
+    pip_install_code_blocks_if_needed, pip_install_interpreter_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,8 +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_interpreter_if_needed(app_handle);
+    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_interpreter_if_needed(&app_handle);
 
     // wait for any model pulls to finish
     block_on(pull_future).map_err(|e| match e {
@@ -32,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/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_interpreter_if_needed(
-    app_handle: tauri::AppHandle,
+    app_handle: &tauri::AppHandle,
 ) -> Result<PathBuf, tauri::Error> {
     // the interpreter requirements.txt
     let requirements_path = app_handle
@@ -47,5 +103,5 @@ pub async fn pip_install_interpreter_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/pdl_ast.d.ts

@@ -2038,6 +2038,11 @@ export type Kind16 = "code"
  *
  */
 export type Lang1 = "python" | "command" | "jinja" | "pdl"
+/**
+ * Pip requirements.txt
+ *
+ */
+export type Requirements = string | string[] | null
 /**
  * Code to execute.
  *
@@ -2984,6 +2989,7 @@ export interface CodeBlock {
   pdl__is_leaf?: PdlIsLeaf16
   kind?: Kind16
   lang: Lang1
+  requirements?: Requirements
   code: Code
 }
 /**

src/pdl/pdl-schema.json

@@ -1699,6 +1699,25 @@
           "title": "Lang",
           "type": "string"
         },
+        "requirements": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "items": {
+                "type": "string"
+              },
+              "type": "array"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "description": "Pip requirements.txt\n    ",
+          "title": "Requirements"
+        },
         "code": {
           "anyOf": [
             {

src/pdl/pdl_ast.py

@@ -15,7 +15,14 @@
     Union,
 )
 
-from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, RootModel
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    RootModel,
+    model_validator,
+)
 from pydantic.json_schema import SkipJsonSchema
 
 from .pdl_lazy import PdlDict, PdlLazy
@@ -479,10 +486,21 @@ class CodeBlock(BaseCodeBlock):
     ]
     """Programming language of the code.
     """
+    requirements: Optional[str | list[str]] = None
+    """Pip requirements.txt
+    """
     code: "BlockType"
     """Code to execute.
     """
 
+    @model_validator(mode="after")
+    def lang_is_python(self):
+        if self.requirements is not None and self.lang != "python":
+            raise ValueError(
+                "CodeBlock requirements field provided for non-python block"
+            )
+        return self
+
 
 class ArgsBlock(BaseCodeBlock):
     """

src/pdl/pdl_dumper.py

@@ -144,6 +144,7 @@ def block_to_dict(  # noqa: C901
         case CodeBlock():
             d["lang"] = block.lang
             d["code"] = block_to_dict(block.code, json_compatible)
+            d["requirements"] = block.requirements
         case GetBlock():
             d["get"] = block.get
         case DataBlock():