feat: tool results assets + misc improvements (#568)

* fix: language when editing * feat: download tool assets * feat: improve cron tasks * feat: improve cron tasks * fix types
dcSpark · Dec 18, 2024 · e325b81 · e325b81
1 parent 9241bf4
commit e325b81
Show file tree

Hide file tree

Showing 15 changed files with 406 additions and 145 deletions.
diff --git a/apps/shinkai-desktop/src/components/cron-task/component/cron-task.tsx b/apps/shinkai-desktop/src/components/cron-task/component/cron-task.tsx
@@ -156,7 +156,10 @@ function CronTask({ mode, initialValues }: CronTaskProps) {
                   .tool_router_key
               : '',
         },
-        llmOrAgentId: defaultAgentId, // TODO: backend doesnt send this atm
+        llmOrAgentId:
+          'CreateJobWithConfigAndMessage' in initialValues.action
+            ? initialValues.action.CreateJobWithConfigAndMessage.llm_provider
+            : defaultAgentId,
       });
     }
   }, [form, initialValues]);
@@ -248,7 +251,7 @@ function CronTask({ mode, initialValues }: CronTaskProps) {
   return (
     <SubpageLayout
       className="max-w-4xl"
-      title={`${mode === 'create' ? 'Create' : 'Edit'}} Cron Task`}
+      title={`${mode === 'create' ? 'Create' : 'Edit'} Cron Task`}
     >
       <p className="text-gray-80 -mt-8 py-3 pb-6 text-center text-sm">
         Schedule recurring tasks at a specified time
@@ -317,7 +320,7 @@ function CronTask({ mode, initialValues }: CronTaskProps) {
                 </span>
 
                 <div className="space-y-4">
-                  <div className="grid grid-cols-[1fr_100px] items-center">
+                  <div className="grid grid-cols-[1fr_auto] items-center">
                     <span className="text-gray-80 text-xs">AI/Agent</span>
                     <AIModelSelector
                       onValueChange={(value) => {

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,
@@ -703,15 +714,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>
@@ -813,3 +842,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
@@ -415,6 +415,14 @@ export const useToolCode = ({
     }, 0);
   };
 
+  const { __created_files__: toolResultFiles, ...toolResultWithoutFiles } =
+    (toolResult || {
+      __created_files__: [],
+    }) as {
+      __created_files__: string[];
+      [key: string]: any;
+    };
+
   return {
     chatInboxId,
     toolCode,
@@ -437,7 +445,8 @@ export const useToolCode = ({
     metadataEditorRef,
     createToolCode,
     executeToolCodeQuery,
-    toolResult,
+    toolResult: toolResultWithoutFiles,
+    toolResultFiles: toolResultFiles,
     forceGenerateMetadata,
     isDirtyCodeEditor,
     setIsDirtyCodeEditor,

diff --git a/apps/shinkai-desktop/src/components/playground-tool/utils/code.ts b/apps/shinkai-desktop/src/components/playground-tool/utils/code.ts
@@ -34,3 +34,14 @@ export function detectLanguage(code: string): string {
   if (isTypeScript) return 'TypeScript';
   return 'Unknown';
 }
+
+export function getLanguage(language: string): CodeLanguage {
+  switch (language) {
+    case 'python':
+      return CodeLanguage.Python;
+    case 'typeScript':
+      return CodeLanguage.Typescript;
+    default:
+      return CodeLanguage.Typescript;
+  }
+}
diff --git a/apps/shinkai-desktop/src/pages/edit-tool.tsx b/apps/shinkai-desktop/src/pages/edit-tool.tsx
@@ -1,12 +1,11 @@
 import validator from '@rjsf/validator-ajv8';
-import { CodeLanguage } from '@shinkai_network/shinkai-message-ts/api/tools/types';
 import { buildInboxIdFromJobId } from '@shinkai_network/shinkai-message-ts/utils/inbox_name_handler';
 import { useGetPlaygroundTool } from '@shinkai_network/shinkai-node-state/v2/queries/getPlaygroundTool/useGetPlaygroundTool';
 import { useMemo } from 'react';
 import { useParams } from 'react-router-dom';
 
 import PlaygroundToolEditor from '../components/playground-tool/components/tool-playground';
-import { detectLanguage } from '../components/playground-tool/utils/code';
+import { getLanguage } from '../components/playground-tool/utils/code';
 import { useAuth } from '../store/auth';
 
 function EditToolPage() {
@@ -74,7 +73,7 @@ function EditToolPage() {
   return (
     <PlaygroundToolEditor
       createToolCodeFormInitialValues={{
-        language: detectLanguage(playgroundTool?.code ?? '') as CodeLanguage,
+        language: getLanguage(playgroundTool?.language ?? ''),
       }}
       initialChatInboxId={chatInboxId}
       mode="edit"