avoid block_on on model calls

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

Author: starpit

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

@@ -38,6 +38,7 @@ minijinja = { version = "2.9.0", features = ["custom_syntax"] }
 ollama-rs = { version = "0.3.0", features = ["tokio"] }
 owo-colors = "4.2.0"
 rustpython-vm = "0.4.0"
+async-recursion = "1.1.1"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 tauri-plugin-cli = "2"

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

@@ -5,7 +5,7 @@ use urlencoding::encode;
 
 use crate::compile;
 use crate::gui::new_window;
-use crate::pdl::interpreter::run_file as runr;
+use crate::pdl::interpreter::run_file_sync as runr;
 use crate::pdl::run::run_pdl_program;
 
 #[cfg(desktop)]

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

@@ -2,7 +2,9 @@ use crate::pdl::interpreter::{pretty_print, run_string};
 
 #[tauri::command]
 pub async fn run_pdl_program(program: String, debug: bool) -> Result<String, String> {
-    let (messages, _) = run_string(&program, debug).map_err(|err| err.to_string())?;
+    let (messages, _) = run_string(&program, debug)
+        .await
+        .map_err(|err| err.to_string())?;
 
     Ok(pretty_print(&messages))
 }

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

@@ -5,6 +5,7 @@ use ::std::error::Error;
 use ::std::fs::File;
 // use ::std::path::PathBuf;
 
+use async_recursion::async_recursion;
 use minijinja::{syntax::SyntaxConfig, Environment};
 use owo_colors::OwoColorize;
 
@@ -27,7 +28,8 @@ use crate::pdl::ast::{
 
 type Context = Vec<ChatMessage>;
 type Scope = HashMap<String, Value>;
-type Interpretation = Result<(Context, PdlBlock), Box<dyn Error>>;
+type Interpretation = Result<(Context, PdlBlock), Box<dyn Error + Send + Sync>>;
+type InterpretationSync = Result<(Context, PdlBlock), Box<dyn Error>>;
 
 struct Interpreter<'a> {
     // batch: u32,
@@ -65,7 +67,7 @@ impl<'a> Interpreter<'a> {
         }
     }
 
