messageblock and arrayblock support

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

Author: starpit

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

@@ -12,8 +12,9 @@ use serde_json::{from_reader, json, to_string, Map, Value};
 use tempfile::Builder;
 
 use crate::pdl::ast::{
-    CallBlock, FunctionBlock, ListOrString, MessageBlock, ModelBlock, ObjectBlock, PdlBaseType,
-    PdlBlock, PdlOptionalType, PdlParser, PdlType, PythonCodeBlock, RepeatBlock, Role, TextBlock,
+    ArrayBlock, CallBlock, FunctionBlock, ListOrString, MessageBlock, ModelBlock, ObjectBlock,
+    PdlBaseType, PdlBlock, PdlOptionalType, PdlParser, PdlType, PythonCodeBlock, RepeatBlock, Role,
+    TextBlock,
 };
 use crate::pdl::pip::pip_install_if_needed;
 use crate::pdl::requirements::BEEAI_FRAMEWORK;
@@ -198,7 +199,7 @@ fn call_tools(model: &String, parameters: &HashMap<String, Value>) -> PdlBlock {
         text: vec![PdlBlock::Model(
             ModelBlock::new(model.as_str())
                 .parameters(&strip_nulls(parameters))
-                .input(PdlBlock::Array {
+                .input(PdlBlock::Array(ArrayBlock {
                     array: vec![PdlBlock::Message(MessageBlock {
                         role: Role::Tool,
                         description: None,
@@ -214,7 +215,7 @@ fn call_tools(model: &String, parameters: &HashMap<String, Value>) -> PdlBlock {
                             args: Some("${ args }".into()), // invoke with arguments as specified by the model
                         })),
                     })],
-                })
+                }))
                 .build(),
         )],
     });

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

@@ -164,7 +164,7 @@ pub struct ModelBlock {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub parameters: Option<HashMap<String, Value>>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub input: Option<Box<PdlBlock>>, // really this should be restricted to be PdlBlock::Array; how do we do this in rust?
+    pub input: Option<Box<PdlBlock>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "modelResponse")]
     pub model_response: Option<String>,
@@ -346,6 +346,13 @@ pub struct IfBlock {
     pub defs: Option<HashMap<String, PdlBlock>>,
 }
 
