Merge pull request #14 from firstbatchxyz/erhant/workflows-install

Use `dkn-workflows` for model parsing logic & its re-export

Author: erhant

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 name = "dkn-oracle"
 description = "Dria Knowledge Network: Oracle Node"
-version = "0.1.5"
+version = "0.1.6"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"
@@ -20,8 +20,8 @@ tokio = { version = "1.39.2", features = [
 ] }
 lazy_static = "1.5.0"
 
-# compute & workflows
-ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows" }
+# workflows
+dkn-workflows = { git = "https://github.com/firstbatchxyz/dkn-compute-node" }
 
 # errors & logging & env
 env_logger = "0.11.5"

Cargo.toml

@@ -7,15 +7,15 @@ endif
 ###############################################################################
 .PHONY: launch #       | Run with INFO logging & release mode
 launch:
-		RUST_LOG=none,dkn_oracle=info cargo run --release start 
+		RUST_LOG=none,dkn_oracle=info,dkn_workflows=info cargo run --release start 
 
 .PHONY: run #          | Run with INFO logging
 run:
-		RUST_LOG=none,dkn_oracle=info cargo run start
+		RUST_LOG=none,dkn_oracle=info,dkn_workflows=info cargo run start
 
 .PHONY: debug #        | Run with crate-level DEBUG logging & info-level workflows
 debug:
-		RUST_LOG=none,dkn_oracle=debug,ollama_workflows=info cargo run start
+		RUST_LOG=none,dkn_oracle=debug,dkn_workflows=debug,ollama_workflows=info cargo run start
 
 ###############################################################################
 .PHONY: install #        | Install to path

Makefile

@@ -5,8 +5,8 @@ use alloy::{
     primitives::{B256, U256},
 };
 use clap::{Parser, Subcommand};
+use dkn_workflows::Model;
 use eyre::{eyre, Context, Result};
-use ollama_workflows::Model;
 use reqwest::Url;
 
 /// `value_parser` to parse a `str` to `OracleKind`.

src/cli.rs

@@ -1,15 +1,17 @@
+use std::time::Duration;
+
 use crate::{
-    compute::{handle_request, ModelConfig},
+    compute::handle_request,
     contracts::{bytes_to_string, string_to_bytes, OracleKind, TaskStatus},
     DriaOracle,
 };
 use alloy::{
     eips::BlockNumberOrTag,
     primitives::{utils::format_ether, U256},
 };
+use dkn_workflows::{DriaWorkflowsConfig, Model};
 use eyre::{eyre, Context, Result};
 use futures_util::StreamExt;
