From 2776200edf1c82ee34aa7c60105fff073c63d24d Mon Sep 17 00:00:00 2001
From: Anirban Kar <thecodacus@gmail.com>
Date: Thu, 14 Nov 2024 05:34:55 +0530
Subject: [PATCH] feat(code-template): code template with AI tool Calling
 Feature added with optin feature

---
 app/components/chat/BaseChat.tsx        | 19 +++++--
 app/components/chat/Chat.client.tsx     |  9 +++-
 app/components/chat/Messages.client.tsx | 16 ++++--
 app/components/chat/ToolManager.tsx     | 32 ++++++++++++
 app/components/chat/ToolMessage.tsx     | 14 ++++--
 app/components/ui/Popover.tsx           | 47 +++++++++++++++++
 app/components/ui/ToggleSwitch.tsx      | 55 ++++++++++++++++++++
 app/lib/.server/llm/stream-text.ts      |  3 +-
 app/lib/stores/tool.ts                  | 67 +++++++++++++++++--------
 app/lib/stores/workbench.ts             | 16 ++----
 app/routes/api.chat.ts                  |  7 ++-
 app/utils/types.ts                      | 10 ++++
 12 files changed, 244 insertions(+), 51 deletions(-)
 create mode 100644 app/components/chat/ToolManager.tsx
 create mode 100644 app/components/ui/Popover.tsx
 create mode 100644 app/components/ui/ToggleSwitch.tsx

diff --git a/app/components/chat/BaseChat.tsx b/app/components/chat/BaseChat.tsx
index 2a266d175..b0fef7380 100644
--- a/app/components/chat/BaseChat.tsx
+++ b/app/components/chat/BaseChat.tsx
@@ -15,6 +15,7 @@ import { APIKeyManager } from './APIKeyManager';
 import Cookies from 'js-cookie';
 
 import styles from './BaseChat.module.scss';
