- feature: add envs support to tool.run (#60)

* - feature: add envs support to tool.run * - fix: add additional test
dcSpark · Nov 12, 2024 · 148b1f0 · 148b1f0
1 parent 8154708
commit 148b1f0
Show file tree

Hide file tree

Showing 8 changed files with 73 additions and 36 deletions.
diff --git a/.github/workflows/release-production.yml b/.github/workflows/release-production.yml
@@ -75,7 +75,7 @@ jobs:
           cargo set-version ${{ needs.prebuild.outputs.version }}
 
       - name: Run NX publish
-        run: npx nx run-many -t publish --verbose
+        run: npx nx run-many -t publish --verbose --parallel=false
         env:
           CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
           EMBEDDING_API_URL: ${{ vars.EMBEDDING_API_URL }}
diff --git a/libs/shinkai-tools-runner/src/lib.test.rs b/libs/shinkai-tools-runner/src/lib.test.rs
@@ -18,7 +18,7 @@ async fn shinkai_tool_echo() {
         None,
     );
     let run_result = tool
-        .run(serde_json::json!({ "message": "valparaíso" }), None)
+        .run(None, serde_json::json!({ "message": "valparaíso" }), None)
         .await
         .unwrap();
     assert_eq!(run_result.data["message"], "echoing: valparaíso");
@@ -37,7 +37,7 @@ async fn shinkai_tool_weather_by_city() {
         None,
     );
     let run_result = tool
-        .run(serde_json::json!({ "city": "valparaíso" }), None)
+        .run(None, serde_json::json!({ "city": "valparaíso" }), None)
         .await;
     assert!(run_result.is_ok());
 }
@@ -55,7 +55,7 @@ async fn shinkai_tool_inline() {
 "#;
     let tool = Tool::new(js_code.to_string(), serde_json::Value::Null, None);
     let run_result = tool
-        .run(serde_json::json!({ "name": "world" }), None)
+        .run(None, serde_json::json!({ "name": "world" }), None)
         .await
         .unwrap();
     assert_eq!(run_result.data["message"], "Hello, world!");
@@ -73,7 +73,7 @@ async fn shinkai_tool_inline_non_json_return() {
         }
 "#;
     let tool = Tool::new(js_code.to_string(), serde_json::Value::Null, None);
-    let run_result = tool.run(serde_json::json!({}), None).await.unwrap();
+    let run_result = tool.run(None, serde_json::json!({}), None).await.unwrap();
     assert_eq!(run_result.data, 5);
 }
 
@@ -91,6 +91,7 @@ async fn shinkai_tool_web3_eth_balance() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!({ "address": "0x388c818ca8b9251b393131c08a736a67ccb19297" }),
             None,
         )
