From 3662b4a1036de0e9fa276a0f31375d0350a81621 Mon Sep 17 00:00:00 2001
From: Guillermo Valin <34345309+guillevalin@users.noreply.github.com>
Date: Mon, 20 Jan 2025 15:57:03 -0300
Subject: [PATCH] fix: tool deletion button (#606)

---
 .../components/remove-tool-button.tsx         |  9 ++++++++-
 .../components/tools/components/tool-card.tsx |  5 ++++-
 .../shinkai-message-ts/src/api/tools/index.ts | 17 +++++++++++++++++
 .../shinkai-message-ts/src/api/tools/types.ts |  4 ++++
 .../src/v2/mutations/removeTool/index.ts      | 19 ++++++++++++++++++-
 .../src/v2/mutations/removeTool/types.ts      |  1 +
 .../v2/mutations/removeTool/useRemoveTool.ts  |  1 +
 7 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/apps/shinkai-desktop/src/components/playground-tool/components/remove-tool-button.tsx b/apps/shinkai-desktop/src/components/playground-tool/components/remove-tool-button.tsx
index 426542e6a..3e858177d 100644
--- a/apps/shinkai-desktop/src/components/playground-tool/components/remove-tool-button.tsx
+++ b/apps/shinkai-desktop/src/components/playground-tool/components/remove-tool-button.tsx
@@ -5,7 +5,13 @@ import { toast } from 'sonner';
 
 import { useAuth } from '../../../store/auth';
 
-export default function RemoveToolButton({ toolKey }: { toolKey: string }) {
+export default function RemoveToolButton({
+  isPlaygroundTool,
+  toolKey,
+}: {
+  isPlaygroundTool: boolean;
+  toolKey: string;
+}) {
   const auth = useAuth((state) => state.auth);
   const navigate = useNavigate();
 
@@ -32,6 +38,7 @@ export default function RemoveToolButton({ toolKey }: { toolKey: string }) {
           toolKey: toolKey ?? '',
           nodeAddress: auth?.node_address ?? '',
           token: auth?.api_v2_key ?? '',
+          isPlaygroundTool,
         });
       }}
       size="sm"
diff --git a/apps/shinkai-desktop/src/components/tools/components/tool-card.tsx b/apps/shinkai-desktop/src/components/tools/components/tool-card.tsx
index c63395b0d..5e09f276b 100644
--- a/apps/shinkai-desktop/src/components/tools/components/tool-card.tsx
+++ b/apps/shinkai-desktop/src/components/tools/components/tool-card.tsx
@@ -245,7 +245,10 @@ export default function ToolCard({
               Go Playground
             </Link>
           )}
-          <RemoveToolButton toolKey={toolKey as string} />
+          <RemoveToolButton
+            isPlaygroundTool={isPlaygroundTool}
+            toolKey={toolKey as string}
+          />
         </div>
       </div>
     </SubpageLayout>
diff --git a/libs/shinkai-message-ts/src/api/tools/index.ts b/libs/shinkai-message-ts/src/api/tools/index.ts
index b16fba6fc..0af77007c 100644
--- a/libs/shinkai-message-ts/src/api/tools/index.ts
+++ b/libs/shinkai-message-ts/src/api/tools/index.ts
@@ -28,6 +28,7 @@ import {
   ImportToolResponse,
   PayInvoiceRequest,
   RemovePlaygroundToolRequest,
+  RemoveToolRequest,
   RemoveToolRequestRequest,
   SaveToolCodeRequest,
   SaveToolCodeResponse,
@@ -333,6 +334,22 @@ export const restoreToolConversation = async (
   return response.data as UndoToolImplementationResponse;
 };
 
+export const removeTool = async (
+  nodeAddress: string,
+  bearerToken: string,
+  payload: RemoveToolRequest,
+) => {
+  const response = await httpClient.delete(
+    urlJoin(nodeAddress, '/v2/remove_tool'),
+    {
+      params: { tool_key: payload.tool_key },
+      headers: { Authorization: `Bearer ${bearerToken}` },
+      responseType: 'json',
+    },
+  );
+  return response.data;
+};
+
 export const removePlaygroundTool = async (
   nodeAddress: string,
   bearerToken: string,
diff --git a/libs/shinkai-message-ts/src/api/tools/types.ts b/libs/shinkai-message-ts/src/api/tools/types.ts
index 4b3305d28..369b9bc6f 100644
--- a/libs/shinkai-message-ts/src/api/tools/types.ts
+++ b/libs/shinkai-message-ts/src/api/tools/types.ts
@@ -305,6 +305,10 @@ export type RemovePlaygroundToolRequest = {
   tool_key: string;
 };
 
+export type RemoveToolRequest = {
+  tool_key: string;
+};
+
 export type ImportToolRequest = {
   url: string;
 };
diff --git a/libs/shinkai-node-state/src/v2/mutations/removeTool/index.ts b/libs/shinkai-node-state/src/v2/mutations/removeTool/index.ts
index 534314ae6..e1d1d41e1 100644
--- a/libs/shinkai-node-state/src/v2/mutations/removeTool/index.ts
+++ b/libs/shinkai-node-state/src/v2/mutations/removeTool/index.ts
@@ -1,4 +1,7 @@
-import { removePlaygroundTool as removePlaygroundToolApi } from '@shinkai_network/shinkai-message-ts/api/tools/index';
+import {
+  removePlaygroundTool as removePlaygroundToolApi,
+  removeTool as removeToolApi,
+} from '@shinkai_network/shinkai-message-ts/api/tools/index';
 
 import { RemoveToolInput } from './types';
 
@@ -6,6 +9,20 @@ export const removeTool = async ({
   nodeAddress,
   token,
   toolKey,
+  isPlaygroundTool,
+}: RemoveToolInput) => {
+  if (isPlaygroundTool) {
+    return await removePlaygroundTool({ nodeAddress, token, toolKey });
+  }
+  return await removeToolApi(nodeAddress, token, {
+    tool_key: toolKey,
+  });
+};
+
+export const removePlaygroundTool = async ({
+  nodeAddress,
+  token,
+  toolKey,
 }: RemoveToolInput) => {
   return await removePlaygroundToolApi(nodeAddress, token, {
     tool_key: toolKey,
diff --git a/libs/shinkai-node-state/src/v2/mutations/removeTool/types.ts b/libs/shinkai-node-state/src/v2/mutations/removeTool/types.ts
index f846e2612..ec3a8304f 100644
--- a/libs/shinkai-node-state/src/v2/mutations/removeTool/types.ts
+++ b/libs/shinkai-node-state/src/v2/mutations/removeTool/types.ts
@@ -7,4 +7,5 @@ export type RemoveToolOutput = {
 export type RemoveToolInput = Token & {
   nodeAddress: string;
   toolKey: string;
+  isPlaygroundTool?: boolean;
 };
diff --git a/libs/shinkai-node-state/src/v2/mutations/removeTool/useRemoveTool.ts b/libs/shinkai-node-state/src/v2/mutations/removeTool/useRemoveTool.ts
index fc51985eb..d4d2ce468 100644
--- a/libs/shinkai-node-state/src/v2/mutations/removeTool/useRemoveTool.ts
+++ b/libs/shinkai-node-state/src/v2/mutations/removeTool/useRemoveTool.ts
@@ -23,6 +23,7 @@ export const useRemoveTool = (options?: Options) => {
           {
             nodeAddress: variables.nodeAddress,
             token: variables.token,
+            isPlaygroundTool: variables.isPlaygroundTool,
           },
         ],
       });