more tests

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

Author: starpit

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

@@ -327,6 +327,14 @@ pub fn run_file(source_file_path: &str, debug: bool) -> Interpretation {
     run(&from_reader(File::open(source_file_path)?)?, debug)
 }
 
+pub fn run_string(source: &str, debug: bool) -> Interpretation {
+    run(&from_str(source)?, debug)
+}
+
+pub fn run_json(source: Value, debug: bool) -> Interpretation {
+    run_string(&to_string(&source)?, debug)
+}
+
 fn pretty_print(messages: &Vec<ChatMessage>) {
     messages
         .into_iter()

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

@@ -2,10 +2,11 @@
 mod tests {
     //    use super::*;
     use ::std::error::Error;
+    use serde_json::json;
 
     use crate::pdl::{
         ast::{PdlBlock, PdlModelBlock, PdlParser, PdlTextBlock},
-        interpreter::run,
+        interpreter::{run, run_json},
     };
 
     use ollama_rs::generation::chat::MessageRole;
@@ -35,42 +36,67 @@ mod tests {
 
     #[test]
     fn single_model_via_text_chain() -> Result<(), Box<dyn Error>> {
-        let (messages, _) = run(
-            &PdlBlock::Text(
-                vec![
-                    "hello".into(),
-                    PdlBlock::Model(PdlModelBlock::new(DEFAULT_MODEL).build()),
+        let (messages, _) = run_json(
+            json!({
+                "text": [
+                    "hello",
+                    { "model": DEFAULT_MODEL }
                 ]
-                .into(),
-            ),
+            }),
             false,
         )?;
         assert_eq!(messages.len(), 2);
         assert_eq!(messages[0].role, MessageRole::User);
-        assert!(messages[0].content.contains("hello"));
+        assert_eq!(messages[0].content, "hello");
         assert_eq!(messages[1].role, MessageRole::Assistant);
         assert!(messages[1].content.contains("Hello!"));
         Ok(())
     }
 
+    #[test]
+    fn two_models_via_text_chain() -> Result<(), Box<dyn Error>> {
+        let (messages, _) = run_json(
+            json!({
+                "text": [
+                    "what is the fastest animal?",
+                    { "model": DEFAULT_MODEL },
+                    "in europe?",
+                    { "model": DEFAULT_MODEL },
+                ]
+            }),
+            false,
+        )?;
+        assert_eq!(messages.len(), 4);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "what is the fastest animal?");
+        assert_eq!(messages[1].role, MessageRole::Assistant);
+        assert!(messages[1].content.contains("cheetah"));
+        assert_eq!(messages[2].role, MessageRole::User);
+        assert_eq!(messages[2].content, "in europe?");
+        assert_eq!(messages[3].role, MessageRole::Assistant);
+        assert!(
+            messages[3].content.contains("peregrine")
+                || messages[3].content.contains("bison")
+                || messages[3].content.contains("hare")
+                || messages[3].content.contains("Eagle")
+                || messages[3].content.contains("Greyhound")
+                || messages[3].content.contains("gazelle")
+                || messages[3].content.contains("lynx")
+        );
+        Ok(())
+    }
+
     #[test]
     fn text_parser_json() -> Result<(), Box<dyn Error>> {
         let json = "{\"key\": \"value\"}";
-        let program = PdlBlock::Text(
-            vec![
-                PdlBlock::Text(
-                    PdlTextBlock::new(vec![json.into()])
-                        .def(&"foo")
-                        .parser(PdlParser::Json)
-                        .build(),
-                ),
-                "${ foo.key }".into(),
+        let program = json!({
+            "text": [
+                { "def": "foo", "parser": "json", "text":  [json] },
+                "${ foo.key }"
             ]
-            .into(),
-        );
-        println!("{}", serde_json::to_string(&program)?);
+        });
 
-        let (messages, _) = run(&program, false)?;
+        let (messages, _) = run_json(program, false)?;
         assert_eq!(messages.len(), 2);
         assert_eq!(messages[0].role, MessageRole::User);
         assert_eq!(messages[0].content, json);