if blocks

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

Author: starpit

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

@@ -302,6 +302,13 @@ pub struct ReadBlock {
     pub multiline: Option<bool>,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(untagged)]
+pub enum StringOrBoolean {
+    String(String),
+    Boolean(bool),
+}
+
 /// Conditional control structure.
 ///
 /// Example:
@@ -317,15 +324,18 @@ pub struct ReadBlock {
 pub struct IfBlock {
     /// The condition to check
     #[serde(rename = "if")]
-    condition: String,
+    pub condition: StringOrBoolean,
 
     /// Branch to execute if the condition is true
-    then: Box<PdlBlock>,
+    pub then: Box<PdlBlock>,
 
     /// Branch to execute if the condition is false.
     #[serde(rename = "else")]
     #[serde(skip_serializing_if = "Option::is_none")]
-    else_: Option<Box<PdlBlock>>,
+    pub else_: Option<Box<PdlBlock>>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub defs: Option<HashMap<String, PdlBlock>>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

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::{
-    CallBlock, ListOrString, ModelBlock, PdlBlock, PdlParser, PdlUsage, PythonCodeBlock, ReadBlock,
-    RepeatBlock, Role, TextBlock,
+    CallBlock, IfBlock, ListOrString, ModelBlock, PdlBlock, PdlParser, PdlUsage, PythonCodeBlock,
+    ReadBlock, RepeatBlock, Role, StringOrBoolean, TextBlock,
 };
 
 type Context = Vec<ChatMessage>;
@@ -82,8 +82,9 @@ impl<'a> Interpreter<'a> {
         let (messages, trace) = match program {
             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::If(block) => self.run_if(block, context).await,
             PdlBlock::Model(block) => self.run_model(block, context).await,
+            PdlBlock::PythonCode(block) => self.run_python_code(block, context).await,
             PdlBlock::Read(block) => self.run_read(block, context).await,
             PdlBlock::Repeat(block) => self.run_repeat(block, context).await,
             PdlBlock::Text(block) => self.run_text(block, context).await,
@@ -238,6 +239,34 @@ impl<'a> Interpreter<'a> {
         res
     }
 
+    // Run a PdlBlock::Call
+    async fn run_if(&mut self, block: &IfBlock, context: Context) -> Interpretation {
+        if self.debug {
+            eprintln!("If {:?}({:?})", block.condition, block.then);
+        }
+
+        self.extend_scope_with_block_map(&block.defs);
+
+        let cond = match &block.condition {
+            StringOrBoolean::Boolean(b) => Value::Bool(*b),
+            StringOrBoolean::String(s) => self.eval::<Value>(s)?,
+        };
+        let res = match cond {
+            Value::Bool(true) => self.run_quiet(&block.then, context).await,
+            Value::Bool(false) => match &block.else_ {
+                Some(else_block) => self.run_quiet(&else_block, context).await,
+                None => Ok((vec![], PdlBlock::If(block.clone()))),
+            },
+            x => Err(Box::from(format!(
+                "if block condition evaluated to non-boolean value: {:?}",
+                x
+            ))),
+        };
+
+        self.scope.pop();
+        res
+    }
+
     fn to_ollama_model_options(
         &self,
         maybe_parameters: &Option<HashMap<String, Value>>,

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

@@ -9,7 +9,7 @@ mod tests {
         interpreter::{run_json_sync as run_json, run_sync as run},
     };
 
-    use ollama_rs::generation::chat::{ChatMessage, MessageRole};
+    use ollama_rs::generation::chat::MessageRole;
 
     const DEFAULT_MODEL: &'static str = "ollama/granite3.2:2b";
 
@@ -294,4 +294,51 @@ mod tests {
         assert_eq!(messages[2].content, "6c");
         Ok(())
     }
+
+    #[test]
+    fn text_if_true() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "if": true,
+            "then": "good"
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "good");
+        Ok(())
+    }
+
+    #[test]
+    fn text_if_false() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "if": false,
+            "then": "bug",
+            "else": "good"
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "good");
+        Ok(())
+    }
+
+    #[test]
+    fn text_if_with_defs() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "defs": {
+                "x": 5
+            },
+            "if": "${x!=5}",
+            "then": "bug",
+            "else": "good"
+        });
+
+        let (messages, _) = run_json(program, false)?;
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "good");
+        Ok(())
+    }
 }

pdl-live-react/src-tauri/tests/cli/if1.pdl

@@ -0,0 +1,2 @@
+if: true
+then: hi

pdl-live-react/src-tauri/tests/cli/if2.pdl

@@ -0,0 +1,5 @@
+defs:
+  x: 5
+if: ${x!=5}
+then: bug
+else: good