Merge pull request #21 from firstbatchxyz/erhant/parser-logic-update

feat: historic messages & refactors

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "dkn-oracle"
 description = "Dria Knowledge Network: Oracle Node"
-version = "0.1.7"
+version = "0.1.8"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"
@@ -36,7 +36,8 @@ bytes = "1.7.1"
 rand = "0.8.5"
 reqwest = "0.12.5"
 
-# hex, serde
+# b64, hex, serde
+base64 = "0.22.1"
 hex = "0.4.3"
 hex-literal = "0.4.1"
 serde = "1.0.204"
@@ -49,4 +50,3 @@ clap = { version = "4.5.13", features = ["derive", "env"] }
 # there are many unused stuff here, but everything breaks if you use the minimal set
 # because Bundlr SDK is not maintained at all
 bundlr-sdk = { version = "0.5.0" }
-base64 = "0.22.1"

src/commands/coordinator.rs

@@ -92,20 +92,21 @@ pub async fn run_oracle(
     }
 
     // handle new tasks with subscription
-    let event_poller = node
-        .subscribe_to_tasks()
-        .await
-        .wrap_err("could not subscribe to tasks")?;
     log::info!(
-        "Subscribed to LLMOracleCoordinator ({}) as {}",
+        "Subscribing to LLMOracleCoordinator ({}) as {}",
         node.addresses.coordinator,
         kinds
             .iter()
             .map(|kind| kind.to_string())
             .collect::<Vec<String>>()
             .join(", ")
     );
+    let event_poller = node
+        .subscribe_to_tasks()
+        .await
+        .wrap_err("could not subscribe to tasks")?;
 
