use workflows subcrate in compute

Author: erhant

Cargo.lock

@@ -16,6 +16,7 @@ parking_lot = "0.12.2"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 async-trait = "0.1.81"
+dotenvy = "0.15.7"
 reqwest = "0.12.5"
 rand = "0.8.5"
 env_logger = "0.11.3"

Cargo.toml

@@ -7,15 +7,15 @@ endif
 ###############################################################################
 .PHONY: launch #       | Run with INFO logs in release mode
 launch:
-		RUST_LOG=none,dkn_compute=info cargo run --release
+		RUST_LOG=none,dkn_compute=info,dkn_workflows=info cargo run --release
 
 .PHONY: run #          | Run with INFO logs
 run:
-		RUST_LOG=none,dkn_compute=info cargo run
+		RUST_LOG=none,dkn_compute=info,dkn_workflows=info cargo run
 
 .PHONY: debug #        | Run with DEBUG logs with INFO log-level workflows
 debug:
-		RUST_LOG=warn,dkn_compute=debug,ollama_workflows=info cargo run
+		RUST_LOG=warn,dkn_compute=debug,,dkn_workflows=debug,ollama_workflows=info cargo run
 
 .PHONY: trace #        | Run with TRACE logs
 trace:

Makefile

@@ -52,11 +52,9 @@ sha2 = "0.10.8"
 sha3 = "0.10.8"
 fastbloom-rs = "0.5.9"
 
-# workflows
-ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows" }
-
-# p2p
+# dria subcrates
 dkn-p2p = { path = "../p2p" }
+dkn-workflows = { path = "../workflows" }
 
 # Vendor OpenSSL so that its easier to build cross-platform packages
 [dependencies.openssl]

compute/Cargo.toml

@@ -1,23 +1,10 @@
-mod models;
-mod ollama;
-mod openai;
-
 use crate::utils::{address_in_use, crypto::to_address};
 use dkn_p2p::libp2p::Multiaddr;
+use dkn_workflows::ModelConfig;
 use eyre::{eyre, Result};
 use libsecp256k1::{PublicKey, SecretKey};
-use models::ModelConfig;
-use ollama::OllamaConfig;
-use ollama_workflows::ModelProvider;
-use openai::OpenAIConfig;
-
-use std::{env, str::FromStr, time::Duration};
-
-/// Timeout duration for checking model performance during a generation.
-const CHECK_TIMEOUT_DURATION: Duration = Duration::from_secs(80);
 
-/// Minimum tokens per second (TPS) for checking model performance during a generation.
-const CHECK_TPS: f64 = 15.0;
+use std::{env, str::FromStr};
 
 #[derive(Debug, Clone)]
 pub struct DriaComputeNodeConfig {
@@ -33,11 +20,6 @@ pub struct DriaComputeNodeConfig {
     pub p2p_listen_addr: Multiaddr,
     /// Available LLM models & providers for the node.
     pub model_config: ModelConfig,
-    /// Even if Ollama is not used, we store the host & port here.
-    /// If Ollama is used, this config will be respected during its instantiations.
-    pub ollama_config: OllamaConfig,
-    /// OpenAI API key & its service check implementation.
-    pub openai_config: OpenAIConfig,
 }
 
 /// The default P2P network listen address.
@@ -97,7 +79,7 @@ impl DriaComputeNodeConfig {
         let address = to_address(&public_key);
         log::info!("Node Address:     0x{}", hex::encode(address));
 
-        let model_config = ModelConfig::new_from_csv(env::var("DKN_MODELS").ok());
+        let model_config = ModelConfig::new_from_csv(&env::var("DKN_MODELS").unwrap_or_default());
         #[cfg(not(test))]
         if model_config.models.is_empty() {
             log::error!("No models were provided, make sure to restart with at least one model provided within DKN_MODELS.");
@@ -118,72 +100,11 @@ impl DriaComputeNodeConfig {
             address,
             model_config,
             p2p_listen_addr,
-            ollama_config: OllamaConfig::new(),
-            openai_config: OpenAIConfig::new(),
-        }
-    }
-
-    /// Check if the required compute services are running.
-    /// This has several steps:
-    ///
-    /// - If Ollama models are used, hardcoded models are checked locally, and for
-    ///   external models, the workflow is tested with a simple task with timeout.
-    /// - If OpenAI models are used, the API key is checked and the models are tested
-    ///
-    /// If both type of models are used, both services are checked.
-    /// In the end, bad models are filtered out and we simply check if we are left if any valid models at all.
-    /// If not, an error is returned.
-    pub async fn check_services(&mut self) -> Result<()> {
-        log::info!("Checking configured services.");
-
-        // TODO: can refactor (provider, model) logic here
-        let unique_providers = self.model_config.get_providers();
-
-        let mut good_models = Vec::new();
-
-        // if Ollama is a provider, check that it is running & Ollama models are pulled (or pull them)
-        if unique_providers.contains(&ModelProvider::Ollama) {
-            let ollama_models = self
-                .model_config
-                .get_models_for_provider(ModelProvider::Ollama);
-
-            // ensure that the models are pulled / pull them if not
-            let good_ollama_models = self
-                .ollama_config
-                .check(ollama_models, CHECK_TIMEOUT_DURATION, CHECK_TPS)
-                .await?;
-            good_models.extend(
-                good_ollama_models
-                    .into_iter()
-                    .map(|m| (ModelProvider::Ollama, m)),
-            );
-        }
-
-        // if OpenAI is a provider, check that the API key is set
-        if unique_providers.contains(&ModelProvider::OpenAI) {
-            let openai_models = self
-                .model_config
-                .get_models_for_provider(ModelProvider::OpenAI);
-
-            let good_openai_models = self.openai_config.check(openai_models).await?;
-            good_models.extend(
-                good_openai_models
-                    .into_iter()
-                    .map(|m| (ModelProvider::OpenAI, m)),
-            );
-        }
-
-        // update good models
-        if good_models.is_empty() {
-            Err(eyre!("No good models found, please check logs for errors."))
-        } else {
-            self.model_config.models = good_models;
-            Ok(())
         }
     }
 
-    // ensure that listen address is free
-    pub fn check_address_in_use(&self) -> Result<()> {
+    /// Asserts that the configured listen address is free.
+    pub fn assert_address_not_in_use(&self) -> Result<()> {
         if address_in_use(&self.p2p_listen_addr) {
             return Err(eyre!(
                 "Listen address {} is already in use.",