use Model instead of (ModelProvider, Model)

Author: erhant

Cargo.lock

@@ -90,8 +90,8 @@ docker compose --profile=ollama-rocm up
 Note that we are very dependent on Ollama packages, and it is important to check their versions if relevant:
 
 ```sh
-@cat Cargo.lock | grep "https://github.com/andthattoo/ollama-workflows"
-@cat Cargo.lock | grep "https://github.com/andthattoo/ollama-rs"
+cat Cargo.lock | grep "https://github.com/andthattoo/ollama-workflows"
+cat Cargo.lock | grep "https://github.com/andthattoo/ollama-rs"
 ```
 
 ### Testing

README.md

@@ -88,13 +88,7 @@ async fn main() -> Result<()> {
     }?;
     log::warn!(
         "Using models: {}",
-        config
-            .workflows
-            .models
-            .iter()
-            .map(|(p, m)| format!("{}/{}", p, m))
-            .collect::<Vec<_>>()
-            .join(", ")
+        config.workflows.get_model_names().join(", ")
     );
 
     // check network-specific configurations

compute/src/main.rs

@@ -47,13 +47,7 @@ impl DriaComputeNode {
         // print models
         diagnostics.push(format!(
             "Models: {}",
-            self.config
-                .workflows
-                .models
-                .iter()
-                .map(|(p, m)| format!("{}/{}", p, m))
-                .collect::<Vec<String>>()
-                .join(", ")
+            self.config.workflows.get_model_names().join(", ")
         ));
 
         // if we have not received pings for a while, we are considered offline

compute/src/node/diagnostic.rs

@@ -58,13 +58,7 @@ impl HeartbeatRequester {
         let heartbeat_request = HeartbeatRequest {
             heartbeat_id: uuid,
             deadline,
-            models: node
-                .config
-                .workflows
-                .models
-                .iter()
-                .map(|m| m.1.to_string())
-                .collect(),
+            models: node.config.workflows.get_model_names(),
             pending_tasks: node.get_pending_task_count(),
         };
 

compute/src/reqres/heartbeat.rs

@@ -65,15 +65,15 @@ impl TaskResponder {
         let task_public_key = PublicKey::parse_slice(&task_public_key_bytes, None)?;
 
         // read model / provider from the task
-        let (model_provider, model) = node
+        let model = node
             .config
             .workflows
             .get_any_matching_model(task.input.model)?;
         let model_name = model.to_string(); // get model name, we will pass it in payload
         log::info!("Using model {} for task {}", model_name, task.task_id);
 
         // prepare workflow executor
-        let (executor, batchable) = if model_provider == ModelProvider::Ollama {
+        let (executor, batchable) = if model.provider() == ModelProvider::Ollama {
             (
                 Executor::new_at(
                     model,

compute/src/reqres/task.rs

@@ -4,13 +4,15 @@ use crate::{
     Model, ModelProvider,
 };
 use dkn_utils::split_csv_line;
-use eyre::{eyre, Result};
+use eyre::{eyre, OptionExt, Result};
 use rand::seq::IteratorRandom; // provides Vec<_>.choose
 
 #[derive(Debug, Clone)]
 pub struct DriaWorkflowsConfig {
-    /// List of models with their providers.
-    pub models: Vec<(ModelProvider, Model)>,
+    /// List of models.
+    ///
+    /// You can do `model.provider()` to get its provider.
+    pub models: Vec<Model>,
     /// Ollama configurations, in case Ollama is used.
     /// Otherwise, can be ignored.
     pub ollama: OllamaConfig,
@@ -40,13 +42,8 @@ impl Default for DriaWorkflowsConfig {
 impl DriaWorkflowsConfig {
     /// Creates a new config with the given models.
     pub fn new(models: Vec<Model>) -> Self {
-        let models_and_providers = models
-            .into_iter()
-            .map(|model| (model.clone().into(), model))
-            .collect::<Vec<_>>();
-
         Self {
-            models: models_and_providers,
+            models,
             ollama: OllamaConfig::new(),
             openai: OpenAIConfig::new(),
             openrouter: OpenRouterConfig::new(),
@@ -84,24 +81,23 @@ impl DriaWorkflowsConfig {
     pub fn get_models_for_provider(&self, provider: ModelProvider) -> Vec<Model> {
         self.models
             .iter()
-            .filter_map(|(p, m)| {
-                if *p == provider {
-                    Some(m.clone())
-                } else {
-                    None
-                }
-            })
+            .filter(|m| m.provider() == provider)
+            .cloned()
             .collect()
     }
 
     /// Returns `true` if the configuration contains models that can be processed in parallel, e.g. API calls.
     pub fn has_batchable_models(&self) -> bool {
-        self.models.iter().any(|(p, _)| *p != ModelProvider::Ollama)
+        self.models
+            .iter()
+            .any(|m| m.provider() != ModelProvider::Ollama)
     }
 
     /// Returns `true` if the configuration contains a model that cant be run in parallel, e.g. a Ollama model.
     pub fn has_non_batchable_models(&self) -> bool {
-        self.models.iter().any(|(p, _)| *p == ModelProvider::Ollama)
+        self.models
+            .iter()
+            .any(|m| m.provider() == ModelProvider::Ollama)
     }
 
     /// Given a raw model name or provider (as a string), returns the first matching model & provider.
@@ -112,51 +108,46 @@ impl DriaWorkflowsConfig {
     /// - If input is a provider, the first matching model in the node config is returned.
     ///
     /// If there are no matching models with this logic, an error is returned.
-    pub fn get_matching_model(&self, model_or_provider: String) -> Result<(ModelProvider, Model)> {
+    pub fn get_matching_model(&self, model_or_provider: String) -> Result<Model> {
         if model_or_provider == "*" {
             // return a random model
             self.models
                 .iter()
                 .choose(&mut rand::thread_rng())
-                .ok_or_else(|| eyre!("No models to randomly pick for '*'."))
+                .ok_or_eyre("could not find models to randomly pick for '*'")
                 .cloned()
         } else if model_or_provider == "!" {
             // return the first model
             self.models
                 .first()
-                .ok_or_else(|| eyre!("No models to choose first for '!'."))
+                .ok_or_eyre("could not find models to choose first for '!'")
                 .cloned()
         } else if let Ok(provider) = ModelProvider::try_from(model_or_provider.clone()) {
             // this is a valid provider, return the first matching model in the config
             self.models
                 .iter()
-                .find(|(p, _)| *p == provider)
-                .ok_or(eyre!(
-                    "Provider {} is not supported by this node.",
-                    provider
+                .find(|&m| m.provider() == provider)
+                .ok_or_eyre(format!(
+                    "Provider {provider} is not supported by this node."
                 ))
                 .cloned()
         } else if let Ok(model) = Model::try_from(model_or_provider.clone()) {
             // this is a valid model, return it if it is supported by the node
             self.models
                 .iter()
-                .find(|(_, m)| *m == model)
-                .ok_or(eyre!("Model {} is not supported by this node.", model))
+                .find(|&m| *m == model)
+                .ok_or_eyre(format!("Model {model} is not supported by this node."))
                 .cloned()
         } else {
             // this is neither a valid provider or model for this node
             Err(eyre!(
-                "Given string '{}' is neither a model nor provider.",
-                model_or_provider
+                "Given string '{model_or_provider}' is neither a model nor provider.",
             ))
         }
     }
 
     /// From a list of model or provider names, return a random matching model & provider.
-    pub fn get_any_matching_model(
-        &self,
-        list_model_or_provider: Vec<String>,
-    ) -> Result<(ModelProvider, Model)> {
+    pub fn get_any_matching_model(&self, list_model_or_provider: Vec<String>) -> Result<Model> {
         // filter models w.r.t supported ones
         let matching_models = list_model_or_provider
             .into_iter()
@@ -182,23 +173,21 @@ impl DriaWorkflowsConfig {
     /// Returns the list of unique providers in the config.
     #[inline]
     pub fn get_providers(&self) -> Vec<ModelProvider> {
-        self.models
-            .iter()
-            .fold(Vec::new(), |mut unique, (provider, _)| {
-                if !unique.contains(provider) {
-                    unique.push(provider.clone());
-                }
-                unique
-            })
+        self.models.iter().fold(Vec::new(), |mut unique, m| {
+            let provider = m.provider();
+
+            if !unique.contains(&provider) {
+                unique.push(provider);
+            }
+
+            unique
+        })
     }
 
-    /// Returns the list of all models in the config.
-    #[inline]
+    /// Returns the names of all models in the config.
+    #[inline(always)]
     pub fn get_model_names(&self) -> Vec<String> {
-        self.models
-            .iter()
-            .map(|(_, model)| model.to_string())
-            .collect()
+        self.models.iter().map(|m| m.to_string()).collect()
     }
 
     /// Check if the required compute services are running.
@@ -226,49 +215,25 @@ impl DriaWorkflowsConfig {
         // if Ollama is a provider, check that it is running & Ollama models are pulled (or pull them)
         if unique_providers.contains(&ModelProvider::Ollama) {
             let provider_models = self.get_models_for_provider(ModelProvider::Ollama);
-            good_models.extend(
-                self.ollama
-                    .check(provider_models)
-                    .await?
-                    .into_iter()
-                    .map(|m| (ModelProvider::Ollama, m)),
-            );
+            good_models.extend(self.ollama.check(provider_models).await?);
         }
 
         // if OpenAI is a provider, check that the API key is set & models are available
         if unique_providers.contains(&ModelProvider::OpenAI) {
             let provider_models = self.get_models_for_provider(ModelProvider::OpenAI);
-            good_models.extend(
-                self.openai
-                    .check(provider_models)
-                    .await?
-                    .into_iter()
-                    .map(|m| (ModelProvider::OpenAI, m)),
-            );
+            good_models.extend(self.openai.check(provider_models).await?);
         }
 
         // if Gemini is a provider, check that the API key is set & models are available
         if unique_providers.contains(&ModelProvider::Gemini) {
             let provider_models = self.get_models_for_provider(ModelProvider::Gemini);
-            good_models.extend(
-                self.gemini
-                    .check(provider_models)
-                    .await?
-                    .into_iter()
-                    .map(|m| (ModelProvider::Gemini, m)),
-            );
+            good_models.extend(self.gemini.check(provider_models).await?);
         }
 
         // if OpenRouter is a provider, check that the API key is set
         if unique_providers.contains(&ModelProvider::OpenRouter) {
             let provider_models = self.get_models_for_provider(ModelProvider::OpenRouter);
-            good_models.extend(
-                self.openrouter
-                    .check(provider_models)
-                    .await?
-                    .into_iter()
-                    .map(|m| (ModelProvider::OpenRouter, m)),
-            );
+            good_models.extend(self.openrouter.check(provider_models).await?);
         }
 
         // update good models
@@ -286,7 +251,7 @@ impl std::fmt::Display for DriaWorkflowsConfig {
         let models_str = self
             .models
             .iter()
-            .map(|(provider, model)| format!("{:?}:{}", provider, model))
+            .map(|model| format!("{}:{}", model.provider(), model))
             .collect::<Vec<_>>()
             .join(",");
         write!(f, "{}", models_str)
@@ -312,15 +277,14 @@ mod tests {
     fn test_model_matching() {
         let cfg = DriaWorkflowsConfig::new_from_csv("gpt-4o,llama3.1:latest");
         assert_eq!(
-            cfg.get_matching_model("openai".to_string()).unwrap().1,
+            cfg.get_matching_model("openai".to_string()).unwrap(),
             Model::GPT4o,
             "Should find existing model"
         );
 
         assert_eq!(
             cfg.get_matching_model("llama3.1:latest".to_string())
-                .unwrap()
-                .1,
+                .unwrap(),
             Model::Llama3_1_8B,
             "Should find existing model"
         );
@@ -347,7 +311,7 @@ mod tests {
             "ollama".to_string(),
         ]);
         assert_eq!(
-            result.unwrap().1,
+            result.unwrap(),
             Model::Llama3_1_8B,
             "Should find existing model"
         );

workflows/src/config.rs

@@ -1,4 +1,4 @@
-use dkn_workflows::{DriaWorkflowsConfig, Model, ModelProvider};
+use dkn_workflows::{DriaWorkflowsConfig, Model};
 use eyre::Result;
 
 #[inline(always)]
@@ -24,10 +24,7 @@ async fn test_ollama_check() -> Result<()> {
     let mut model_config = DriaWorkflowsConfig::new(models);
     model_config.check_services().await?;
 
-    assert_eq!(
-        model_config.models[0],
-        (ModelProvider::Ollama, Model::Phi3_5Mini)
-    );
+    assert_eq!(model_config.models[0], Model::Phi3_5Mini);
 
     Ok(())
 }
@@ -41,10 +38,7 @@ async fn test_openai_check() -> Result<()> {
     let mut model_config = DriaWorkflowsConfig::new(models);
     model_config.check_services().await?;
 
-    assert_eq!(
-        model_config.models[0],
-        (ModelProvider::OpenAI, Model::GPT4Turbo)
-    );
+    assert_eq!(model_config.models[0], Model::GPT4Turbo);
     Ok(())
 }
 
@@ -57,10 +51,7 @@ async fn test_gemini_check() -> Result<()> {
     let mut model_config = DriaWorkflowsConfig::new(models);
     model_config.check_services().await?;
 
-    assert_eq!(
-        model_config.models[0],
-        (ModelProvider::Gemini, Model::Gemini15Flash)
-    );
+    assert_eq!(model_config.models[0], Model::Gemini15Flash);
     Ok(())
 }
 
@@ -73,10 +64,7 @@ async fn test_openrouter_check() -> Result<()> {
     let mut model_config = DriaWorkflowsConfig::new(models);
     model_config.check_services().await?;
 
-    assert_eq!(
-        model_config.models[0],
-        (ModelProvider::OpenRouter, Model::ORDeepSeek2_5)
-    );
+    assert_eq!(model_config.models[0], Model::ORDeepSeek2_5);
     Ok(())
 }