feat(chat): add cancel button in chat to stop generation (#36)

Author: gespispace

src/common/chat/cloudChat.ts

@@ -11,6 +11,7 @@ import { configuration } from "../utils/configuration";
 type Parameters = {
   temperature: number;
   n_predict: number;
+  controller?: AbortController;
 };
 
 export const sendChatRequestCloud = async (
@@ -45,6 +46,9 @@ export const sendChatRequestCloud = async (
   });
 
   const stream = await model.pipe(parser).stream(messages, {
+    configurable: {
+      signal: parameters.controller,
+    },
     maxConcurrency: 1,
   });
 

src/common/chat/index.ts

@@ -26,6 +26,7 @@ export async function* chat(
   history: ChatMessage[],
   config?: {
     provideHighlightedText?: boolean;
+    abortController: AbortController;
   }
 ) {
   const loggerCompletion = logCompletion();
@@ -36,6 +37,7 @@ export async function* chat(
     n_predict: 4096,
     stop: [],
     temperature: 0.7,
+    controller: config?.abortController,
   };
 
   const { stopTask } = statusBar.startTask();

src/common/chat/localChat.ts

@@ -53,7 +53,10 @@ export async function* sendChatRequestLocal(
     const startTime = performance.now();
 
     let timings;
-    for await (const chunk of llama(prompt, parametersForCompletion, { url })) {
+    for await (const chunk of llama(prompt, parametersForCompletion, {
+      url,
+      controller: parameters.controller,
+    })) {
       // @ts-ignore
       if (chunk.data) {
         // @ts-ignore

src/common/panel/chat.ts

@@ -1,8 +1,9 @@
-import { Disposable, Webview, window, Uri } from "vscode";
+import { Disposable, Webview, Uri } from "vscode";
 import * as vscode from "vscode";
 import { getUri } from "../utils/getUri";
 import { getNonce } from "../utils/getNonce";
 import { chat } from "../chat";
+import { ChatMessage } from "../prompt/promptChat";
 
 export type MessageType =
   | {
@@ -13,15 +14,16 @@ export type MessageType =
     }
   | {
       type: "e2w-response";
+      id: string;
       command: string;
-      messageId: string;
       done: boolean;
       data: any;
     };
 
 export class ChatPanel implements vscode.WebviewViewProvider {
   private disposables: Disposable[] = [];
   private webview: Webview | undefined;
+  private messageCallback: Record<string, any> = {};
 
   constructor(private readonly extensionUri: vscode.Uri) {}
 
@@ -94,26 +96,19 @@ export class ChatPanel implements vscode.WebviewViewProvider {
   private setWebviewMessageListener(webview: Webview) {
     webview.onDidReceiveMessage(
       async (message: any) => {
-        const sendResponse = (messageToResponse: any, done: boolean) => {
-          this.postMessage({
-            type: "e2w-response",
-            command: message.type,
-            messageId: message.messageId,
-            data: messageToResponse,
-            done: done,
-          });
-        };
+        if (message.type in this.messageCallback) {
+          this.messageCallback[message.type]();
+          return;
+        }
         const type = message.type;
-        const data = message.data;
 
         switch (type) {
           case "sendMessage":
-            for await (const message of chat(data, {
-              provideHighlightedText: true,
-            })) {
-              sendResponse(message, false);
-            }
-            sendResponse("", true);
+            await this.handleStartGeneration({
+              chatMessage: message.data,
+              messageId: message.messageId,
+              messageType: message.type,
+            });
             return;
         }
       },
@@ -122,6 +117,47 @@ export class ChatPanel implements vscode.WebviewViewProvider {
     );
   }
 
+  private addMessageListener(
+    commandOrMessageId: string,
+    callback: (message: any) => void
+  ) {
+    this.messageCallback[commandOrMessageId] = callback;
+  }
+
+  private async handleStartGeneration({
+    messageId,
+    messageType,
+    chatMessage,
+  }: {
+    messageId: string;
+    messageType: string;
+    chatMessage: ChatMessage[];
+  }) {
+    const sendResponse = (messageToResponse: any, done: boolean) => {
+      this.postMessage({
+        type: "e2w-response",
+        id: messageId,
+        command: messageType,
+        data: messageToResponse,
+        done: done,
+      });
+    };
+    const abortController = new AbortController();
+
+    this.addMessageListener("abort-generate", () => {
+      abortController.abort();
+    });
+
+    for await (const message of chat(chatMessage, {
+      provideHighlightedText: true,
+      abortController,
+    })) {
+      sendResponse(message, false);
+    }
+
+    sendResponse("", true);
+  }
+
   public async sendMessageToWebview(
     command: MessageType["command"],
     data: MessageType["data"]

webviews/src/App.tsx

@@ -15,6 +15,7 @@ export const App = () => {
     input,
     setInput,
     startNewChat,
+    stop,
   } = useChat();
 
   useMessageListener("startNewChat", () => {
@@ -50,14 +51,11 @@ export const App = () => {
             buttonEnd={
               <VSCodeButton
                 appearance="icon"
-                disabled={isLoading}
-                onClick={handleSubmit}
+                onClick={isLoading ? stop : handleSubmit}
               >
                 <span
                   className={`codicon ${
-                    isLoading
-                      ? "codicon-loading codicon-modifier-spin codicon-modifier-disabled"
-                      : "codicon-send"
+                    isLoading ? "codicon-debug-stop" : "codicon-send"
                   }`}
                 ></span>
               </VSCodeButton>

webviews/src/components/TextArea/index.tsx

@@ -54,7 +54,7 @@ const TextArea = ({
               onSubmit();
             }
 
-            event.preventDefault(); // Prevents the addition of a new line in the text field
+            event.preventDefault();
           }
         }}
       ></textarea>
@@ -63,5 +63,4 @@ const TextArea = ({
   );
 };
 
-// 24 42 61
 export default TextArea;

webviews/src/hooks/useChat.ts

@@ -1,26 +1,58 @@
-import { useCallback, useState } from "react";
+import { useCallback, useRef, useState } from "react";
 import { randomMessageId } from "../utilities/messageId";
 import { vscode } from "../utilities/vscode";
 
+export type ChatMessage = {
+  role: string;
+  content: string;
+  chatMessageId: string;
+};
+
 export const useChat = () => {
-  const [chatMessages, setChatMessages] = useState<
-    {
-      role: string;
-      content: string;
-      chatMessageId: string;
-    }[]
-  >([]);
+  const [chatMessages, setChatMessages] = useState<ChatMessage[]>([]);
 
   const [input, setInput] = useState("");
   const [isLoading, setIsLoading] = useState(false);
 
+  const abortController = useRef(new AbortController());
+
+  const sendMessage = async (chatHistoryLocal: ChatMessage[]) => {
+    const messageId = randomMessageId();
+    for await (const newMessage of vscode.startGeneration(chatHistoryLocal, {
+      signal: abortController.current.signal,
+    })) {
+      setChatMessages((chatHistoryLocal) => {
+        const messages = chatHistoryLocal.filter(
+          (message) => message.chatMessageId !== messageId
+        );
+
+        const currentChatMessage = chatHistoryLocal.find(
+          (message) => message.chatMessageId === messageId
+        );
+
+        return [
+          ...messages,
+          {
+            role: "ai",
+            content: (currentChatMessage?.content || "") + newMessage,
+            chatMessageId: messageId,
+          },
+        ];
+      });
+    }
+    setIsLoading(false);
+  };
+
   const handleSubmit = () => {
     if (isLoading) {
       return;
     }
     if (input === "") {
       return;
     }
+    if (abortController.current.signal.aborted) {
+      abortController.current = new AbortController();
+    }
 
     setChatMessages((value) => {
       const messageId = randomMessageId();
@@ -41,40 +73,10 @@ export const useChat = () => {
     setInput("");
   };
 
-  const sendMessage = async (chatHistoryLocal: any) => {
-    const messageId = randomMessageId();
-    await vscode.postMessageCallback(
-      {
-        type: "sendMessage",
-        data: chatHistoryLocal,
-      },
-      (newMessage) => {
-        setChatMessages((chatHistoryLocal) => {
-          const messages = chatHistoryLocal.filter(
-            (message) => message.chatMessageId !== messageId
-          );
-
-          const currentChatMessage = chatHistoryLocal.find(
-            (message) => message.chatMessageId === messageId
-          );
-
-          if (newMessage.done) {
-            setIsLoading(false);
-            return chatHistoryLocal;
-          }
-
-          return [
-            ...messages,
-            {
-              role: "ai",
-              content: (currentChatMessage?.content || "") + newMessage.data,
-              chatMessageId: messageId,
-            },
-          ];
-        });
-      }
-    );
-  };
+  const stop = useCallback(() => {
+    abortController.current.abort();
+    setIsLoading(false);
+  }, [abortController]);
 
   const startNewChat = useCallback(() => {
     setChatMessages([]);
@@ -87,5 +89,6 @@ export const useChat = () => {
     setInput,
     handleSubmit,
     startNewChat,
+    stop,
   };
 };

webviews/src/utilities/transformCallback2AsyncGenerator.ts

@@ -0,0 +1,30 @@
+export class Transform<T> {
+  private open = true;
+  private queue: T[] = [];
+  private resolve: (() => void) | undefined;
+
+  async *stream(): AsyncGenerator<T> {
+    this.open = true;
+
+    while (this.open) {
+      if (this.queue.length) {
+        yield this.queue.shift()!;
+        continue;
+      }
+
+      await new Promise<void>((resolveLocal) => {
+        this.resolve = resolveLocal;
+      });
+    }
+  }
+
+  push(data: T): void {
+    this.queue.push(data);
+    this.resolve?.();
+  }
+
+  close(): void {
+    this.open = false;
+    this.resolve?.();
+  }
+}