repeat blocks

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

Author: starpit

pdl-live-react/src-tauri/.taurignore

@@ -1,4 +1,6 @@
 # emacs temp files
 *~
 \#*\#
-.\#*
+.\#*
+tests/**/*.txt
+tests/**/*.pdl

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

@@ -12,9 +12,9 @@ use serde_json::{from_reader, json, to_string, Map, Value};
 use tempfile::Builder;
 
 use crate::pdl::ast::{
-    PdlBaseType, PdlBlock, PdlCallBlock, PdlFunctionBlock, PdlMessageBlock, PdlModelBlock,
-    PdlObjectBlock, PdlOptionalType, PdlParser, PdlPythonCodeBlock, PdlRepeatBlock, PdlTextBlock,
-    PdlType, Role,
+    PdlBaseType, PdlBlock, PdlCallBlock, PdlFunctionBlock, PdlListOrString, PdlMessageBlock,
+    PdlModelBlock, PdlObjectBlock, PdlOptionalType, PdlParser, PdlPythonCodeBlock, PdlRepeatBlock,
+    PdlTextBlock, PdlType, Role,
 };
 use crate::pdl::pip::pip_install_if_needed;
 use crate::pdl::requirements::BEEAI_FRAMEWORK;
@@ -58,7 +58,7 @@ struct BeeAiToolState {
     name: String,
     description: Option<String>,
     input_schema: BeeAiToolSchema,
-    options: Option<HashMap<String, Value>>,
+    // options: Option<HashMap<String, Value>>,
 }
 #[derive(Deserialize, Debug)]
 struct BeeAiTool {
@@ -223,7 +223,7 @@ fn call_tools(model: &String, parameters: &HashMap<String, Value>) -> PdlBlock {
     let mut for_ = HashMap::new();
     for_.insert(
         "tool".to_string(),
-        "${ response.choices[0].message.tool_calls }".to_string(),
+        PdlListOrString::String("${ response.choices[0].message.tool_calls }".to_string()),
     );
 
     // response.choices[0].message.tool_calls

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

@@ -197,10 +197,17 @@ impl PdlModelBlock {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(untagged)]
+pub enum PdlListOrString {
+    String(String),
+    List(Vec<Value>),
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct PdlRepeatBlock {
     #[serde(rename = "for")]
-    pub for_: HashMap<String, String>,
+    pub for_: HashMap<String, PdlListOrString>,
     pub repeat: Box<PdlBlock>,
 }
 

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

@@ -25,8 +25,8 @@ use serde_json::{from_str, to_string, Map, Value};
 use serde_norway::{from_reader, from_str as from_yaml_str};
 
 use crate::pdl::ast::{
-    PdlBlock, PdlCallBlock, PdlModelBlock, PdlParser, PdlPythonCodeBlock, PdlReadBlock,
-    PdlRepeatBlock, PdlTextBlock, PdlUsage, Role,
+    PdlBlock, PdlCallBlock, PdlListOrString, PdlModelBlock, PdlParser, PdlPythonCodeBlock,
+    PdlReadBlock, PdlRepeatBlock, PdlTextBlock, PdlUsage, Role,
 };
 
 type Context = Vec<ChatMessage>;
@@ -112,7 +112,7 @@ impl<'a> Interpreter<'a> {
         self.run_with_emit(program, context, self.emit).await
     }
 
-    // Evaluate as a Jinja2 expression
+    // Evaluate String as a Jinja2 expression
     fn eval<T: serde::de::DeserializeOwned + ::std::convert::From<String>>(
         &self,
         expr: &String,
@@ -134,6 +134,43 @@ impl<'a> Interpreter<'a> {
         }))
     }
 
+    fn eval_complex(&self, expr: &Value) -> Result<Value, Box<dyn Error + Send + Sync>> {
+        match expr {
+            Value::String(s) => self.eval(s),
+            Value::Array(a) => Ok(Value::Array(
+                a.iter()
+                    .map(|v| self.eval_complex(v))
+                    .collect::<Result<_, _>>()?,
+            )),
+            Value::Object(o) => Ok(Value::Object(
+                o.iter()
+                    .map(|(k, v)| match self.eval_complex(v) {
+                        Ok(v) => Ok((k.clone(), v)),
+                        Err(e) => Err(e),
+                    })
+                    .collect::<Result<_, _>>()?,
+            )),
+            v => Ok(v.clone()),
+        }
+    }
+
+    // Evaluate an string or list of Values into a list of Values
+    fn eval_list_or_string(
+        &self,
+        expr: &PdlListOrString,
+    ) -> Result<Vec<Value>, Box<dyn Error + Send + Sync>> {
+        match expr {
+            PdlListOrString::String(s) => match self.eval::<Value>(s)? {
+                Value::Array(a) => Ok(a),
+                x => Err(Box::from(format!(
+                    "Jinja string expanded to non-list. {} -> {:?}",
+                    s, x
+                ))),
+            },
+            PdlListOrString::List(l) => l.iter().map(|v| self.eval_complex(v)).collect(),
+        }
+    }
+
     // Run a PdlBlock::String
     async fn run_string(&self, msg: &String, _context: Context) -> Interpretation {
         let trace = self.eval::<PdlBlock>(msg)?;
@@ -179,53 +216,26 @@ impl<'a> Interpreter<'a> {
             eprintln!("Call {:?}({:?})", block.call, block.args);
         }
 
-        let args = match &block.args {
-            Some(x) => match x {
-                // args is a string; eval it and see if we get an Object out the other side
-                Value::String(s) => match self.eval::<Value>(&s)? {
-                    // 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 + Send + Sync>::from(format!(
-                        "Invalid arguments to call {:?}",
-                        y
-                    ))),
-                },
-                // args is already an Object
-                Value::Object(m) => Ok(Some(self.to_pdl(&m))),
-                // args is something we don't understand
-                y => Err(Box::from(format!("Invalid arguments to call {:?}", y))),
-            },
-            // no args... that's ok (TODO: check against function schema)
-            None => Ok(None),
-        }?;
-        self.extend_scope_with_map(&args);
+        if let Some(args) = &block.args {
+            match self.eval_complex(args)? {
+                Value::Object(m) => Ok(self.extend_scope_with_json_map(m)),
+                x => Err(Box::<dyn Error + Send + Sync>::from(format!(
+                    "Call arguments not a map: {:?}",
+                    x
+                ))),
+            }?;
+        }
 
         let res = match self.eval::<PdlBlock>(&block.call)? {
             PdlBlock::Function(f) => self.run(&f.return_, context.clone()).await,
             _ => Err(Box::from(format!("call of non-function {:?}", &block.call))),
         };
