Merge pull request #288 from thedadams/save-assistant

feat: add ability to save an assistant

Author: thedadams

components/chat.tsx

@@ -19,6 +19,7 @@ import ScriptToolsDropdown from '@/components/scripts/tool-dropdown';
 import AssistantNotFound from '@/components/assistant-not-found';
 import { generateThreadName, renameThread } from '@/actions/threads';
 import KnowledgeDropdown from '@/components/scripts/knowledge-dropdown';
+import SaveScriptDropdown from '@/components/scripts/script-save';
 
 interface ScriptProps {
   className?: string;
@@ -44,7 +45,7 @@ const Chat: React.FC<ScriptProps> = ({
   const {
     script,
     scriptDisplayName,
-    tool,
+    rootTool,
     showForm,
     setShowForm,
     formValues,
@@ -80,7 +81,7 @@ const Chat: React.FC<ScriptProps> = ({
       socket?.emit(
         'run',
         `${await getGatewayUrl()}/${script}`,
-        tool.name,
+        rootTool.name,
         formValues,
         workspace,
         thread
@@ -125,7 +126,7 @@ const Chat: React.FC<ScriptProps> = ({
           >
             {showForm && hasParams ? (
               <ToolForm
-                tool={tool}
+                tool={rootTool}
                 formValues={formValues}
                 handleInputChange={handleInputChange}
               />
@@ -139,6 +140,7 @@ const Chat: React.FC<ScriptProps> = ({
                     <div className="flex gap-2">
                       <ScriptToolsDropdown />
                       <KnowledgeDropdown />
+                      <SaveScriptDropdown />
                     </div>
                   </div>
                 )}
@@ -152,11 +154,11 @@ const Chat: React.FC<ScriptProps> = ({
               <Button
                 className="mt-4 w-full"
                 type="submit"
-                color={tool.chat ? 'primary' : 'secondary'}
+                color={rootTool.chat ? 'primary' : 'secondary'}
                 onPress={handleFormSubmit}
                 size="lg"
               >
-                {tool.chat ? 'Start chat' : 'Run script'}
+                {rootTool.chat ? 'Start chat' : 'Run script'}
               </Button>
             ) : (
               <ChatBar

components/chat/chatBar.tsx

@@ -34,7 +34,7 @@ const ChatBar = ({
     generating,
     interrupt,
     hasParams,
-    tool,
+    rootTool,
     setShowForm,
     messages,
     setMessages,
@@ -89,7 +89,7 @@ const ChatBar = ({
     if (generating) setLocked(false);
   }, [generating]);
 
-  if (!tool.chat) {
+  if (!rootTool.chat) {
     if (hasParams)
       return (
         <Button

components/scripts/knowledge-dropdown.tsx

@@ -45,7 +45,7 @@ const ScriptKnowledgeDropdown = () => {
           const newKnowledgeFiles = knowledgeFiles.filter((f) => f !== file);
           setKnowledgeFiles(newKnowledgeFiles);
           if (newKnowledgeFiles.length === 0) {
-            socket?.emit('removeTool', gatewayTool());
+            socket?.emit('removeTool', gatewayTool(), true);
           }
         }
       );

components/scripts/script-save.tsx

@@ -0,0 +1,67 @@
+import {
+  Button,
+  DropdownMenu,
+  Dropdown,
+  DropdownTrigger,
+  DropdownItem,
+} from '@nextui-org/react';
+import React, { useContext } from 'react';
+import { ChatContext } from '@/contexts/chat';
+import { PiFloppyDiskThin } from 'react-icons/pi';
+import { getScript, Script, updateScript } from '@/actions/me/scripts';
+import { stringify } from '@/actions/gptscript';
+import { gatewayTool } from '@/actions/knowledge/util';
+
+const SaveScriptDropdown = () => {
+  const { scriptId, rootTool, setRootTool, tools, socket, setScriptContent } =
+    useContext(ChatContext);
+
+  const knowledgeTool = gatewayTool();
+
+  async function saveScript() {
+    if (!rootTool.tools || rootTool.tools.length === 0) {
+      rootTool.tools = [];
+    }
+
+    // The knowledge tool is dynamic and not controlled by the user. Don't add it to the saved tool.
+    const addedTools = tools.filter((t) => t !== knowledgeTool);
+
+    rootTool.tools = rootTool.tools
+      .filter((t) => !addedTools.includes(t))
+      .concat(...addedTools);
+    try {
+      await updateScript({
+        content: await stringify([rootTool]),
+        id: parseInt(scriptId!),
+      } as Script);
+
+      socket?.emit('removeTool', addedTools, false);
+      setRootTool(rootTool);
+
+      setScriptContent((await getScript(scriptId!))?.script || null);
+    } catch (e) {
+      console.error(e);
+    }
+  }
+
+  return (
+    <Dropdown placement="bottom">
+      <DropdownTrigger>
+        <Button variant="light" isIconOnly>
+          <PiFloppyDiskThin className="size-5" />
+        </Button>
+      </DropdownTrigger>
+      <DropdownMenu>
+        {scriptId ? (
+          <DropdownItem key="save" onPress={saveScript}>
+            Save Assistant
+          </DropdownItem>
+        ) : (
+          (null as any)
+        )}
+      </DropdownMenu>
+    </Dropdown>
+  );
+};
+
+export default SaveScriptDropdown;

components/scripts/tool-dropdown.tsx

@@ -61,7 +61,7 @@ const ScriptToolsDropdown = () => {
   }
 
   function removeTool(tool: string) {
-    socket?.emit('removeTool', tool);
+    socket?.emit('removeTool', tool, true);
     setMessages((prev) => [
       ...prev,
       {

contexts/chat.tsx

@@ -31,8 +31,8 @@ interface ChatContextState {
   setSubTool: React.Dispatch<React.SetStateAction<string>>;
   setScript: React.Dispatch<React.SetStateAction<string>>;
   setScriptContent: React.Dispatch<React.SetStateAction<Block[] | null>>;
-  tool: Tool;
-  setTool: React.Dispatch<React.SetStateAction<Tool>>;
+  rootTool: Tool;
+  setRootTool: React.Dispatch<React.SetStateAction<Tool>>;
   program: Program | null;
   showForm: boolean;
   setShowForm: React.Dispatch<React.SetStateAction<boolean>>;
@@ -346,8 +346,8 @@ const ChatContextProvider: React.FC<ChatContextProps> = ({
         setScriptContent,
         workspace,
         setWorkspace,
-        tool,
-        setTool,
+        rootTool: tool,
+        setRootTool: setTool,
         program,
         subTool,
         setSubTool,

server/app.mjs

@@ -313,31 +313,40 @@ const mount = async (
     socket.emit('toolAdded', state.tools);
   });
 
-  socket.on('removeTool', async (tool) => {
+  socket.on('removeTool', async (tool, alterChatState) => {
     if (runningScript) {
       await runningScript.close();
       runningScript = null;
     }
 
-    const stateTools = (state.tools || []).filter((t) => t !== tool);
-    // find the root tool and then remove the tool
-    for (let block of script) {
-      if (block.type === 'tool') {
-        if (!block.tools) break;
-        block.tools = [...new Set(block.tools.filter((t) => t !== tool))];
-        break;
+    const stateTools = (state.tools || []).filter((t) => {
+      if (Array.isArray(tool)) {
+        return !tool.includes(t);
       }
-    }
+      return t !== tool;
+    });
 
     socket.emit('removingTool');
 
-    const loaded = await gptscript.loadTools(script, true);
-    const newTools = toChatStateTools(loaded?.program?.toolSet);
-    const currentState = JSON.parse(state.chatState);
-    currentState.continuation.state.completion.tools = newTools;
+    if (alterChatState) {
+      // find the root tool and then remove the tool
+      for (let block of script) {
+        if (block.type === 'tool') {
+          if (!block.tools) break;
+          block.tools = [...new Set(block.tools.filter((t) => t !== tool))];
+          break;
+        }
+      }
+
+      const loaded = await gptscript.loadTools(script, true);
+      const newTools = toChatStateTools(loaded?.program?.toolSet);
+      const currentState = JSON.parse(state.chatState);
+      currentState.continuation.state.completion.tools = newTools;
+
+      opts.chatState = JSON.stringify(currentState);
+      state.chatState = JSON.stringify(currentState);
+    }
 
-    opts.chatState = JSON.stringify(currentState);
-    state.chatState = JSON.stringify(currentState);
     state.tools = stateTools;
 
     if (threadID) {