feat: download tool assets

dcSpark · Dec 18, 2024 · aba46af · aba46af
1 parent a112c18
commit aba46af
Show file tree

Hide file tree

Showing 8 changed files with 192 additions and 10 deletions.
diff --git a/apps/shinkai-desktop/src/components/playground-tool/components/tool-playground.tsx b/apps/shinkai-desktop/src/components/playground-tool/components/tool-playground.tsx
@@ -6,6 +6,7 @@ import {
   CodeLanguage,
   ToolMetadata,
 } from '@shinkai_network/shinkai-message-ts/api/tools/types';
+import { useGetShinkaiFileProtocol } from '@shinkai_network/shinkai-node-state/v2/queries/getShinkaiFileProtocol/useGetShinkaiFileProtocol';
 import {
   Badge,
   Button,
@@ -29,20 +30,29 @@ import {
   TooltipPortal,
   TooltipTrigger,
 } from '@shinkai_network/shinkai-ui';
-import { SendIcon } from '@shinkai_network/shinkai-ui/assets';
+import {
+  fileIconMap,
+  FileTypeIcon,
+  SendIcon,
+} from '@shinkai_network/shinkai-ui/assets';
 import { cn } from '@shinkai_network/shinkai-ui/utils';
+import { save } from '@tauri-apps/plugin-dialog';
+import * as fs from '@tauri-apps/plugin-fs';
+import { BaseDirectory } from '@tauri-apps/plugin-fs';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
   ArrowUpRight,
   Loader2,
   LucideArrowLeft,
+  Paperclip,
   Play,
   Redo2Icon,
   Save,
   Undo2Icon,
 } from 'lucide-react';
 import { useEffect, useRef, useState } from 'react';
 import { Link, To } from 'react-router-dom';
+import { toast } from 'sonner';
 
 import { useAuth } from '../../../store/auth';
 import { AIModelSelector } from '../../chat/chat-action-bar/ai-update-selection-action-bar';