-        self.scope.pop();
 
-        res
-    }
+        if let Some(_) = block.args {
+            self.scope.pop();
+        }
 
-    fn to_pdl(&self, m: &Map<String, Value>) -> HashMap<String, PdlBlock> {
-        m.into_iter()
-            .map(|(k, v)| {
-                (
-                    k.clone(),
-                    match v {
-                        Value::String(s) => PdlBlock::String(s.clone()),
-                        Value::Number(n) => PdlBlock::Number(n.clone()),
-                        x => {
-                            eprintln!("Unhandled arg value {:?}", x);
-                            "error".into()
-                        }
-                    },
-                )
-            })
-            .collect()
+        res
     }
 
     fn to_ollama_model_options(
@@ -428,19 +438,41 @@ impl<'a> Interpreter<'a> {
     }
 
     // Run a PdlBlock::Repeat
-    async fn run_repeat(&mut self, block: &PdlRepeatBlock, _context: Context) -> Interpretation {
-        let for_ = block
+    async fn run_repeat(&mut self, block: &PdlRepeatBlock, context: Context) -> Interpretation {
+        // { i:[1,2,3], j: [4,5,6]} -> ([i,j], [[1,2,3],[4,5,6]])
+        //        let (variables, values): (Vec<_>, Vec<Vec<_>>) = block
+        //            .into_iter()
+        //            .unzip();
+        let map = block
             .for_
             .iter()
-            .map(|(var, values)| (var, self.eval::<PdlBlock>(&values)));
+            .map(|(var, values)| match self.eval_list_or_string(values) {
+                Ok(value) => Ok((var.clone(), value)),
+                Err(e) => Err(e),
+            })
+            .collect::<Result<HashMap<_, _>, _>>()?;
 
         if self.debug {
-            eprintln!("Repeat {:?}", &for_);
+            eprintln!("Repeat {:?}", map);
         }
-        Ok((
-            vec![ChatMessage::user("TODO".into())],
-            PdlBlock::Repeat(block.clone()),
-        ))
+
+        let mut messages = vec![];
+        let mut trace = vec![];
+        if let Some(n) = map.iter().map(|(_, v)| v.len()).min() {
+            for iter in 0..n {
+                let scope: HashMap<String, Value> = map
+                    .iter()
+                    .map(|(k, v)| (k.clone(), v[iter].clone()))
+                    .collect();
+                self.extend_scope_with_map(scope);
+                let (ms, t) = self.run_quiet(&block.repeat, context.clone()).await?;
+                messages.extend(ms);
+                trace.push(t);
+                self.scope.pop();
+            }
+        }
+
+        Ok((messages, PdlBlock::Repeat(block.clone())))
     }
 
     fn to_ollama_role(&self, role: &Role) -> MessageRole {
@@ -462,7 +494,18 @@ impl<'a> Interpreter<'a> {
         }
     }
 
-    fn extend_scope_with_map(&mut self, map: &Option<HashMap<String, PdlBlock>>) {
+    fn extend_scope_with_map(&mut self, new_scope: HashMap<String, Value>) {
+        self.scope.push(new_scope);
+    }
+
+    fn extend_scope_with_json_map(&mut self, new_scope: Map<String, Value>) {
+        let mut scope = self.scope.last().unwrap_or(&HashMap::new()).clone();
+        // TODO figure out iterators
+        scope.extend(new_scope.into_iter().collect::<HashMap<String, Value>>());
+        self.extend_scope_with_map(scope);
+    }
+
+    fn extend_scope_with_block_map(&mut self, map: &Option<HashMap<String, PdlBlock>>) {
         let cur_scope = self.scope.last().unwrap_or(&HashMap::new()).clone();
         let new_scope = match map {
             Some(defs) => {
@@ -479,7 +522,7 @@ impl<'a> Interpreter<'a> {
             None => cur_scope,
         };
 
-        self.scope.push(new_scope);
+        self.extend_scope_with_map(new_scope);
     }
 
     // Run a PdlBlock::Text
@@ -498,7 +541,7 @@ impl<'a> Interpreter<'a> {
         let mut output_messages = vec![];
         let mut output_blocks = vec![];
 
-        self.extend_scope_with_map(&block.defs);
+        self.extend_scope_with_block_map(&block.defs);
         let mut iter = block.text.iter();
         while let Some(block) = iter.next() {
             // run each element of the Text block

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

@@ -5,11 +5,11 @@ mod tests {
     use serde_json::json;
 
     use crate::pdl::{
-        ast::{PdlBlock, PdlModelBlock, PdlParser, PdlTextBlock},
+        ast::{PdlBlock, PdlModelBlock},
         interpreter::{run_json_sync as run_json, run_sync as run},
     };
 
-    use ollama_rs::generation::chat::MessageRole;
+    use ollama_rs::generation::chat::{ChatMessage, MessageRole};
 
     const DEFAULT_MODEL: &'static str = "ollama/granite3.2:2b";
 
@@ -220,4 +220,78 @@ mod tests {
         assert_eq!(messages[0].content, "this should be foo\n");
         Ok(())
     }
+
+    #[test]
+    fn text_repeat_numbers_1d() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "for": {
+                "x": [1,2,3]
+            },
+            "repeat": {
+                "text": [
+                    "${ x + 1 }"
+                ]
+            }
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 3);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "2");
+        assert_eq!(messages[1].role, MessageRole::User);
+        assert_eq!(messages[1].content, "3");
+        assert_eq!(messages[2].role, MessageRole::User);
+        assert_eq!(messages[2].content, "4");
+        Ok(())
+    }
+
+    #[test]
+    fn text_repeat_numbers_2d() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "for": {
+                "x": [1,2,3],
+                "y": [4,5,6]
+            },
+            "repeat": {
+                "text": [
+                    "${ x + y }"
+                ]
+            }
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 3);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "5");
+        assert_eq!(messages[1].role, MessageRole::User);
+        assert_eq!(messages[1].content, "7");
+        assert_eq!(messages[2].role, MessageRole::User);
+        assert_eq!(messages[2].content, "9");
+        Ok(())
+    }
+
+    #[test]
+    fn text_repeat_mix_2d() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "for": {
+                "x": [{"z": 4}, {"z": 5}, {"z": 6}],
+                "y": ["a","b","c"]
+            },
+            "repeat": {
+                "text": [
+                    "${ x.z ~ y }" // ~ is string concatenation in jinja
+                ]
+            }
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 3);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "4a");
+        assert_eq!(messages[1].role, MessageRole::User);
+        assert_eq!(messages[1].content, "5b");
+        assert_eq!(messages[2].role, MessageRole::User);
+        assert_eq!(messages[2].content, "6c");
+        Ok(())
+    }
 }

pdl-live-react/src-tauri/tests/cli/call-no-args.pdl

@@ -0,0 +1,9 @@
+defs:
+  foo:
+    function: {}
+    return:
+      description: nullary function
+      text:
+        - hello world
+text:
+  - call: ${ foo }

pdl-live-react/src-tauri/tests/cli/call-with-args.pdl

@@ -0,0 +1,12 @@
+defs:
+  foo:
+    function:
+      x: int
+    return:
+      description: nullary function
+      text:
+        - hello world ${x+1}
+text:
+  - call: ${ foo }
+    args:
+      x: 3

pdl-live-react/src-tauri/tests/cli/code-python.pdl

@@ -0,0 +1,2 @@
+lang: python
+code: 'print(''hi ho''); result = {"foo": 3}'