+/// Return the array of values computed by each block of the list of blocks
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ArrayBlock {
+    /// Elements of the array
+    pub array: Vec<PdlBlock>,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(untagged)]
 pub enum PdlBlock {
@@ -355,7 +362,7 @@ pub enum PdlBlock {
     If(IfBlock),
     Object(ObjectBlock),
     Call(CallBlock),
-    Array { array: Vec<PdlBlock> },
+    Array(ArrayBlock),
     Message(MessageBlock),
     Repeat(RepeatBlock),
     Text(TextBlock),

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

@@ -26,8 +26,9 @@ use serde_json::{from_str, to_string, Number, Value};
 use serde_norway::{from_reader, from_str as from_yaml_str};
 
 use crate::pdl::ast::{
-    CallBlock, FunctionBlock, IfBlock, ListOrString, ModelBlock, ObjectBlock, PdlBlock, PdlParser,
-    PdlUsage, PythonCodeBlock, ReadBlock, RepeatBlock, Role, StringOrBoolean, TextBlock,
+    ArrayBlock, CallBlock, FunctionBlock, IfBlock, ListOrString, MessageBlock, ModelBlock,
+    ObjectBlock, PdlBlock, PdlParser, PdlUsage, PythonCodeBlock, ReadBlock, RepeatBlock, Role,
+    StringOrBoolean, TextBlock,
 };
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -48,23 +49,11 @@ pub enum PdlResult {
     Dict(HashMap<String, PdlResult>),
 }
 impl ::std::fmt::Display for PdlResult {
-    // This trait requires `fmt` with this exact signature.
     fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
         let s = to_string(&self).unwrap(); // TODO: .map_err(|e| e.to_string())?;
         write!(f, "{}", s)
     }
 }
-/*impl From<&Value> for PdlResult {
-    fn from(v: &Value) -> Self {
-        match v {
-            Value::Bool(b) => PdlResult::Bool(*b),
-            Value::String(s) => PdlResult::String(s.clone()),
-            Value::Number(n) => PdlResult::Number(n.clone()),
-            Value::Array(a) => PdlResult::List(a.into_iter().map(|v| v.into())),
-            Value::Object(m) => PdlResult::Dict(m.into_iter().map(|(k,v)| (k, v.into()))),
-        }
-    }
-}*/
 impl From<&str> for PdlResult {
     fn from(s: &str) -> Self {
         PdlResult::String(s.to_string())
@@ -93,7 +82,6 @@ struct Interpreter<'a> {
     cwd: PathBuf,
     // id_stack: Vec<String>,
     jinja_env: Environment<'a>,
-    // rt: Runtime,
     scope: Vec<Scope>,
     debug: bool,
     emit: bool,
@@ -116,7 +104,6 @@ impl<'a> Interpreter<'a> {
             cwd: current_dir().unwrap_or(PathBuf::from("/")),
             // id_stack: vec![],
             jinja_env: jinja_env,
-            // rt: Runtime::new().unwrap(),
             scope: vec![Scope::new()],
             debug: false,
             emit: true,
@@ -160,14 +147,15 @@ impl<'a> Interpreter<'a> {
             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,
+            PdlBlock::Array(block) => self.run_array(block, context).await,
+            PdlBlock::Message(block) => self.run_message(block, context).await,
             _ => Err(Box::from(format!("Unsupported block {:?}", program))),
         }?;
 
         if match program {
-            PdlBlock::Call(_) | PdlBlock::Text(_) => false,
+            PdlBlock::Call(_) | PdlBlock::Model(_) | PdlBlock::Text(_) => false,
             _ => self.emit,
         } {
-            // eprintln!("{:?}", program);
             println!("{}", pretty_print(&messages));
         }
         self.emit = prior_emit;
@@ -204,6 +192,17 @@ impl<'a> Interpreter<'a> {
         }))
     }
 
+    /// Evaluate String as a Jinja2 expression, expecting a string in response
+    fn eval_to_string(&self, expr: &String) -> Result<String, PdlError> {
+        match self.eval(expr)? {
+            PdlResult::String(s) => Ok(s),
+            x => Err(Box::from(format!(
+                "Expression {expr} evaluated to non-string {:?}",
+                x
+            ))),
+        }
+    }
+
     fn eval_complex(&self, expr: &Value) -> Result<PdlResult, PdlError> {
         match expr {
             Value::Null => Ok("".into()),
@@ -499,7 +498,7 @@ impl<'a> Interpreter<'a> {
                     println!("Model options {:?}", options);
                 }
 
-                let messages = match &block.input {
+                let input_messages = match &block.input {
                     Some(input) => {
                         // TODO ignoring result, trace
                         let (_result, messages, _trace) = self.run_quiet(&*input, context).await?;
@@ -508,7 +507,7 @@ impl<'a> Interpreter<'a> {
                     None => context,
                 };
                 let (prompt, history_slice): (&ChatMessage, &[ChatMessage]) =
-                    match messages.split_last() {
+                    match input_messages.split_last() {
                         Some(x) => x,
                         None => (&ChatMessage::user("".into()), &[]),
                     };
@@ -523,6 +522,10 @@ impl<'a> Interpreter<'a> {
                     );
                 }
 
+                if self.emit {
+                    println!("{}", pretty_print(&input_messages));
+                }
+
                 let req = ChatMessageRequest::new(model.into(), vec![prompt.clone()])
                     .options(options)
                     .tools(tools);
@@ -571,6 +574,7 @@ impl<'a> Interpreter<'a> {
                     response_string += res.message.content.as_str();
                     last_res = Some(res);
                 }
+                stdout.write_all(b"\n").await?;
 
                 let mut trace = block.clone();
                 trace.pdl_result = Some(response_string.clone());
@@ -584,11 +588,10 @@ impl<'a> Interpreter<'a> {
                             completion_nanos: usage.eval_duration,
                         });
                     }
-                    let mut message = res.message.clone();
-                    message.content = response_string;
+                    let output_messages = vec![ChatMessage::assistant(response_string)];
                     Ok((
-                        message.content.clone().into(),
-                        vec![message],
+                        res.message.content.into(),
+                        output_messages,
                         PdlBlock::Model(trace),
                     ))
                 } else {
@@ -673,7 +676,7 @@ impl<'a> Interpreter<'a> {
         match role {
             Role::User => MessageRole::User,
             Role::Assistant => MessageRole::Assistant,
-            Role::System => MessageRole::Assistant,
+            Role::System => MessageRole::System,
             Role::Tool => MessageRole::Tool,
         }
     }
@@ -786,6 +789,70 @@ impl<'a> Interpreter<'a> {
             PdlBlock::Text(trace),
         ))
     }
+
+    async fn run_array(&mut self, block: &ArrayBlock, context: Context) -> Interpretation {
+        let mut result_items = vec![];
+        let mut all_messages = vec![];
+        let mut trace_items = vec![];
+
+        let mut iter = block.array.iter();
+        while let Some(item) = iter.next() {
+            // TODO accumulate messages
+            let (result, messages, trace) = self.run_quiet(item, context.clone()).await?;
+            result_items.push(result);
+            all_messages.extend(messages);
+            trace_items.push(trace);
+        }
+
+        Ok((
+            PdlResult::List(result_items),
+            all_messages,
+            PdlBlock::Array(ArrayBlock { array: trace_items }),
+        ))
+    }
+
+    async fn run_message(&mut self, block: &MessageBlock, context: Context) -> Interpretation {
+        let (content_result, content_messages, content_trace) =
+            self.run(&block.content, context).await?;
+        let name = if let Some(name) = &block.name {
+            Some(self.eval_to_string(&name)?)
+        } else {
+            None
+        };
+        let tool_call_id = if let Some(tool_call_id) = &block.tool_call_id {
+            Some(self.eval_to_string(&tool_call_id)?)
+        } else {
+            None
+        };
+
+        let mut dict: HashMap<String, PdlResult> = HashMap::new();
+        dict.insert("role".into(), PdlResult::String(to_string(&block.role)?));
+        dict.insert("content".into(), content_result);
+        if let Some(name) = &name {
+            dict.insert("name".into(), PdlResult::String(name.clone()));
+        }
+        if let Some(tool_call_id) = &tool_call_id {
+            dict.insert(
+                "tool_call_id".into(),
+                PdlResult::String(tool_call_id.clone()),
+            );
+        }
+
+        Ok((
+            PdlResult::Dict(dict),
+            content_messages
+                .into_iter()
+                .map(|m| ChatMessage::new(self.to_ollama_role(&block.role), m.content))
+                .collect(),
+            PdlBlock::Message(MessageBlock {
+                role: block.role.clone(),
+                content: Box::new(content_trace),
+                description: block.description.clone(),
+                name: name,
+                tool_call_id: tool_call_id,
+            }),
+        ))
+    }
 }
 
 pub async fn run(program: &PdlBlock, cwd: Option<PathBuf>, debug: bool) -> Interpretation {

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

@@ -23,7 +23,7 @@ mod tests {
     }
 
     #[test]
-    fn single_model_via_input() -> Result<(), Box<dyn Error>> {
+    fn single_model_via_input_string() -> Result<(), Box<dyn Error>> {
         let (_, messages, _) = run(
             &PdlBlock::Model(ModelBlock::new(DEFAULT_MODEL).input_str("hello").build()),
             None,
@@ -54,6 +54,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn single_model_via_input_array() -> Result<(), Box<dyn Error>> {
+        let (_, messages, _) = run_json(
+            json!({
+                "model": DEFAULT_MODEL,
+                "input": {
+                    "array": [
+                        { "role": "system", "content": "answer as if you live in europe" },
+                        { "role": "user", "content": "what is the fastest animal where you live?" },
+                    ]
+                }
+            }),
+            false,
+        )?;
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].role, MessageRole::Assistant);
+        let m = messages[0].content.to_lowercase();
+        assert!(
+            m.contains("pronghorn")
+                || m.contains("falcon")
+                || m.contains("bison")
+                || m.contains("native")
+        );
+        Ok(())
+    }
+
     #[test]
     fn two_models_via_text_chain() -> Result<(), Box<dyn Error>> {
         let (_, messages, _) = run_json(

pdl-live-react/src-tauri/tests/cli/model-input-array.pdl

@@ -0,0 +1,7 @@
+model: ollama/granite3.2:2b
+input:
+  array:
+    - role: system
+      content: answer as if you live in europe
+    - role: user
+      content: what is the fastest animal where i live?

pdl-live-react/src-tauri/tests/cli/model-input-string.pdl

@@ -0,0 +1,2 @@
+model: ollama/granite3.2:2b
+input: what is the fastest animal?