+import { ToolManager } from './ToolManager';
 
 const EXAMPLE_PROMPTS = [
   { text: 'Build a todo app in React using Tailwind' },
@@ -92,6 +93,8 @@ interface BaseChatProps {
   handleInputChange?: (event: React.ChangeEvent<HTMLTextAreaElement>) => void;
   enhancePrompt?: () => void;
   addToolResult?: ({ toolCallId, result }: { toolCallId: string; result: any }) => void;
+  toolConfig: IToolsConfig;
+  onToolConfigChange?: (val: IToolsConfig) => void;
 }
 
 export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
@@ -116,6 +119,8 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       enhancePrompt,
       handleStop,
       addToolResult,
+      toolConfig,
+      onToolConfigChange,
     },
     ref,
   ) => {
@@ -208,11 +213,15 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                   setProvider={setProvider}
                   providerList={providerList}
                 />
-                <APIKeyManager
-                  provider={provider}
-                  apiKey={apiKeys[provider] || ''}
-                  setApiKey={(key) => updateApiKey(provider, key)}
-                />
+                <div className="flex justify-between">
+                  <ToolManager toolConfig={toolConfig} onConfigChange={onToolConfigChange} />
+
+                  <APIKeyManager
+                    provider={provider}
+                    apiKey={apiKeys[provider] || ''}
+                    setApiKey={(key) => updateApiKey(provider, key)}
+                  />
+                </div>
                 <div
                   className={classNames(
                     'shadow-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background backdrop-filter backdrop-blur-[8px] rounded-lg overflow-hidden transition-all',
diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx
index 80e5e2c22..31f19d7e7 100644
--- a/app/components/chat/Chat.client.tsx
+++ b/app/components/chat/Chat.client.tsx
@@ -14,6 +14,7 @@ import { cubicEasingFn } from '~/utils/easings';
 import { createScopedLogger, renderLogger } from '~/utils/logger';
 import { BaseChat } from './BaseChat';
 import Cookies from 'js-cookie';
+import { toolStore } from '~/lib/stores/tool';
 
 const toastAnimation = cssTransition({
   enter: 'animated fadeInRight',
@@ -83,6 +84,8 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
 
   const { showChat } = useStore(chatStore);
 
+  const toolConfig = useStore(toolStore.config);
+
   const [animationScope, animate] = useAnimate();
 
   const [apiKeys, setApiKeys] = useState<Record<string, string>>({});
@@ -91,7 +94,9 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
     api: '/api/chat',
     body: {
       apiKeys,
+      toolEnabled: toolConfig.enabled,
     },
+    maxSteps: 3,
     onError: (error) => {
       logger.error('Request failed\n\n', error);
       toast.error('There was an error processing your request');
@@ -104,7 +109,7 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
       if (toolCall.toolName == 'askForConfirmation') return;
       logger.debug('Calling Tool:', toolCall);
       try {
-        let result = await workbenchStore.handleToolCall({
+        let result = await toolStore.handleToolCall({
           toolName: toolCall.toolName,
           args: toolCall.args,
           toolCallId: toolCall.toolCallId,
@@ -276,6 +281,8 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
       handleInputChange={handleInputChange}
       handleStop={abort}
       addToolResult={addToolResult}
+      toolConfig={toolConfig}
+      onToolConfigChange={(config) => toolStore.setConfig(config)}
       messages={messages.map((message, i) => {
         if (message.role === 'user') {
           return message;
diff --git a/app/components/chat/Messages.client.tsx b/app/components/chat/Messages.client.tsx
index d262ba359..0c1b60ff2 100644
--- a/app/components/chat/Messages.client.tsx
+++ b/app/components/chat/Messages.client.tsx
@@ -1,5 +1,5 @@
 import type { JSONValue, Message, ToolInvocation } from 'ai';
-import React from 'react';
+import React, { Fragment } from 'react';
 import { classNames } from '~/utils/classNames';
 import { AssistantMessage } from './AssistantMessage';
 import { UserMessage } from './UserMessage';
@@ -24,15 +24,23 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
             const isUserMessage = role === 'user';
             const isFirst = index === 0;
             const isLast = index === messages.length - 1;
+
+            // checking for annotated message marked as "hidden"
             if (message.annotations) {
               let isHidden = message.annotations.find((annotation: JSONValue) => {
                 if (typeof annotation !== 'object' || typeof annotation?.length === 'number') return false;
                 let object = annotation as any;
                 return object.visibility === 'hidden';
               });
-              console.log('isHidden', isHidden, message);
-
-              if (isHidden) return <></>;
+              if (isHidden) return <Fragment key={index}></Fragment>;
+            }
+            //hide confirmation message that has been answered
+            if (
+              message.toolInvocations?.length == 1 &&
+              message.toolInvocations[0].toolName === 'askForConfirmation' &&
+              (message.toolInvocations[0] as any).result
+            ) {
+              return <Fragment key={index}></Fragment>;
             }
             return (
               <div
diff --git a/app/components/chat/ToolManager.tsx b/app/components/chat/ToolManager.tsx
new file mode 100644
index 000000000..79db8be96
--- /dev/null
+++ b/app/components/chat/ToolManager.tsx
@@ -0,0 +1,32 @@
+import { IconButton } from '../ui/IconButton';
+import { Popover, PopoverContent, PopoverTrigger } from '../ui/Popover';
+import { ToggleSwitch } from '../ui/ToggleSwitch';
+import type { IToolsConfig } from '~/utils/types';
+
+interface ToolManagerProps {
+  toolConfig: IToolsConfig;
+  onConfigChange?: (val: IToolsConfig) => void;
+}
+
+export function ToolManager({ toolConfig, onConfigChange }: ToolManagerProps) {
+  return (
+    <>
+      {toolConfig && (
+        <div className="grid gap-4 text-sm">
+          <div className="flex items-center gap-2">
+            <label className="text-sm text-bolt-elements-textSecondary">Tool Calling</label>
+            <ToggleSwitch
+              checked={toolConfig.enabled}
+              onCheckedChange={(e) => {
+                onConfigChange?.({
+                  enabled: e,
+                  config: toolConfig.config,
+                });
+              }}
+            />
+          </div>
+        </div>
+      )}
+    </>
+  );
+}
diff --git a/app/components/chat/ToolMessage.tsx b/app/components/chat/ToolMessage.tsx
index 285810acc..80bce311c 100644
--- a/app/components/chat/ToolMessage.tsx
+++ b/app/components/chat/ToolMessage.tsx
@@ -15,7 +15,8 @@ const AskForConfirmation = (props: {
   addToolResult?: ({ toolCallId, result }: { toolCallId: string; result: any }) => void;
 }) => {
   const toolCallId = props.toolInvocation.toolCallId;
-  const addResult = (result: string) => (props.addToolResult ? props.addToolResult({ toolCallId, result }) : null);
+  const addResult = (result: string) =>
+    props.addToolResult ? props.addToolResult({ toolCallId, result: result }) : null;
 
   return (
     <div className="overflow-hidden pt-[4px] flex flex-col gap-2">
@@ -53,8 +54,11 @@ export function ToolMessage({ content, data: result, addToolResult }: ToolMessag
       result: (content as any).result,
     });
   }, [content]);
-  if (data?.name == 'askForConfirmation' && (data.state == 'result' || data.state == 'call')) {
-    return <AskForConfirmation toolInvocation={content} addToolResult={addToolResult}></AskForConfirmation>;
+  if (data?.name == 'askForConfirmation') {
+    if (!data.result) {
+      return <AskForConfirmation toolInvocation={content} addToolResult={addToolResult}></AskForConfirmation>;
+    }
+    return <></>;
   }
   return (
     <div className="overflow-hidden pt-[4px]">
@@ -86,7 +90,9 @@ export function ToolMessage({ content, data: result, addToolResult }: ToolMessag
                 )}
                 {/* //results in dimmed tone */}
                 {data.state == 'result' && (
-                  <div className="text-bolt-elements-textSecondary">Output: {(content as any).result}</div>
+                  <div className="text-bolt-elements-textSecondary">
+                    Output: {`${(content as any).result || ''}`.split('---')[0]}
+                  </div>
                 )}
               </div>
             </>
diff --git a/app/components/ui/Popover.tsx b/app/components/ui/Popover.tsx
new file mode 100644
index 000000000..f1947ea5f
--- /dev/null
+++ b/app/components/ui/Popover.tsx
@@ -0,0 +1,47 @@
+import * as React from 'react';
+import * as PopoverPrimitive from '@radix-ui/react-popover';
+
+const Popover = PopoverPrimitive.Root;
+const PopoverTrigger = PopoverPrimitive.Trigger;
+
+const PopoverContent = React.forwardRef<
+  React.ElementRef<typeof PopoverPrimitive.Content>,
+  React.ComponentPropsWithoutRef<typeof PopoverPrimitive.Content>
+>(({ className, align = 'center', sideOffset = 4, ...props }, ref) => (
+  <PopoverPrimitive.Content
+    ref={ref}
+    align={align}
+    sideOffset={sideOffset}
+    className={[
+      'z-50',
+      'w-auto',
+      'rounded-md',
+      'border border-bolt-elements-borderColor',
+      'bg-bolt-elements-background-depth-2', // Using depth-2 for better contrast
+      'p-4',
+      'text-bolt-elements-textPrimary',
+      'shadow-md',
+      'outline-none',
+      // Focus and ring states could use your theme colors
+      //   'focus-visible:(outline-none ring-2 ring-bolt-elements-item-contentAccent ring-offset-2 ring-offset-bolt-elements-background-depth-1)',
+      // Transitions
+      'transition-all duration-800',
+      'transform-gpu',
+      // State animations
+      'data-[state=open]:(opacity-100 scale-100)',
+      'data-[state=closed]:(opacity-0 scale-95)',
+      // Position animations
+      'data-[side=bottom]:translate-y-2',
+      'data-[side=left]:translate-x-2',
+      'data-[side=right]:-translate-x-2',
+      'data-[side=top]:-translate-y-2',
+      className,
+    ]
+      .filter(Boolean)
+      .join(' ')}
+    {...props}
+  />
+));
+PopoverContent.displayName = PopoverPrimitive.Content.displayName;
+
+export { Popover, PopoverTrigger, PopoverContent };
diff --git a/app/components/ui/ToggleSwitch.tsx b/app/components/ui/ToggleSwitch.tsx
new file mode 100644
index 000000000..1e68c7b50
--- /dev/null
+++ b/app/components/ui/ToggleSwitch.tsx
@@ -0,0 +1,55 @@
+import * as React from 'react';
+import * as SwitchPrimitives from '@radix-ui/react-switch';
+
+const ToggleSwitch = React.forwardRef<
+  React.ElementRef<typeof SwitchPrimitives.Root>,
+  React.ComponentPropsWithoutRef<typeof SwitchPrimitives.Root>
+>(({ className, ...props }, ref) => (
+  <SwitchPrimitives.Root
+    className={[
+      'peer',
+      'inline-flex',
+      'h-4',
+      'w-9',
+      'shrink-0',
+      'cursor-pointer',
+      'items-center',
+      'rounded-full',
+      'border-2',
+      'border-transparent',
+      'transition-colors duration-200 bolt-ease-cubic-bezier',
+      // Focus styles
+      'focus-visible:(outline-none ring-1)',
+      // Disabled styles
+      'disabled:(cursor-not-allowed opacity-50)',
+      // State styles
+      'data-[state=checked]:bg-bolt-elements-item-contentAccent',
+      'data-[state=unchecked]:bg-bolt-elements-button-secondary-background',
+      'hover:data-[state=unchecked]:bg-bolt-elements-button-secondary-backgroundHover',
+      className,
+    ]
+      .filter(Boolean)
+      .join(' ')}
+    {...props}
+    ref={ref}
+  >
+    <SwitchPrimitives.Thumb
+      className={[
+        'pointer-events-none',
+        'block',
+        'h-3',
+        'w-3',
+        'rounded-full',
+        'bg-bolt-elements-textPrimary',
+        'shadow-lg',
+        'ring-0',
+        'transition-transform duration-200 bolt-ease-cubic-bezier',
+        'data-[state=checked]:translate-x-5',
+        'data-[state=unchecked]:translate-x-0',
+      ].join(' ')}
+    />
+  </SwitchPrimitives.Root>
+));
+ToggleSwitch.displayName = SwitchPrimitives.Root.displayName;
+
+export { ToggleSwitch };
diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
index 4554d7f1b..e5500c13d 100644
--- a/app/lib/.server/llm/stream-text.ts
+++ b/app/lib/.server/llm/stream-text.ts
@@ -76,8 +76,7 @@ export function streamText(
     ...options,
     // toolChoice: messages.length == 1 ? { type: 'tool', toolName: 'selectCodeTemplate' } : 'none',
     toolChoice: 'auto',
-    tools: tools,
     experimental_toolCallStreaming: true,
-    maxSteps: 2
+    maxSteps: 3
   });
 }
diff --git a/app/lib/stores/tool.ts b/app/lib/stores/tool.ts
index 65c683090..e8238ff44 100644
--- a/app/lib/stores/tool.ts
+++ b/app/lib/stores/tool.ts
@@ -1,21 +1,36 @@
-import type { WebContainer, FileNode } from '@webcontainer/api';
-import type { FileMap, FilesStore } from './files';
+import type { WebContainer } from '@webcontainer/api';
 import { TEMPLATE_LIST } from '~/utils/constants';
 import * as nodePath from 'node:path';
-import type { EditorStore } from './editor';
-import type { WorkbenchStore } from './workbench';
-
-
-
+import { workbenchStore } from './workbench';
+import { webcontainer } from '../webcontainer';
+import { map, type MapStore } from 'nanostores';
+import type { IToolsConfig } from '~/utils/types';
+import Cookies from 'js-cookie';
 
 export class ToolStore {
     #webcontainer: Promise<WebContainer>;
-    #workbench: WorkbenchStore;
-    #editorStore: EditorStore;
-    constructor(webcontainerPromise: Promise<WebContainer>, workbench: WorkbenchStore, editorStore: EditorStore) {
+    config: MapStore<IToolsConfig> = map<IToolsConfig>({
+        enabled: false,
+        config: {}
+    })
+    constructor(webcontainerPromise: Promise<WebContainer>) {
         this.#webcontainer = webcontainerPromise;
-        this.#workbench = workbench;
-        this.#editorStore = editorStore;
+        let configString = Cookies.get('toolsConfig');
+        if (configString) {
+            try {
+                let config = JSON.parse(configString);
+                this.config.set(config);
+            } catch (error) {
+                console.error('Error parsing tools config:', error);
+            }
+        }
+    }
+    enableTools(enable: boolean) {
+        this.config.setKey('enabled', enable);
+    }
+    setConfig(config: IToolsConfig) {
+        this.config.set(config);
+        Cookies.set('toolsConfig', JSON.stringify(config));
     }
 
     async handleToolCall(payload: { toolName: string, args: any, toolCallId: string; }): Promise<string> {
@@ -59,18 +74,21 @@ export class ToolStore {
                 console.log("Writing to file", fullPath);
 
                 await webcontainer.fs.writeFile(file.path, file.content);
-                await this.#workbench.files.setKey(fullPath, { type: 'file', content: '', isBinary: false })
-                await this.#editorStore.updateFile(fullPath, file.content)
-                await this.#workbench.saveFile(file.path)
-                // await this.#editorStore.setSelectedFile(file.path)
-
-
+                await workbenchStore.files.setKey(fullPath, { type: 'file', content: '', isBinary: false })
+                await workbenchStore.updateFile(fullPath, file.content)
+                await workbenchStore.saveFile(file.path)
             }
+            return this.generateFormattedResult(`template imported successfully`, `
+                here is the imported content,
+                these files are loaded into the bolt. to not write them again, if it don't require changes
+                you only need to write the files that needs changing 
+                
+                ${JSON.stringify(files, null, 2)}
+            `)
         } catch (error) {
             console.error('error importing template', error);
             return 'error fetching template';
         }
-        return 'templace imported successfully';
     }
     private async getGitHubRepoContent(repoName: string, path: string = ''): Promise<{ name: string, path: string, content: string }[]> {
 
@@ -140,4 +158,13 @@ export class ToolStore {
             throw error;
         }
     }
-}
\ No newline at end of file
+    private generateFormattedResult(uiResult: string, aiResult?: string) {
+        return `
+        ${uiResult}
+        ---
+        ${aiResult || ""}
+        `
+    }
+}
+
+export const toolStore = new ToolStore(webcontainer);
\ No newline at end of file
diff --git a/app/lib/stores/workbench.ts b/app/lib/stores/workbench.ts
index 733774542..544112812 100644
--- a/app/lib/stores/workbench.ts
+++ b/app/lib/stores/workbench.ts
@@ -13,8 +13,6 @@ import JSZip from 'jszip';
 import { saveAs } from 'file-saver';
 import { Octokit, type RestEndpointMethodTypes } from "@octokit/rest";
 import * as nodePath from 'node:path';
-import type { WebContainerProcess } from '@webcontainer/api';
-import { ToolStore } from './tool';
 
 export interface ArtifactState {
   id: string;
@@ -34,7 +32,6 @@ export class WorkbenchStore {
   #filesStore = new FilesStore(webcontainer);
   #editorStore = new EditorStore(this.#filesStore);
   #terminalStore = new TerminalStore(webcontainer);
-  #toolStore = new ToolStore(webcontainer, this, this.#editorStore);
   artifacts: Artifacts = import.meta.hot?.data.artifacts ?? map({});
 
   showWorkbench: WritableAtom<boolean> = import.meta.hot?.data.showWorkbench ?? atom(false);
@@ -112,7 +109,9 @@ export class WorkbenchStore {
       }
     }
   }
-
+  updateFile(filePath: string, newContent: string) {
+    this.#editorStore.updateFile(filePath, newContent);
+  }
   setShowWorkbench(show: boolean) {
     this.showWorkbench.set(show);
   }
@@ -471,15 +470,6 @@ export class WorkbenchStore {
       console.error('Error pushing to GitHub:', error instanceof Error ? error.message : String(error));
     }
   }
-
-  async handleToolCall(payload: { toolName: string, args: any, toolCallId: string; }): Promise<string> {
-    try {
-      return await this.#toolStore.handleToolCall(payload);
-    } catch (error: any) {
-      console.error('Error calling tool:', error);
-      return 'error:' + error.message || 'error';
-    }
-  }
 }
 
 export const workbenchStore = new WorkbenchStore();
diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index 79cb8edbe..aefab033d 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -5,6 +5,7 @@ import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants';
 import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts';
 import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
 import SwitchableStream from '~/lib/.server/llm/switchable-stream';
+import tools from '~/lib/.server/llm/tools';
 import { z } from 'zod';
 
 export async function action(args: ActionFunctionArgs) {
@@ -12,9 +13,10 @@ export async function action(args: ActionFunctionArgs) {
 }
 
 async function chatAction({ context, request }: ActionFunctionArgs) {
-  const { messages, apiKeys } = await request.json<{
+  const { messages, apiKeys, toolEnabled } = await request.json<{
     messages: Messages,
     apiKeys: Record<string, string>
+    toolEnabled: boolean
   }>();
 
   const stream = new SwitchableStream();
@@ -43,7 +45,8 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
         const result = await streamText(messages, context.cloudflare.env, options, apiKeys);
 
         return stream.switchSource(result.toDataStream());
-      }
+      },
+      tools: toolEnabled ? tools : undefined,
     };
 
     const result = await streamText(messages, context.cloudflare.env, options, apiKeys);
diff --git a/app/utils/types.ts b/app/utils/types.ts
index b0cdd5e5f..fb9b77234 100644
--- a/app/utils/types.ts
+++ b/app/utils/types.ts
@@ -32,4 +32,14 @@ export interface TemplateInfo {
   name: string;
   label: string;
   githubRepo: string;
+}
+
+export interface IToolsConfig {
+  enabled: boolean;
+  //this section will be usefull for future features
+  config: Record<string, {
+    name: string;
+    id: string;
+    enabled: boolean;
+  }>
 }
\ No newline at end of file