@@ -113,6 +114,7 @@ async fn shinkai_tool_web3_eth_uniswap() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!({
                 "fromToken": "ETH",
                 "toToken": "USDC",
@@ -141,6 +143,7 @@ async fn shinkai_tool_download_page() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!({
                 "urls": "https://shinkai.com"
             }),
@@ -175,7 +178,7 @@ async fn max_execution_time() {
 "#;
     let tool = Tool::new(js_code.to_string(), serde_json::Value::Null, None);
     let run_result = tool
-        .run(serde_json::json!({ "timeoutMs": 3200 }), Some(3000))
+        .run(None, serde_json::json!({ "timeoutMs": 3200 }), Some(3000))
         .await;
     assert!(run_result.is_err());
     assert!(run_result.err().unwrap().message().contains("timed out"));
@@ -200,6 +203,7 @@ async fn shinkai_tool_download_page_stack_overflow() {
                     None,
                 );
                 tool.run(
+                    None,
                     serde_json::json!({
                         "url": "https://en.wikipedia.org/wiki/Prospect_Park_(Brooklyn)"
                     }),
@@ -311,7 +315,7 @@ async fn shinkai_tool_leiden() {
       "edges": edges
     });
     let start_time = std::time::Instant::now(); // Start measuring time
-    let run_result = tool.run(params, None).await.unwrap();
+    let run_result = tool.run(None, params, None).await.unwrap();
     let elapsed_time = start_time.elapsed(); // Measure elapsed time
 
     println!("Execution time: {:?}", elapsed_time); // Print the elapsed time
@@ -337,6 +341,7 @@ async fn shinkai_tool_duckduckgo_search() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!({ "message": "best movie of all time" }),
             None,
         )
@@ -367,6 +372,7 @@ async fn shinkai_tool_playwright_example() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!({
                 "url": "https://shinkai.com"
             }),
@@ -397,6 +403,7 @@ async fn shinkai_tool_defillama_lending_tvl_rankings() {
     );
     let run_result = tool
         .run(
+            None,
             serde_json::json!(      {
               "top10": false,
               "categoryName": "Liquid Staking",
@@ -450,6 +457,7 @@ async fn shinkai_tool_youtube_summary() {
     let tool = Tool::new(tool_definition.code.clone().unwrap(), configurations, None);
     let run_result = tool
         .run(
+            None,
             serde_json::json!({ "url": "https://www.youtube.com/watch?v=GQ9yRPfsDPk" }),
             None,
         )
@@ -536,6 +544,7 @@ async fn shinkai_tool_json_to_md() {
     .to_string();
     let run_result = tool
         .run(
+            None,
             json!({
                 "message": message,
             "template": "# Introduction{%- for sentence in answer.brief_introduction.sentences %}{{ sentence }}{%- endfor %}\\# Body{%- for section in answer.extensive_body %}## Section {{ loop.index }}{%- for sentence in section.sentences %}{{ sentence }}{%- endfor %}{%- endfor %}\\# Conclusion{%- for section in answer.conclusion %}{{ section.sentences[0] }}{%- endfor %}\\# Citations{%- for citation in relevantSentencesFromText %}[{{ citation.citation_id }}]: {{ citation.relevantSentenceFromDocument }}{%- endfor %}"}),

diff --git a/libs/shinkai-tools-runner/src/tools/deno_runner.rs b/libs/shinkai-tools-runner/src/tools/deno_runner.rs
@@ -21,7 +21,7 @@ impl DenoRunner {
     pub async fn run(
         &mut self,
         code: &str,
-        envs: HashMap<String, String>,
+        envs: Option<HashMap<String, String>>,
         max_execution_time_s: Option<u64>,
     ) -> Result<String, std::io::Error> {
         log::info!("using deno binary at path: {:?}", self.options.binary_path);
@@ -41,8 +41,10 @@ impl DenoRunner {
             .arg(temp_file.path().to_str().unwrap())
             .stdout(std::process::Stdio::piped())
             .stderr(std::process::Stdio::piped())
-            .envs(envs)
             .kill_on_drop(true);
+        if let Some(envs) = envs {
+            command.envs(envs);
+        }
         log::info!("prepared command with arguments: {:?}", command);
         let child = command.spawn()?;
 

diff --git a/libs/shinkai-tools-runner/src/tools/deno_runner.test.rs b/libs/shinkai-tools-runner/src/tools/deno_runner.test.rs
@@ -9,23 +9,21 @@ async fn test_run_echo_tool() {
       console.log('{"message":"hello world"}');
     "#;
 
-    let result = deno_runner
-        .run(
-            code,
-            HashMap::from([
-                (
-                    "configurations".to_string(),
-                    serde_json::json!({}).to_string(),
-                ),
-                ("parameters".to_string(), serde_json::json!({}).to_string()),
-            ]),
-            None,
-        )
-        .await
-        .unwrap();
+    let result = deno_runner.run(code, None, None).await.unwrap();
 
-    assert_eq!(
-        result,
-        "{\"message\":\"hello world\"}\n"
-    );
+    assert_eq!(result, "{\"message\":\"hello world\"}\n");
+}
+
+#[tokio::test]
+async fn test_run_with_env() {
+    let mut deno_runner = DenoRunner::default();
+    let code = r#"
+      console.log(process.env.HELLO_WORLD);
+    "#;
+
+    let mut envs = HashMap::<String, String>::new();
+    envs.insert("HELLO_WORLD".to_string(), "hello world!".to_string()); // Insert the key-value pair
+    let result = deno_runner.run(code, Some(envs), None).await.unwrap();
+
+    assert_eq!(result, "hello world!\n");
 }
diff --git a/libs/shinkai-tools-runner/src/tools/tool.rs b/libs/shinkai-tools-runner/src/tools/tool.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use serde_json::Value;
 
 use super::{
@@ -33,7 +35,6 @@ impl Tool {
         let mut deno_runner = DenoRunner::new(self.deno_runner_options.clone());
 
         // Empty envs when get definition
-        let envs = std::collections::HashMap::new();
         let code = format!(
             r#"
             {}
@@ -44,7 +45,7 @@ impl Tool {
             &self.code.to_string()
         );
         let result = deno_runner
-            .run(&code, envs, None)
+            .run(&code, None, None)
             .await
             .map_err(|e| ExecutionError::new(format!("failed to run deno: {}", e), None))?;
 
@@ -72,6 +73,7 @@ impl Tool {
 
     pub async fn run(
         &self,
+        envs: Option<HashMap<String, String>>,
         parameters: Value,
         max_execution_time_s: Option<u64>,
     ) -> Result<RunResult, ExecutionError> {
@@ -80,8 +82,6 @@ impl Tool {
         log::info!("parameters: {}", parameters.to_string());
 
         let mut deno_runner = DenoRunner::new(self.deno_runner_options.clone());
-        // Empty envs when get definition
-        let envs = std::collections::HashMap::new();
         let code = format!(
             r#"
             {}
@@ -94,8 +94,12 @@ impl Tool {
             console.log("</shinkai-tool-result>");
         "#,
             &self.code.to_string(),
-            serde_json::to_string(&self.configurations).unwrap().replace("\\", "\\\\"),
-            serde_json::to_string(&parameters).unwrap().replace("\\", "\\\\"),
+            serde_json::to_string(&self.configurations)
+                .unwrap()
+                .replace("\\", "\\\\"),
+            serde_json::to_string(&parameters)
+                .unwrap()
+                .replace("\\", "\\\\"),
         );
         let result = deno_runner
             .run(&code, envs, max_execution_time_s)

diff --git a/libs/shinkai-tools-runner/src/tools/tool.test.rs b/libs/shinkai-tools-runner/src/tools/tool.test.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use serde_json::Value;
 
 use crate::tools::tool::Tool;
@@ -33,6 +35,7 @@ async fn run_tool() {
 
     let result = tool
         .run(
+            None,
             serde_json::json!({
                 "message": "hello world"
             }),
@@ -46,3 +49,24 @@ async fn run_tool() {
         serde_json::json!({ "message": "echoing: hello world"})
     );
 }
+
+#[tokio::test]
+async fn shinkai_tool_with_env() {
+    let _ = env_logger::builder()
+        .filter_level(log::LevelFilter::Info)
+        .is_test(true)
+        .try_init();
+    let js_code = r#"
+        function run(configurations, params) {
+            return { foo: process.env.BAR };
+        }
+"#;
+    let tool = Tool::new(js_code.to_string(), serde_json::Value::Null, None);
+    let mut envs = HashMap::<String, String>::new();
+    envs.insert("BAR".to_string(), "bar".to_string());
+    let run_result = tool
+        .run(Some(envs), serde_json::json!({ "name": "world" }), None)
+        .await
+        .unwrap();
+    assert_eq!(run_result.data["foo"], "bar");
+}
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@shinkai_protocol/source",
-  "version": "0.8.1",
+  "version": "0.8.2",
   "description": "This repository serves as the ecosystem to execute Shinkai tools, provided by the Shinkai team or third-party developers, in a secure environment. It provides a sandboxed space for executing these tools, ensuring that they run safely and efficiently, while also allowing for seamless integration with Rust code.",
   "main": "index.js",
   "author": "",