-use ollama_workflows::Model;
 
 /// Runs the main loop of the oracle node.
 pub async fn run_oracle(
@@ -35,11 +37,17 @@ pub async fn run_oracle(
     }
 
     // prepare model config
-    let model_config = ModelConfig::new(models);
-    if model_config.models_providers.is_empty() {
+    let mut model_config = DriaWorkflowsConfig::new(models);
+    if model_config.models.is_empty() {
         return Err(eyre!("No models provided."))?;
     }
-    model_config.check_providers().await?;
+    let ollama_config = model_config.ollama.clone();
+    model_config = model_config.with_ollama_config(
+        ollama_config
+            .with_min_tps(5.0)
+            .with_timeout(Duration::from_secs(150)),
+    );
+    model_config.check_services().await?;
 
     // check previous tasks
     if from_block.clone().into() != BlockNumberOrTag::Latest {

src/commands/coordinator.rs

@@ -2,14 +2,14 @@ use crate::{
     compute::WorkflowsExt,
     contracts::{bytes32_to_string, bytes_to_string},
     data::Arweave,
-    mine_nonce, DriaOracle, ModelConfig,
+    mine_nonce, DriaOracle,
 };
 use alloy::{
     primitives::{FixedBytes, U256},
     rpc::types::TransactionReceipt,
 };
+use dkn_workflows::{DriaWorkflowsConfig, Executor};
 use eyre::{Context, Result};
-use ollama_workflows::Executor;
 
 /// Handles a generation request.
 ///
@@ -19,7 +19,7 @@ use ollama_workflows::Executor;
 /// 2. Then, we check if our models are compatible with the request. If not, we return an error.
 pub async fn handle_generation(
     node: &DriaOracle,
-    models: &ModelConfig,
+    workflows: &DriaWorkflowsConfig,
     task_id: U256,
     protocol: FixedBytes<32>,
 ) -> Result<Option<TransactionReceipt>> {
@@ -43,7 +43,8 @@ pub async fn handle_generation(
     // choose model based on the request
     log::debug!("Choosing model to use");
     let models_string = bytes_to_string(&request.models)?;
-    let (_, model) = models.get_any_matching_model_from_csv(&models_string)?;
+    let models_vec = models_string.split(',').map(|s| s.to_string()).collect();
+    let (_, model) = workflows.get_any_matching_model(models_vec)?;
     log::debug!("Using model: {} from {}", model, models_string);
 
     // execute task

src/compute/handlers/generation.rs

@@ -1,9 +1,9 @@
-use super::ModelConfig;
 use crate::{
     contracts::{OracleCoordinator::StatusUpdate, OracleKind, TaskStatus},
     DriaOracle,
 };
 use alloy::rpc::types::TransactionReceipt;
+use dkn_workflows::DriaWorkflowsConfig;
 use eyre::Result;
 
 mod generation;
@@ -21,15 +21,15 @@ use validation::*;
 pub async fn handle_request(
     node: &DriaOracle,
     kinds: &[OracleKind],
-    model_config: &ModelConfig,
+    workflows: &DriaWorkflowsConfig,
     event: StatusUpdate,
 ) -> Result<Option<TransactionReceipt>> {
     log::debug!("Received event for task {} ()", event.taskId);
 
     let response_tx_hash = match TaskStatus::try_from(event.statusAfter)? {
         TaskStatus::PendingGeneration => {
             if kinds.contains(&OracleKind::Generator) {
-                handle_generation(node, model_config, event.taskId, event.protocol).await?
+                handle_generation(node, workflows, event.taskId, event.protocol).await?
             } else {
                 log::debug!(
                     "Ignoring generation task {} as you are not generator.",
@@ -40,7 +40,7 @@ pub async fn handle_request(
         }
         TaskStatus::PendingValidation => {
             if kinds.contains(&OracleKind::Validator) {
-                handle_validation(node, model_config, event.taskId).await?
+                handle_validation(node, workflows, event.taskId).await?
             } else {
                 log::debug!(
                     "Ignoring generation task {} as you are not validator.",

src/compute/handlers/mod.rs

@@ -1,16 +1,16 @@
-use super::ModelConfig;
 use crate::{mine_nonce, DriaOracle};
 use alloy::{
     primitives::{utils::parse_ether, Bytes, U256},
     rpc::types::TransactionReceipt,
 };
+use dkn_workflows::DriaWorkflowsConfig;
 use eyre::{eyre, Context, Result};
 
 /// Handles a validation request.
 #[allow(unused)]
 pub async fn handle_validation(
     node: &DriaOracle,
-    models: &ModelConfig,
+    workflows: &DriaWorkflowsConfig,
     task_id: U256,
 ) -> Result<Option<TransactionReceipt>> {
     log::info!("Handling validation task {}", task_id);

src/compute/handlers/validation.rs

@@ -5,4 +5,4 @@ mod nonce;
 pub use nonce::mine_nonce;
 
 mod workflows;
-pub use workflows::{ModelConfig, WorkflowsExt};
+pub use workflows::WorkflowsExt;

src/compute/mod.rs

@@ -1,7 +1,7 @@
 use alloy::primitives::Bytes;
 use async_trait::async_trait;
+use dkn_workflows::{Entry, Executor, ProgramMemory, Workflow};
 use eyre::{Context, Result};
-use ollama_workflows::{Entry, Executor, ProgramMemory, Workflow};
 
 use super::postprocess::*;
 use crate::data::{Arweave, OracleExternalData};

src/compute/workflows/executor.rs

@@ -1,19 +1,13 @@
 mod executor;
 pub use executor::WorkflowsExt;
 
-mod models;
-pub use models::ModelConfig;
-
-mod providers;
-pub use providers::*;
-
 mod postprocess;
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use alloy::primitives::Bytes;
-    use ollama_workflows::{Entry, Executor, Model, ProgramMemory};
+    use dkn_workflows::{Entry, Executor, Model, ProgramMemory};
 
     #[tokio::test]
     #[ignore = "run this manually"]