Merge pull request #720 from dcSpark/feature/assets

Save/List/Delete assets for tool
dcSpark · Dec 18, 2024 · b9142df · b9142df
2 parents 0819a87 + 50ccec5
commit b9142df
Show file tree

Hide file tree

Showing 18 changed files with 811 additions and 120 deletions.
diff --git a/shinkai-bin/shinkai-node/src/managers/tool_router.rs b/shinkai-bin/shinkai-node/src/managers/tool_router.rs
@@ -425,6 +425,7 @@ async def run(c: CONFIG, p: INPUTS) -> OUTPUT:
             sql_queries: None,
             file_inbox: None,
             oauth: None,
+            assets: None,
         };
         python_tool
     }
@@ -552,6 +553,7 @@ async def run(c: CONFIG, p: INPUTS) -> OUTPUT:
                         tool_id,
                         node_name,
                         false,
+                        None,
                     )
                     .map_err(|e| LLMProviderError::FunctionExecutionError(e.to_string()))?;
                 let result_str = serde_json::to_string(&result)
@@ -626,9 +628,10 @@ async def run(c: CONFIG, p: INPUTS) -> OUTPUT:
                         function_config_vec,
                         node_storage_path,
                         app_id,
-                        tool_id,
+                        tool_id.clone(),
                         node_name,
                         false,
+                        Some(tool_id),
                     )
                     .map_err(|e| LLMProviderError::FunctionExecutionError(e.to_string()))?;
                 let result_str = serde_json::to_string(&result)
@@ -959,10 +962,11 @@ async def run(c: CONFIG, p: INPUTS) -> OUTPUT:
                 function_config_vec,
                 node_storage_path,
                 app_id,
-                tool_id,
+                tool_id.clone(),
                 // TODO Is this correct?
                 requester_node_name,
                 true,
+                Some(tool_id),
             )
             .map_err(|e| LLMProviderError::FunctionExecutionError(e.to_string()))?;
         let result_str =

diff --git a/shinkai-bin/shinkai-node/src/network/handle_commands_list.rs b/shinkai-bin/shinkai-node/src/network/handle_commands_list.rs
@@ -2591,14 +2591,16 @@ impl Node {
                 res,
             } => {
                 let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
                 tokio::spawn(async move {
-                    let _ = Node::v2_api_export_tool(db_clone, bearer, tool_key_path, res).await;
+                    let _ = Node::v2_api_export_tool(db_clone, bearer, node_env, tool_key_path, res).await;
                 });
             }
             NodeCommand::V2ApiImportTool { bearer, url, res } => {
                 let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
                 tokio::spawn(async move {
-                    let _ = Node::v2_api_import_tool(db_clone, bearer, url, res).await;
+                    let _ = Node::v2_api_import_tool(db_clone, bearer, node_env, url, res).await;
                 });
             }
             NodeCommand::V2ApiRemoveTool {
@@ -2765,10 +2767,18 @@ impl Node {
                     let _ = Node::v2_api_search_shinkai_tool(db_clone, bearer, query, agent_or_llm, res).await;
                 });
             }
-            NodeCommand::V2ApiSetPlaygroundTool { bearer, payload, res } => {
+            NodeCommand::V2ApiSetPlaygroundTool {
+                bearer,
+                payload,
+                tool_id,
+                app_id,
+                res,
+            } => {
                 let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
                 tokio::spawn(async move {
-                    let _ = Node::v2_api_set_playground_tool(db_clone, bearer, payload, res).await;
+                    let _ = Node::v2_api_set_playground_tool(db_clone, bearer, payload, node_env, tool_id, app_id, res)
+                        .await;
                 });
             }
             NodeCommand::V2ApiListPlaygroundTools { bearer, res } => {
@@ -2811,6 +2821,49 @@ impl Node {
                     let _ = Node::v2_api_set_oauth_token(db_clone, bearer, code, state, res).await;
                 });
             }
+            NodeCommand::V2ApiUploadToolAsset {
+                bearer,
+                tool_id,
+                app_id,
+                file_name,
+                file_data,
+                res,
+            } => {
+                let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
+                tokio::spawn(async move {
+                    let _ = Node::v2_api_upload_tool_asset(
+                        db_clone, bearer, tool_id, app_id, file_name, file_data, node_env, res,
+                    )
+                    .await;
+                });
+            }
+            NodeCommand::V2ApiListToolAssets {
+                bearer,
+                tool_id,
+                app_id,
+                res,
+            } => {
+                let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
+                tokio::spawn(async move {
+                    let _ = Node::v2_api_list_tool_assets(db_clone, bearer, tool_id, app_id, node_env, res).await;
+                });
+            }
+            NodeCommand::V2ApiDeleteToolAsset {
+                bearer,
+                tool_id,
+                app_id,
+                file_name,
+                res,
+            } => {
+                let db_clone = Arc::clone(&self.db);
+                let node_env = fetch_node_environment();
+                tokio::spawn(async move {
+                    let _ = Node::v2_api_delete_tool_asset(db_clone, bearer, tool_id, app_id, file_name, node_env, res)
+                        .await;
+                });
+            }
             _ => (),
         }
     }