+    log::info!("Listening for events...");
     event_poller
         .into_stream()
         .for_each(|log| async {

src/commands/token.rs

@@ -43,6 +43,7 @@ pub async fn claim_rewards(node: &DriaOracle) -> Result<()> {
         // transfer rewards
         node.transfer_from(node.addresses.coordinator, node.address(), allowance.amount)
             .await?;
+        log::info!("Rewards claimed: {}.", allowance);
     }
 
     Ok(())

src/compute/handlers/generation.rs

@@ -1,5 +1,5 @@
 use crate::{
-    compute::WorkflowsExt,
+    compute::Request,
     contracts::{bytes32_to_string, bytes_to_string},
     data::Arweave,
     mine_nonce, DriaOracle,
@@ -8,7 +8,7 @@ use alloy::{
     primitives::{FixedBytes, U256},
     rpc::types::TransactionReceipt,
 };
-use dkn_workflows::{DriaWorkflowsConfig, Executor};
+use dkn_workflows::DriaWorkflowsConfig;
 use eyre::{Context, Result};
 
 /// Handles a generation request.
@@ -50,14 +50,17 @@ pub async fn handle_generation(
     // execute task
     log::debug!("Executing the workflow");
     let protocol_string = bytes32_to_string(&protocol)?;
-    let executor = Executor::new(model);
-    let (output, metadata, use_storage) = executor
-        .execute_raw(&request.input, &protocol_string)
-        .await?;
+    let mut input = Request::try_parse_bytes(&request.input).await?;
+    let output = input.execute(model, Some(node)).await?;
+    log::debug!("Output: {}", output);
+
+    // post-processing
+    log::debug!("Post-processing the output");
+    let (output, metadata, use_storage) = Request::post_process(output, &protocol_string).await?;
 
-    // do the Arweave trick for large inputs
-    log::debug!("Uploading to Arweave if required");
-    let arweave = Arweave::new_from_env().wrap_err("could not create Arweave instance")?;
+    // uploading to storage
+    log::debug!("Uploading output to storage");
+    let arweave = Arweave::new_from_env()?;
     let output = if use_storage {
         arweave.put_if_large(output).await?
     } else {

src/compute/mod.rs

@@ -5,4 +5,4 @@ mod nonce;
 pub use nonce::mine_nonce;
 
 mod workflows;
-pub use workflows::WorkflowsExt;
+pub use workflows::*;

src/compute/workflows/mod.rs

@@ -1,56 +1,6 @@
-mod executor;
-pub use executor::WorkflowsExt;
-
 mod postprocess;
 
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use alloy::primitives::Bytes;
-    use dkn_workflows::{Entry, Executor, Model, ProgramMemory};
-
-    #[tokio::test]
-    #[ignore = "run this manually"]
-    async fn test_ollama_generation() {
-        dotenvy::dotenv().unwrap();
-        let executor = Executor::new(Model::Llama3_1_8B);
-        let (output, _, _) = executor
-            .execute_raw(&Bytes::from_static(b"What is the result of 2 + 2?"), "")
-            .await
-            .unwrap();
-
-        // funny test but it should pass
-        println!("Output:\n{}", output);
-        // assert!(output.contains('4')); // FIXME: make this use bytes
-    }
-
-    #[tokio::test]
-    #[ignore = "run this manually"]
-    async fn test_openai_generation() {
-        dotenvy::dotenv().unwrap();
-        let executor = Executor::new(Model::Llama3_1_8B);
-        let (output, _, _) = executor
-            .execute_raw(&Bytes::from_static(b"What is the result of 2 + 2?"), "")
-            .await
-            .unwrap();
-
-        // funny test but it should pass
-        println!("Output:\n{}", output);
-        // assert!(output.contains('4')); // FIXME: make this use bytes
-    }
+mod presets;
 
-    /// Test the generation workflow with a plain input.
-    #[tokio::test]
-    async fn test_workflow_plain() {
-        dotenvy::dotenv().unwrap();
-        let executor = Executor::new(Model::GPT4o);
-        let mut memory = ProgramMemory::new();
-        let workflow = executor.get_generation_workflow().unwrap();
-        let input = Entry::try_value_or_str("What is 2 + 2?");
-        let output = executor
-            .execute(Some(&input), workflow, &mut memory)
-            .await
-            .unwrap();
-        println!("Output:\n{}", output);
-    }
-}
+mod requests;
+pub use requests::Request;

src/compute/workflows/presets/chat.json

@@ -0,0 +1,52 @@
+{
+  "name": "LLM generation",
+  "description": "Directly generate text with input",
+  "config": {
+    "max_steps": 10,
+    "max_time": 50,
+    "tools": [""]
+  },
+  "external_memory": {
+    "context": [""],
+    "question": [""],
+    "answer": [""]
+  },
+  "tasks": [
+    {
+      "id": "A",
+      "name": "Generate with history",
+      "description": "Expects an array of messages for generation",
+      "messages": [],
+      "inputs": [],
+      "operator": "generation",
+      "outputs": [
+        {
+          "type": "write",
+          "key": "result",
+          "value": "__result"
+        }
+      ]
+    },
+    {
+      "id": "__end",
+      "name": "end",
+      "description": "End of the task",
+      "messages": [{ "role": "user", "content": "End of the task" }],
+      "inputs": [],
+      "operator": "end",
+      "outputs": []
+    }
+  ],
+  "steps": [
+    {
+      "source": "A",
+      "target": "__end"
+    }
+  ],
+  "return_value": {
+    "input": {
+      "type": "read",
+      "key": "result"
+    }
+  }
+}

src/compute/workflows/presets/generation.json

@@ -15,8 +15,13 @@
     {
       "id": "A",
       "name": "Generate",
-      "description": "",
-      "prompt": "{text}",
+      "description": "Executes a simple generation request",
+      "messages": [
+        {
+          "role": "user",
+          "content": "{{text}}"
+        }
+      ],
       "inputs": [
         {
           "name": "text",
@@ -40,7 +45,7 @@
       "id": "__end",
       "name": "end",
       "description": "End of the task",
-      "prompt": "End of the task",
+      "messages": [{ "role": "user", "content": "End of the task" }],
       "inputs": [],
       "operator": "end",
       "outputs": []