some fixes and better interface

Author: erhant

Cargo.lock

@@ -21,3 +21,6 @@ eyre.workspace = true
 
 # ollama-rs is re-exported from ollama-workflows
 ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows" }
+
+[dev-dependencies]
+env_logger.workspace = true

workflows/Cargo.toml

@@ -5,3 +5,15 @@ This crate handles the configurations of models to be used, and implements vario
 
 - **OpenAI**: We check that the chosen models are enabled for the user's profile by fetching their models with their API key. We filter out the disabled models.
 - **Ollama**: We provide a sample workflow to measure TPS and then pick models that are above some TPS threshold. While calculating TPS, there is also a timeout so that beyond that timeout the TPS is not even considered and the model becomes invalid.
+
+## Environment Variables
+
+DKN Workflows make use of several environment variables, respecting the providers.
+
+- `OPENAI_API_KEY` is used for OpenAI requests
+- `OLLAMA_HOST` is used to connect to Ollama server
+- `OLLAMA_PORT` is used to connect to Ollama server
+- `OLLAMA_AUTO_PULL` indicates whether we should pull missing models automatically or not
+
+SERPER_API_KEY=
+JINA_API_KEY=

workflows/README.md

@@ -15,27 +15,29 @@ pub struct ModelConfig {
 }
 
 impl ModelConfig {
-    /// Parses Ollama-Workflows compatible models from a comma-separated values string.
-    pub fn new_from_csv(input: Option<String>) -> Self {
-        let models_str = split_comma_separated(input);
-
-        let models = models_str
+    pub fn new(models: Vec<Model>) -> Self {
+        let models_and_providers = models
             .into_iter()
-            .filter_map(|s| match Model::try_from(s) {
-                Ok(model) => Some((model.clone().into(), model)),
-                Err(e) => {
-                    log::warn!("Error parsing model: {}", e);
-                    None
-                }
-            })
+            .map(|model| (model.clone().into(), model))
             .collect::<Vec<_>>();
 
         Self {
-            models,
+            models: models_and_providers,
             openai: OpenAIConfig::new(),
             ollama: OllamaConfig::new(),
         }
     }
+    /// Parses Ollama-Workflows compatible models from a comma-separated values string.
+    pub fn new_from_csv(input: &str) -> Self {
+        let models_str = split_comma_separated(input);
+
+        let models = models_str
+            .into_iter()
+            .filter_map(|s| Model::try_from(s).ok())
+            .collect();
+
+        Self::new(models)
+    }
 
     /// Returns the models that belong to a given providers from the config.
     pub fn get_models_for_provider(&self, provider: ModelProvider) -> Vec<Model> {
@@ -209,19 +211,18 @@ mod tests {
 
     #[test]
     fn test_csv_parser() {
-        let cfg =
-            ModelConfig::new_from_csv(Some("idontexist,i dont either,i332287648762".to_string()));
+        let cfg = ModelConfig::new_from_csv("idontexist,i dont either,i332287648762");
         assert_eq!(cfg.models.len(), 0);
 
-        let cfg = ModelConfig::new_from_csv(Some(
-            "gemma2:9b-instruct-q8_0,phi3:14b-medium-4k-instruct-q4_1,balblablabl".to_string(),
-        ));
+        let cfg = ModelConfig::new_from_csv(
+            "gemma2:9b-instruct-q8_0,phi3:14b-medium-4k-instruct-q4_1,balblablabl",
+        );
         assert_eq!(cfg.models.len(), 2);
     }
 
     #[test]
     fn test_model_matching() {
-        let cfg = ModelConfig::new_from_csv(Some("gpt-4o,llama3.1:latest".to_string()));
+        let cfg = ModelConfig::new_from_csv("gpt-4o,llama3.1:latest");
         assert_eq!(
             cfg.get_matching_model("openai".to_string()).unwrap().1,
             Model::GPT4o,
@@ -250,7 +251,7 @@ mod tests {
 
     #[test]
     fn test_get_any_matching_model() {
-        let cfg = ModelConfig::new_from_csv(Some("gpt-3.5-turbo,llama3.1:latest".to_string()));
+        let cfg = ModelConfig::new_from_csv("gpt-3.5-turbo,llama3.1:latest");
         let result = cfg.get_any_matching_model(vec![
             "i-dont-exist".to_string(),
             "llama3.1:latest".to_string(),

workflows/src/config.rs

@@ -2,7 +2,7 @@ mod utils;
 pub use utils::*;
 
 mod providers;
-pub use providers::*;
+use providers::*;
 
 mod config;
 pub use config::ModelConfig;

workflows/src/lib.rs

@@ -227,7 +227,7 @@ mod tests {
     use ollama_workflows::{Executor, Model, ProgramMemory, Workflow};
 
     #[tokio::test]
-    #[ignore = "run this manually"]
+    #[ignore = "requires Ollama"]
     async fn test_ollama_prompt() {
         let model = Model::default().to_string();
         let ollama = Ollama::default();
@@ -246,7 +246,7 @@ mod tests {
     }
 
     #[tokio::test]
-    #[ignore = "run this manually"]
+    #[ignore = "requires Ollama"]
     async fn test_ollama_workflow() {
         let workflow = r#"{
         "name": "Simple",

workflows/src/providers/ollama.rs

@@ -1,5 +1,6 @@
 use eyre::{eyre, Context, Result};
 use ollama_workflows::Model;
+use reqwest::Client;
 use serde::Deserialize;
 
 const OPENAI_MODELS_API: &str = "https://api.openai.com/v1/models";
@@ -29,6 +30,7 @@ struct OpenAIModelsResponse {
 
 #[derive(Debug, Clone, Default)]
 pub struct OpenAIConfig {
+    /// API key, if available.
     pub(crate) api_key: Option<String>,
 }
 
@@ -50,17 +52,17 @@ impl OpenAIConfig {
         };
 
         // fetch models
-        let client = reqwest::Client::new();
+        let client = Client::new();
         let request = client
             .get(OPENAI_MODELS_API)
             .header("Authorization", format!("Bearer {}", api_key))
             .build()
-            .wrap_err("Failed to build request")?;
+            .wrap_err("failed to build request")?;
 
         let response = client
             .execute(request)
             .await
-            .wrap_err("Failed to send request")?;
+            .wrap_err("failed to send request")?;
 
         // parse response
         if response.status().is_client_error() {

workflows/src/providers/openai.rs

@@ -2,33 +2,35 @@
 ///
 /// - Trims `"` from both ends at the start
 /// - For each item, trims whitespace from both ends
-pub fn split_comma_separated(input: Option<String>) -> Vec<String> {
-    match input {
-        Some(s) => s
-            .trim_matches('"')
-            .split(',')
-            .filter_map(|s| {
-                let s = s.trim().to_string();
-                if s.is_empty() {
-                    None
-                } else {
-                    Some(s)
-                }
-            })
-            .collect::<Vec<_>>(),
-        None => vec![],
-    }
+pub fn split_comma_separated(input: &str) -> Vec<String> {
+    input
+        .trim_matches('"')
+        .split(',')
+        .filter_map(|s| {
+            let s = s.trim().to_string();
+            if s.is_empty() {
+                None
+            } else {
+                Some(s)
+            }
+        })
+        .collect::<Vec<_>>()
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_split_comma_separated() {
+    fn test_example() {
         // should ignore whitespaces and `"` at both ends, and ignore empty items
-        let input = Some("\"a,    b , c ,,  \"".to_string());
+        let input = "\"a,    b , c ,,  \"";
         let expected = vec!["a".to_string(), "b".to_string(), "c".to_string()];
         assert_eq!(split_comma_separated(input), expected);
     }
+
+    #[test]
+    fn test_empty() {
+        assert!(split_comma_separated(Default::default()).is_empty());
+    }
 }

workflows/src/utils.rs

@@ -1,14 +1,38 @@
+use std::env;
+
 use dkn_workflows::ModelConfig;
 use eyre::Result;
+use ollama_workflows::Model;
+
+#[tokio::test]
+#[ignore = "requires Ollama"]
+async fn test_ollama() -> Result<()> {
+    env::set_var("RUST_LOG", "none,dkn_workflows=debug");
+    let _ = env_logger::try_init();
+
+    let models = vec![Model::Phi3_5Mini];
+    let mut model_config = ModelConfig::new(models);
+
+    model_config.check_services().await
+}
+
+#[tokio::test]
+async fn test_openai() -> Result<()> {
+    env::set_var("RUST_LOG", "debug");
+    let _ = env_logger::try_init();
+
+    let models = vec![Model::GPT4Turbo];
+    let mut model_config = ModelConfig::new(models);
+
+    model_config.check_services().await
+}
 
-// #[tokio::test]
-// async fn test_ollama() -> Result<()> {}
+#[tokio::test]
+async fn test_empty() -> Result<()> {
+    let mut model_config = ModelConfig::new(vec![]);
 
-// #[tokio::test]
-// async fn test_openai() -> Result<()> {}
+    let result = model_config.check_services().await;
+    assert!(result.is_err());
 
-// #[tokio::test]
-// async fn test_empty() -> Result<()> {
-//     let mut model_config = ModelConfig::default();
-//     model_config.check_services().await
-// }
+    Ok(())
+}