-    fn run_with_emit(
+    async fn run_with_emit(
         &mut self,
         program: &PdlBlock,
         context: Context,
@@ -75,12 +77,12 @@ impl<'a> Interpreter<'a> {
         self.emit = emit;
 
         let (messages, trace) = match program {
-            PdlBlock::String(s) => self.run_string(s, context),
-            PdlBlock::Call(block) => self.run_call(block, context),
-            PdlBlock::PythonCode(block) => self.run_python_code(block, context),
-            PdlBlock::Model(block) => self.run_model(block, context),
-            PdlBlock::Repeat(block) => self.run_repeat(block, context),
-            PdlBlock::Text(block) => self.run_text(block, context),
+            PdlBlock::String(s) => self.run_string(s, context).await,
+            PdlBlock::Call(block) => self.run_call(block, context).await,
+            PdlBlock::PythonCode(block) => self.run_python_code(block, context).await,
+            PdlBlock::Model(block) => self.run_model(block, context).await,
+            PdlBlock::Repeat(block) => self.run_repeat(block, context).await,
+            PdlBlock::Text(block) => self.run_text(block, context).await,
             _ => Err(Box::from(format!("Unsupported block {:?}", program))),
         }?;
 
@@ -96,19 +98,21 @@ impl<'a> Interpreter<'a> {
         Ok((messages, trace))
     }
 
-    fn run_quiet(&mut self, program: &PdlBlock, context: Context) -> Interpretation {
-        self.run_with_emit(program, context, false)
+    #[async_recursion]
+    async fn run_quiet(&mut self, program: &PdlBlock, context: Context) -> Interpretation {
+        self.run_with_emit(program, context, false).await
     }
 
-    fn run(&mut self, program: &PdlBlock, context: Context) -> Interpretation {
-        self.run_with_emit(program, context, self.emit)
+    #[async_recursion]
+    async fn run(&mut self, program: &PdlBlock, context: Context) -> Interpretation {
+        self.run_with_emit(program, context, self.emit).await
     }
 
     // Evaluate as a Jinja2 expression
     fn eval<T: serde::de::DeserializeOwned + ::std::convert::From<String>>(
         &self,
         expr: &String,
-    ) -> Result<T, Box<dyn Error>> {
+    ) -> Result<T, Box<dyn Error + Send + Sync>> {
         let result = self
             .jinja_env
             .render_str(expr.as_str(), self.scope.last().unwrap_or(&HashMap::new()))?;
@@ -127,7 +131,7 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::String
-    fn run_string(&self, msg: &String, _context: Context) -> Interpretation {
+    async fn run_string(&self, msg: &String, _context: Context) -> Interpretation {
         let trace = self.eval::<PdlBlock>(msg)?;
         if self.debug {
             eprintln!("String {} -> {:?}", msg, trace);
@@ -142,7 +146,7 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::Call
-    fn run_call(&mut self, block: &PdlCallBlock, context: Context) -> Interpretation {
+    async fn run_call(&mut self, block: &PdlCallBlock, context: Context) -> Interpretation {
         if self.debug {
             eprintln!("Call {:?}({:?})", block.call, block.args);
         }
@@ -154,7 +158,7 @@ impl<'a> Interpreter<'a> {
                     // args was a string that eval'd to an Object
                     Value::Object(m) => Ok(Some(self.to_pdl(&m))),
                     // args was a string that eval'd to something we don't understand
-                    y => Err(Box::<dyn Error>::from(format!(
+                    y => Err(Box::<dyn Error + Send + Sync>::from(format!(
                         "Invalid arguments to call {:?}",
                         y
                     ))),
@@ -170,7 +174,7 @@ impl<'a> Interpreter<'a> {
         self.extend_scope_with_map(&args);
 
         let res = match self.eval::<PdlBlock>(&block.call)? {
-            PdlBlock::Function(f) => self.run(&f.return_, context.clone()),
+            PdlBlock::Function(f) => self.run(&f.return_, context.clone()).await,
             _ => Err(Box::from(format!("call of non-function {:?}", &block.call))),
         };
         self.scope.pop();
@@ -230,7 +234,11 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::PythonCode
-    fn run_python_code(&mut self, block: &PdlPythonCodeBlock, context: Context) -> Interpretation {
+    async fn run_python_code(
+        &mut self,
+        block: &PdlPythonCodeBlock,
+        context: Context,
+    ) -> Interpretation {
         use rustpython_vm as vm;
         vm::Interpreter::without_stdlib(Default::default()).enter(|vm| -> Interpretation {
             let scope = vm.new_scope_with_builtins();
@@ -275,7 +283,7 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::Model
-    fn run_model(&mut self, block: &PdlModelBlock, context: Context) -> Interpretation {
+    async fn run_model(&mut self, block: &PdlModelBlock, context: Context) -> Interpretation {
         match &block.model {
             pdl_model
                 if pdl_model.starts_with("ollama/") || pdl_model.starts_with("ollama_chat/") =>
@@ -295,7 +303,7 @@ impl<'a> Interpreter<'a> {
                 let messages = match &block.input {
                     Some(input) => {
                         // TODO ignoring trace
-                        let (messages, _trace) = self.run_quiet(&*input, context)?;
+                        let (messages, _trace) = self.run_quiet(&*input, context).await?;
                         messages
                     }
                     None => context,
@@ -310,24 +318,27 @@ impl<'a> Interpreter<'a> {
                     eprintln!(
                         "Ollama {:?} model={:?} prompt={:?} history={:?}",
                         block.description.clone().unwrap_or("".into()),
-                        &block.model,
-                        &prompt,
-                        &history
+                        block.model,
+                        prompt,
+                        history
                     );
                 }
 
                 let req = ChatMessageRequest::new(model.into(), vec![prompt.clone()])
                     .options(options)
                     .tools(tools);
-                let res = /*self.rt.*/tauri::async_runtime::block_on(ollama.send_chat_messages_with_history(
-                    &mut history,
-                    req,
-                    //ollama.generate(GenerationRequest::new(model.into(), prompt),
-                ))?;
+                let res = ollama
+                    .send_chat_messages_with_history(
+                        &mut history,
+                        req,
+                        //ollama.generate(GenerationRequest::new(model.into(), prompt),
+                    )
+                    .await?;
                 // dbg!("Model result {:?}", &res);
 
                 let mut trace = block.clone();
                 trace.pdl_result = Some(res.message.content.clone());
+
                 if let Some(usage) = res.final_data {
                     trace.pdl_usage = Some(PdlUsage {
                         prompt_tokens: usage.prompt_eval_count,
@@ -344,7 +355,7 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::Repeat
-    fn run_repeat(&mut self, block: &PdlRepeatBlock, _context: Context) -> Interpretation {
+    async fn run_repeat(&mut self, block: &PdlRepeatBlock, _context: Context) -> Interpretation {
         let for_ = block
             .for_
             .iter()
@@ -368,7 +379,11 @@ impl<'a> Interpreter<'a> {
         }
     }
 
-    fn parse_result(&self, parser: &PdlParser, result: &String) -> Result<Value, Box<dyn Error>> {
+    fn parse_result(
+        &self,
+        parser: &PdlParser,
+        result: &String,
+    ) -> Result<Value, Box<dyn Error + Send + Sync>> {
         match parser {
             PdlParser::Json => Ok(from_str(result)?),
         }
@@ -395,7 +410,7 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::Text
-    fn run_text(&mut self, block: &PdlTextBlock, context: Context) -> Interpretation {
+    async fn run_text(&mut self, block: &PdlTextBlock, context: Context) -> Interpretation {
         if self.debug {
             eprintln!(
                 "Text {:?}",
@@ -411,14 +426,14 @@ impl<'a> Interpreter<'a> {
         let mut output_blocks = vec![];
 
         self.extend_scope_with_map(&block.defs);
-        block.text.iter().try_for_each(|block| {
+        let mut iter = block.text.iter();
+        while let Some(block) = iter.next() {
             // run each element of the Text block
-            let (this_messages, trace) = self.run(&block, input_messages.clone())?;
+            let (this_messages, trace) = self.run(&block, input_messages.clone()).await?;
             input_messages.extend(this_messages.clone());
             output_messages.extend(this_messages);
             output_blocks.push(trace);
-            Ok::<(), Box<dyn Error>>(())
-        })?;
+        }
         self.scope.pop();
 
         let mut trace = block.clone();
@@ -455,22 +470,37 @@ impl<'a> Interpreter<'a> {
     }
 }
 
-pub fn run(program: &PdlBlock, debug: bool) -> Interpretation {
+pub async fn run(program: &PdlBlock, debug: bool) -> Interpretation {
     let mut interpreter = Interpreter::new();
     interpreter.debug = debug;
-    interpreter.run(&program, vec![])
+    interpreter.run(&program, vec![]).await
+}
+
+pub fn run_sync(program: &PdlBlock, debug: bool) -> InterpretationSync {
+    tauri::async_runtime::block_on(run(program, debug))
+        .map_err(|err| Box::<dyn ::std::error::Error>::from(err.to_string()))
+}
+
+pub async fn run_file(source_file_path: &str, debug: bool) -> Interpretation {
+    run(&from_reader(File::open(source_file_path)?)?, debug).await
+}
+
+pub fn run_file_sync(source_file_path: &str, debug: bool) -> InterpretationSync {
+    tauri::async_runtime::block_on(run_file(source_file_path, debug))
+        .map_err(|err| Box::<dyn ::std::error::Error>::from(err.to_string()))
 }
 
-pub fn run_file(source_file_path: &str, debug: bool) -> Interpretation {
-    run(&from_reader(File::open(source_file_path)?)?, debug)
+pub async fn run_string(source: &str, debug: bool) -> Interpretation {
+    run(&from_yaml_str(source)?, debug).await
 }
 
-pub fn run_string(source: &str, debug: bool) -> Interpretation {
-    run(&from_yaml_str(source)?, debug)
+pub async fn run_json(source: Value, debug: bool) -> Interpretation {
+    run_string(&to_string(&source)?, debug).await
 }
 
-pub fn run_json(source: Value, debug: bool) -> Interpretation {
-    run_string(&to_string(&source)?, debug)
+pub fn run_json_sync(source: Value, debug: bool) -> InterpretationSync {
+    tauri::async_runtime::block_on(run_json(source, debug))
+        .map_err(|err| Box::<dyn ::std::error::Error>::from(err.to_string()))
 }
 
 pub fn pretty_print(messages: &Vec<ChatMessage>) -> String {

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

@@ -6,7 +6,7 @@ mod tests {
 
     use crate::pdl::{
         ast::{PdlBlock, PdlModelBlock, PdlParser, PdlTextBlock},
-        interpreter::{run, run_json},
+        interpreter::{run_json_sync as run_json, run_sync as run},
     };
 
     use ollama_rs::generation::chat::MessageRole;

pdl-live-react/src-tauri/src/pdl/interpreter_tests.rs

@@ -28,6 +28,7 @@ const initialInput = `text:
     def: foo
   - \${ foo.key }`
 export default function Run() {
+  const [running, setRunning] = useState(false)
   const [input, setInput] = useState(initialInput)
   const [_error, setError] = useState(false)
 
@@ -69,6 +70,7 @@ export default function Run() {
 
   const run = async () => {
     try {
+      setRunning(true)
       term?.reset()
       const result = await invoke("run_pdl_program", {
         program: input,
@@ -79,6 +81,8 @@ export default function Run() {
     } catch (err) {
       term?.write(String(err))
       setError(true)
+    } finally {
+      setRunning(false)
     }
   }
 
@@ -132,7 +136,9 @@ export default function Run() {
         <Toolbar>
           <ToolbarContent>
             <ToolbarItem>
-              <Button onClick={run}>Run</Button>
+              <Button onClick={run} isLoading={running}>
+                Run
+              </Button>
             </ToolbarItem>
           </ToolbarContent>
         </Toolbar>