Merge branch 'main' of github.com:andthattoo/ollama-workflows

src/lib.rs

@@ -36,3 +36,5 @@ pub use program::{
     models::{Model, ModelProvider},
     workflow::Workflow,
 };
+
+pub use program::atomics::{MessageInput, Task};

src/program/atomics.rs

@@ -1,6 +1,6 @@
 use crate::program::io::{Input, InputValue, Output};
 use crate::ProgramMemory;
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
 
@@ -77,9 +77,12 @@ pub enum Operator {
     End,
 }
 
-#[derive(Clone, Debug, Deserialize)]
+/// A message entry.
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct MessageInput {
+    /// Role, usually `user`, `assistant` or `system`.
     pub role: String,
+    /// Message content.
     pub content: String,
 }
 
@@ -103,6 +106,24 @@ pub struct Task {
     pub schema: Option<String>,
 }
 
+impl Task {
+    /// Creates a new chat history entry with the given content for `assistant` role.
+    pub fn append_assistant_message(&mut self, content: impl Into<String>) {
+        self.messages.push(MessageInput {
+            role: "assistant".to_string(),
+            content: content.into(),
+        });
+    }
+
+    /// Creates a new chat history entry with the given content for `user` role.
+    pub fn append_user_message(&mut self, content: impl Into<String>) {
+        self.messages.push(MessageInput {
+            role: "user".to_string(),
+            content: content.into(),
+        });
+    }
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(untagged)]
 pub enum TaskOutputInput {

src/program/models.rs

@@ -330,7 +330,7 @@ impl fmt::Display for ModelProvider {
 mod tests {
     use super::*;
 
-    const MODEL_NAME: &str = "phi3:3.8b";
+    const MODEL_NAME: &str = "phi3.5:3.8b";
     const PROVIDER_NAME: &str = "openai";
     #[test]
     fn test_model_string_conversion() {

src/program/workflow.rs

@@ -100,6 +100,10 @@ impl Workflow {
     pub fn get_tasks(&self) -> &Vec<Task> {
         &self.tasks
     }
+    /// Returns a mutable reference to the tasks of the workflow.
+    pub fn get_tasks_mut(&mut self) -> &mut Vec<Task> {
+        &mut self.tasks
+    }
     /// Returns a reference to the steps of the workflow.
     pub fn get_workflow(&self) -> &Vec<Edge> {
         &self.steps
@@ -109,8 +113,8 @@ impl Workflow {
         &self.return_value
     }
     /// Returns a reference to the task at the specified index.
-    pub fn get_step(&self, index: u32) -> Option<&Edge> {
-        self.steps.get(index as usize)
+    pub fn get_step(&self, index: usize) -> Option<&Edge> {
+        self.steps.get(index)
     }
     /// Returns a reference to the step for specified task_id.
     pub fn get_step_by_id(&self, task_id: &str) -> Option<&Edge> {
@@ -120,4 +124,8 @@ impl Workflow {
     pub fn get_tasks_by_id(&self, task_id: &str) -> Option<&Task> {
         self.tasks.iter().find(|task| task.id == task_id)
     }
+    /// Returns a mutable reference to the task at the specified task_id.
+    pub fn get_tasks_by_id_mut(&mut self, task_id: &str) -> Option<&mut Task> {
+        self.tasks.iter_mut().find(|task| task.id == task_id)
+    }
 }

tests/task_test.rs

@@ -0,0 +1,15 @@
+use ollama_workflows::Workflow;
+
+#[test]
+fn test_task_mutable() {
+    let mut workflow = Workflow::new_from_json("./tests/test_workflows/search.json").unwrap();
+
+    let task_id = "E";
+    let task = workflow.get_tasks_by_id_mut(task_id).unwrap();
+    assert_eq!(task.id, task_id);
+
+    assert_eq!(task.messages.len(), 1);
+    task.append_assistant_message("This is your response.");
+    task.append_user_message("Thank you.");
+    assert_eq!(task.messages.len(), 3);
+}

tests/test_workflows/search.json

@@ -11,7 +11,12 @@
       "id": "A",
       "name": "Web Search Query",
       "description": "Write a web search query to collect useful information for the given question",
-      "messages": [{"role": "user", "content":"You are tasked with generating a single search query to gather useful information for answering a given question. Your goal is to create a concise, clear, and creative query that avoids repeating previously asked questions.\n\nHere is the query you need to address:\n<query>\n{{query}}\n</query>\n\nHere are the previous questions that have been asked:\n<history>\n{{history}}\n</history>\n\nWhen creating your search query, follow these guidelines:\n1. Analyze the given query and identify key concepts or terms.\n2. Review the previous questions to avoid repetition.\n3. Be creative in your approach, considering alternative phrasings or related concepts.\n4. Keep the query concise and clear, typically no more than 5-7 words.\n5. Focus on gathering information that will be most useful in answering the original query.\n\nProvide your search query inside <search_query> tags. Do not include any explanation or additional text outside these tags."}],
+      "messages": [
+        {
+          "role": "user",
+          "content": "You are tasked with generating a single search query to gather useful information for answering a given question. Your goal is to create a concise, clear, and creative query that avoids repeating previously asked questions.\n\nHere is the query you need to address:\n<query>\n{{query}}\n</query>\n\nHere are the previous questions that have been asked:\n<history>\n{{history}}\n</history>\n\nWhen creating your search query, follow these guidelines:\n1. Analyze the given query and identify key concepts or terms.\n2. Review the previous questions to avoid repetition.\n3. Be creative in your approach, considering alternative phrasings or related concepts.\n4. Keep the query concise and clear, typically no more than 5-7 words.\n5. Focus on gathering information that will be most useful in answering the original query.\n\nProvide your search query inside <search_query> tags. Do not include any explanation or additional text outside these tags."
+        }
+      ],
       "inputs": [
         {
           "name": "query",
@@ -48,7 +53,7 @@
       "id": "B",
       "name": "Web Search",
       "description": "Search the web with the given query",
-      "messages": [{"role": "user", "content": "{query}"}],
+      "messages": [{ "role": "user", "content": "{query}" }],
       "inputs": [
         {
           "name": "query",
@@ -77,7 +82,12 @@
       "id": "C",
       "name": "Candidate Website",
       "description": "Pick the most useful link from web search results for the given query",
-      "messages": [{"role": "user", "content":"Pick the most useful URL to scrape information for Query: {query} \n\n ###Sarch Results:{web_results} \n\n Only output the selected URL: ###Selected URL:"}],
+      "messages": [
+        {
+          "role": "user",
+          "content": "Pick the most useful URL to scrape information for Query: {query} \n\n ###Sarch Results:{web_results} \n\n Only output the selected URL: ###Selected URL:"
+        }
+      ],
       "inputs": [
         {
           "name": "web_results",
@@ -109,7 +119,7 @@
       "id": "D",
       "name": "Scrape Website",
       "description": "Scrape the selected website",
-      "messages": [{"role": "user", "content":"scrape {search_url}"}],
+      "messages": [{ "role": "user", "content": "scrape {search_url}" }],
       "inputs": [
         {
           "name": "search_url",
@@ -134,7 +144,12 @@
       "id": "E",
       "name": "Summarize Website",
       "description": "Summarize website content",
-      "messages": [{"role": "user", "content":"Answer given question completely based on following context. Don't assume anything. Use the provided information to answer the question. If context is not enough, say 'I dont know.' ###Content: {content} \n\n ###Question: {query} \n\n ###Answer:"}],
+      "messages": [
+        {
+          "role": "user",
+          "content": "Answer given question completely based on following context. Don't assume anything. Use the provided information to answer the question. If context is not enough, say 'I dont know.' ###Content: {content} \n\n ###Question: {query} \n\n ###Answer:"
+        }
+      ],
       "inputs": [
         {
           "name": "content",
@@ -167,7 +182,12 @@
       "id": "F",
       "name": "Evaluate Result",
       "description": "Evaluate if the result is satisfying",
-      "messages": [{"role": "user", "content":"Evaluate the result if it answers the query. Write 'Yes' if valid, if not 'No'.Write nothing else but strictly 'Yes'|'No'. query:{query} result:{answer}. ###Evaluation:"}],
+      "messages": [
+        {
+          "role": "user",
+          "content": "Evaluate the result if it answers the query. Write 'Yes' if valid, if not 'No'.Write nothing else but strictly 'Yes'|'No'. query:{query} result:{answer}. ###Evaluation:"
+        }
+      ],
       "inputs": [
         {
           "name": "answer",
@@ -199,7 +219,7 @@
       "id": "__end",
       "name": "end",
       "description": "End of the task",
-      "messages": [{"role": "user", "content": "End of the task"}],
+      "messages": [{ "role": "user", "content": "End of the task" }],
       "inputs": [],
       "operator": "end",
       "outputs": []
@@ -233,14 +253,14 @@
       "source": "F",
       "target": "end",
       "condition": {
-        "input":{
-            "type":"read",
-            "key": "final_result"
+        "input": {
+          "type": "read",
+          "key": "final_result"
         },
         "expression": "Equal",
         "expected": "Yes",
         "target_if_not": "A"
-    }
+      }
     }
   ],
   "return_value": {