@@ -112,6 +122,7 @@ function PlaygroundToolEditor({
     metadataEditorRef,
     executeToolCodeQuery,
     toolResult,
+    toolResultFiles,
     isDirtyCodeEditor,
     setIsDirtyCodeEditor,
     forceGenerateMetadata,
@@ -699,15 +710,33 @@ function PlaygroundToolEditor({
                                     <div ref={toolResultBoxRef}>
                                       {executeToolCodeQuery.isSuccess &&
                                         toolResult && (
-                                          <ToolCodeEditor
-                                            language="json"
-                                            readOnly
-                                            value={JSON.stringify(
-                                              toolResult,
-                                              null,
-                                              2,
+                                          <div className="space-y-3 py-3">
+                                            {toolResultFiles.length > 0 && (
+                                              <div className="inline-flex items-center gap-4">
+                                                <h1>Generated Files</h1>
+                                                <div className="flex flex-wrap gap-2">
+                                                  {toolResultFiles?.map(
+                                                    (file) => (
+                                                      <ToolResultFileCard
+                                                        filePath={file}
+                                                        key={file}
+                                                      />
+                                                    ),
+                                                  )}
+                                                </div>
+                                              </div>
                                             )}
-                                          />
+
+                                            <ToolCodeEditor
+                                              language="json"
+                                              readOnly
+                                              value={JSON.stringify(
+                                                toolResult,
+                                                null,
+                                                2,
+                                              )}
+                                            />
+                                          </div>
                                         )}
                                     </div>
                                   </motion.div>
@@ -809,3 +838,70 @@ function PlaygroundToolEditor({
 }
 
 export default PlaygroundToolEditor;
+
+function ToolResultFileCard({ filePath }: { filePath: string }) {
+  const auth = useAuth((state) => state.auth);
+  const { refetch } = useGetShinkaiFileProtocol(
+    {
+      nodeAddress: auth?.node_address ?? '',
+      token: auth?.api_v2_key ?? '',
+      file: filePath,
+    },
+    {
+      enabled: false,
+    },
+  );
+
+  const fileNameBase = filePath.split('/')?.at(-1) ?? 'untitled_tool';
+  const fileExtension = fileNameBase.split('.')?.at(-1) ?? '';
+
+  return (
+    <Button
+      className="h-[30px] gap-1.5 rounded-lg text-xs"
+      onClick={async () => {
+        const response = await refetch();
+        const file = new Blob([response.data ?? ''], {
+          type: 'application/octet-stream',
+        });
+
+        const arrayBuffer = await file.arrayBuffer();
+        const content = new Uint8Array(arrayBuffer);
+
+        const savePath = await save({
+          defaultPath: `${fileNameBase}.${fileExtension}`,
+          filters: [
+            {
+              name: 'File',
+              extensions: [fileExtension],
+            },
+          ],
+        });
+
+        if (!savePath) {
+          toast.info('File saving cancelled');
+          return;
+        }
+
+        await fs.writeFile(savePath, content, {
+          baseDir: BaseDirectory.Download,
+        });
+
+        toast.success(`${fileNameBase} downloaded successfully`);
+      }}
+      size="auto"
+      variant="outline"
+    >
+      <div className="flex shrink-0 items-center justify-center">
+        {fileExtension && fileIconMap[fileExtension] ? (
+          <FileTypeIcon
+            className="text-gray-80 h-[18px] w-[18px] shrink-0"
+            type={fileExtension}
+          />
+        ) : (
+          <Paperclip className="text-gray-80 h-3.5 w-3.5 shrink-0" />
+        )}
+      </div>
+      <div className="text-left text-xs">{filePath.split('/')?.at(-1)}</div>
+    </Button>
+  );
+}
diff --git a/apps/shinkai-desktop/src/components/playground-tool/hooks/use-tool-code.ts b/apps/shinkai-desktop/src/components/playground-tool/hooks/use-tool-code.ts
@@ -410,6 +410,14 @@ export const useToolCode = ({
     }, 0);
   };
 
+  const { __created_files__: toolResultFiles, ...toolResultWithoutFiles } =
+    (toolResult || {
+      __created_files__: [],
+    }) as {
+      __created_files__: string[];
+      [key: string]: any;
+    };
+
   return {
     chatInboxId,
     toolCode,
@@ -432,7 +440,8 @@ export const useToolCode = ({
     metadataEditorRef,
     createToolCode,
     executeToolCodeQuery,
-    toolResult,
+    toolResult: toolResultWithoutFiles,
+    toolResultFiles: toolResultFiles,
     forceGenerateMetadata,
     isDirtyCodeEditor,
     setIsDirtyCodeEditor,

diff --git a/libs/shinkai-message-ts/src/api/tools/index.ts b/libs/shinkai-message-ts/src/api/tools/index.ts
@@ -17,6 +17,8 @@ import {
   GetPlaygroundToolRequest,
   GetPlaygroundToolResponse,
   GetPlaygroundToolsResponse,
+  GetShinkaiFileProtocolRequest,
+  GetShinkaiFileProtocolResponse,
   GetToolResponse,
   GetToolsResponse,
   ImportToolRequest,
@@ -380,3 +382,19 @@ export const exportTool = async (
   );
   return response.data as ExportToolResponse;
 };
+
+export const getShinkaiFileProtocol = async (
+  nodeAddress: string,
+  bearerToken: string,
+  payload: GetShinkaiFileProtocolRequest,
+) => {
+  const response = await httpClient.get(
+    urlJoin(nodeAddress, '/v2/resolve_shinkai_file_protocol'),
+    {
+      headers: { Authorization: `Bearer ${bearerToken}` },
+      params: { file: payload.file },
+      responseType: 'blob',
+    },
+  );
+  return response.data as GetShinkaiFileProtocolResponse;
+};
diff --git a/libs/shinkai-message-ts/src/api/tools/types.ts b/libs/shinkai-message-ts/src/api/tools/types.ts
@@ -275,6 +275,7 @@ export type PlaygroundTool = {
   job_id: string;
   job_id_history: string[];
   code: string;
+  language: string;
 };
 export type GetPlaygroundToolsResponse = PlaygroundTool[];
 export type GetPlaygroundToolRequest = { tool_key: string };
@@ -316,3 +317,5 @@ export type ExportToolRequest = {
 };
 
 export type ExportToolResponse = Blob;
+export type GetShinkaiFileProtocolRequest = { file: string };
+export type GetShinkaiFileProtocolResponse = Blob;
diff --git a/libs/shinkai-node-state/src/v2/constants.ts b/libs/shinkai-node-state/src/v2/constants.ts
@@ -38,6 +38,7 @@ export enum FunctionKeyV2 {
   GET_RECURRING_TASKS = 'GET_RECURRING_TASKS',
   GET_RECURRING_TASK = 'GET_RECURRING_TASK',
   GET_RECURRING_TASK_LOGS = 'GET_RECURRING_TASK_LOGS',
+  GET_SHINKAI_FILE_PROTOCOL = 'GET_SHINKAI_FILE_PROTOCOL',
 }
 
 export const DEFAULT_CHAT_CONFIG = {

diff --git a/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/index.ts b/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/index.ts
@@ -0,0 +1,14 @@
+import { getShinkaiFileProtocol as getShinkaiFileProtocolApi } from '@shinkai_network/shinkai-message-ts/api/tools/index';
+
+import type { GetShinkaiFileProtocolInput } from './types';
+
+export const getShinkaiFileProtocol = async ({
+  nodeAddress,
+  token,
+  file,
+}: GetShinkaiFileProtocolInput) => {
+  const result = await getShinkaiFileProtocolApi(nodeAddress, token, {
+    file,
+  });
+  return result;
+};
diff --git a/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/types.ts b/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/types.ts
@@ -0,0 +1,9 @@
+import { Token } from '@shinkai_network/shinkai-message-ts/api/general/types';
+import { GetShinkaiFileProtocolResponse } from '@shinkai_network/shinkai-message-ts/api/tools/types';
+
+export type GetShinkaiFileProtocolInput = Token & {
+  nodeAddress: string;
+  file: string;
+};
+
+export type GetShinkaiFileProtocolOutput = GetShinkaiFileProtocolResponse;
diff --git a/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/useGetShinkaiFileProtocol.ts b/libs/shinkai-node-state/src/v2/queries/getShinkaiFileProtocol/useGetShinkaiFileProtocol.ts
@@ -0,0 +1,32 @@
+import { QueryObserverOptions, useQuery } from '@tanstack/react-query';
+
+import { FunctionKeyV2 } from '../../constants';
+import { getShinkaiFileProtocol } from './index';
+import {
+  GetShinkaiFileProtocolInput,
+  GetShinkaiFileProtocolOutput,
+} from './types';
+
+export type UseGetShinkaiFileProtocol = [
+  FunctionKeyV2.GET_SHINKAI_FILE_PROTOCOL,
+  GetShinkaiFileProtocolInput,
+];
+type Options = QueryObserverOptions<
+  GetShinkaiFileProtocolOutput,
+  Error,
+  GetShinkaiFileProtocolOutput,
+  GetShinkaiFileProtocolOutput,
+  UseGetShinkaiFileProtocol
+>;
+
+export const useGetShinkaiFileProtocol = (
+  input: GetShinkaiFileProtocolInput,
+  options?: Omit<Options, 'queryKey' | 'queryFn'>,
+) => {
+  const response = useQuery({
+    queryKey: [FunctionKeyV2.GET_SHINKAI_FILE_PROTOCOL, input],
+    queryFn: () => getShinkaiFileProtocol(input),
+    ...options,
+  });
+  return response;
+};